[ROCm][CI] Fix ROCm multimodal Qwen2.5-VL activation compile and Phi4MM ragged image mask handling (#43647)

Signed-off-by: Andreas Karatzas <akaratza@amd.com>

vllm/model_executor/layers/activation.py

@@ -761,10 +761,16 @@ _ACTIVATION_AND_MUL_REGISTRY: LazyDict[nn.Module] = LazyDict(
 )
 
 
-def get_act_and_mul_fn(act_fn_name: str) -> nn.Module:
+def get_act_and_mul_fn(act_fn_name: str, *, compile_native: bool = True) -> nn.Module:
     """Get an activation-and-mul (i.e. SiluAndMul) function by name."""
     act_fn_name = act_fn_name.lower()
-    if act_fn_name not in _ACTIVATION_AND_MUL_REGISTRY:
-        raise ValueError(f"Activation function {act_fn_name!r} is not supported.")
 
-    return _ACTIVATION_AND_MUL_REGISTRY[act_fn_name]
+    if not compile_native and act_fn_name in ("silu", "swish"):
+        return SiluAndMul(compile_native=False)
+
+    try:
+        return _ACTIVATION_AND_MUL_REGISTRY[act_fn_name]
+    except KeyError:
+        raise ValueError(
+            f"Activation function {act_fn_name!r} is not supported."
+        ) from None

vllm/model_executor/models/phi4mm.py

@@ -482,8 +482,8 @@ class Phi4MMImagePixelInputs(TensorSchema):
     ]
 
     image_attention_mask: Annotated[
-        torch.Tensor,
-        TensorShape("bn", "nc", 32, 32),  # H_mask, W_mask
+        torch.Tensor | list[torch.Tensor],
+        TensorShape("bn", "nc", 32, 32, dynamic_dims={"nc"}),  # H_mask, W_mask
     ]
 
 
@@ -547,6 +547,39 @@ def cat_with_pad(tensors, dim, padding_value=0):
     return output
 
 
+def stack_with_pad(
+    tensors: torch.Tensor | list[torch.Tensor],
+    padding_value: int | float = 0,
+) -> torch.Tensor:
+    """
+    Stack tensors, padding dimensions that differ across items.
+    """
+    if isinstance(tensors, torch.Tensor):
+        return tensors
+
+    assert len(tensors) > 0, "Cannot stack an empty tensor list"
+    first_shape = tensors[0].shape
+    if all(t.shape == first_shape for t in tensors):
+        return torch.stack(tensors)
+
+    ndim = tensors[0].dim()
+    assert all(t.dim() == ndim for t in tensors[1:]), (
+        "All tensors must have the same number of dimensions"
+    )
+
+    out_size = [
+        len(tensors),
+        *(max(t.shape[i] for t in tensors) for i in range(ndim)),
+    ]
+    output = tensors[0].new_full(out_size, padding_value)
+
+    for idx, tensor in enumerate(tensors):
+        slices = [idx, *(slice(0, size) for size in tensor.shape)]
+        output[tuple(slices)] = tensor
+
+    return output
+
+
 class Phi4MMProcessingInfo(BaseProcessingInfo):
     @property
     def image_tokens(self) -> list[str]:
@@ -1176,9 +1209,9 @@ class Phi4MMForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal):
         self, image_input: Phi4MMImagePixelInputs
     ) -> list[torch.Tensor]:
         dtype = next(self.vision_encoder.parameters()).dtype
-        pixel_values = image_input["pixel_values"].to(dtype)
+        pixel_values = stack_with_pad(image_input["pixel_values"]).to(dtype)
         image_sizes = image_input["image_sizes"]
-        image_attention_mask = image_input["image_attention_mask"]
+        image_attention_mask = stack_with_pad(image_input["image_attention_mask"])
         image_embeds = self.vision_encoder(
             pixel_values, image_sizes, image_attention_mask
         )

vllm/model_executor/models/qwen2_5_vl.py

@@ -81,6 +81,7 @@ from vllm.multimodal.inputs import (
 )
 from vllm.multimodal.parse import MultiModalDataItems
 from vllm.multimodal.processing import PromptReplacement, PromptUpdate
+from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
 from vllm.utils.platform_utils import is_pin_memory_available
 from vllm.utils.tensor_schema import TensorSchema, TensorShape
@@ -641,7 +642,10 @@ class Qwen2_5_VisionTransformer(nn.Module):
                     dim=self.hidden_size,
                     num_heads=self.num_heads,
                     mlp_hidden_dim=vision_config.intermediate_size,
-                    act_fn=get_act_and_mul_fn(vision_config.hidden_act),
+                    act_fn=get_act_and_mul_fn(
+                        vision_config.hidden_act,
+                        compile_native=not current_platform.is_rocm(),
+                    ),
                     norm_layer=norm_layer,
                     quant_config=quant_config,
                     prefix=f"{prefix}.blocks.{layer_idx}",