A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 19:30 UTC
[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>
Andreas Karatzas committed
May 26, 2026 at 21:53 UTC
5bdb181df5bddcc9b2173292d3ed9b4770a46716
3 files changed
+52
-9
vllm/model_executor/layers/activation.py
+10
-4
@@ -761,10 +761,16 @@ _ACTIVATION_AND_MUL_REGISTRY: LazyDict[nn.Module] = LazyDict(
761
)
762
763
764
-def get_act_and_mul_fn(act_fn_name: str) -> nn.Module:
764
+def get_act_and_mul_fn(act_fn_name: str, *, compile_native: bool = True) -> nn.Module:
765
"""Get an activation-and-mul (i.e. SiluAndMul) function by name."""
766
act_fn_name = act_fn_name.lower()
767
- if act_fn_name not in _ACTIVATION_AND_MUL_REGISTRY:
768
- raise ValueError(f"Activation function {act_fn_name!r} is not supported.")
767
770
- return _ACTIVATION_AND_MUL_REGISTRY[act_fn_name]
768
+ if not compile_native and act_fn_name in ("silu", "swish"):
769
+ return SiluAndMul(compile_native=False)
770
+
771
+ try:
772
+ return _ACTIVATION_AND_MUL_REGISTRY[act_fn_name]
773
+ except KeyError:
774
+ raise ValueError(
775
+ f"Activation function {act_fn_name!r} is not supported."
776
+ ) from None
vllm/model_executor/models/phi4mm.py
+37
-4
@@ -482,8 +482,8 @@ class Phi4MMImagePixelInputs(TensorSchema):
482
]
483
484
image_attention_mask: Annotated[
485
- torch.Tensor,
486
- TensorShape("bn", "nc", 32, 32), # H_mask, W_mask
485
+ torch.Tensor | list[torch.Tensor],
486
+ TensorShape("bn", "nc", 32, 32, dynamic_dims={"nc"}), # H_mask, W_mask
487
]
488
489
@@ -547,6 +547,39 @@ def cat_with_pad(tensors, dim, padding_value=0):
547
return output
548
549
550
+def stack_with_pad(
551
+ tensors: torch.Tensor | list[torch.Tensor],
552
+ padding_value: int | float = 0,
553
+) -> torch.Tensor:
554
+ """
555
+ Stack tensors, padding dimensions that differ across items.
556
+ """
557
+ if isinstance(tensors, torch.Tensor):
558
+ return tensors
559
+
560
+ assert len(tensors) > 0, "Cannot stack an empty tensor list"
561
+ first_shape = tensors[0].shape
562
+ if all(t.shape == first_shape for t in tensors):
563
+ return torch.stack(tensors)
564
+
565
+ ndim = tensors[0].dim()
566
+ assert all(t.dim() == ndim for t in tensors[1:]), (
567
+ "All tensors must have the same number of dimensions"
568
+ )
569
+
570
+ out_size = [
571
+ len(tensors),
572
+ *(max(t.shape[i] for t in tensors) for i in range(ndim)),
573
+ ]
574
+ output = tensors[0].new_full(out_size, padding_value)
575
+
576
+ for idx, tensor in enumerate(tensors):
577
+ slices = [idx, *(slice(0, size) for size in tensor.shape)]
578
+ output[tuple(slices)] = tensor
579
+
580
+ return output
581
+
582
+
583
class Phi4MMProcessingInfo(BaseProcessingInfo):
584
@property
585
def image_tokens(self) -> list[str]:
@@ -1176,9 +1209,9 @@ class Phi4MMForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal):
1209
self, image_input: Phi4MMImagePixelInputs
1210
) -> list[torch.Tensor]:
1211
dtype = next(self.vision_encoder.parameters()).dtype
1179
- pixel_values = image_input["pixel_values"].to(dtype)
1212
+ pixel_values = stack_with_pad(image_input["pixel_values"]).to(dtype)
1213
image_sizes = image_input["image_sizes"]
1181
- image_attention_mask = image_input["image_attention_mask"]
1214
+ image_attention_mask = stack_with_pad(image_input["image_attention_mask"])
1215
image_embeds = self.vision_encoder(
1216
pixel_values, image_sizes, image_attention_mask
1217
)
vllm/model_executor/models/qwen2_5_vl.py
+5
-1
@@ -81,6 +81,7 @@ from vllm.multimodal.inputs import (
81
)
82
from vllm.multimodal.parse import MultiModalDataItems
83
from vllm.multimodal.processing import PromptReplacement, PromptUpdate
84
+from vllm.platforms import current_platform
85
from vllm.sequence import IntermediateTensors
86
from vllm.utils.platform_utils import is_pin_memory_available
87
from vllm.utils.tensor_schema import TensorSchema, TensorShape
@@ -641,7 +642,10 @@ class Qwen2_5_VisionTransformer(nn.Module):
642
dim=self.hidden_size,
643
num_heads=self.num_heads,
644
mlp_hidden_dim=vision_config.intermediate_size,
644
- act_fn=get_act_and_mul_fn(vision_config.hidden_act),
645
+ act_fn=get_act_and_mul_fn(
646
+ vision_config.hidden_act,
647
+ compile_native=not current_platform.is_rocm(),
648
+ ),
649
norm_layer=norm_layer,
650
quant_config=quant_config,
651
prefix=f"{prefix}.blocks.{layer_idx}",