[Multimodal] Avoid extra d2d for encoder cudagraph with fused input norm (#56711)
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
Isotr0py committed
Sep 28, 2026 at 21:04 UTC
522101f6ede5ffd57844a8729cf0f2f6fe1a7a47
4 files changed
+29
-6
vllm/model_executor/layers/fusion/mm_input_norm.py
+8
@@ -195,6 +195,10 @@ class IdentityInputNorm(nn.Module):
195
``visual_dtype`` here.
196
"""
197
198
+ @property
199
+ def input_dtype(self) -> torch.dtype | None:
200
+ return None
201
+
202
def forward(
203
self, pixel_values: torch.Tensor, visual_dtype: torch.dtype
204
) -> torch.Tensor:
@@ -251,6 +255,10 @@ class FusedMMInputNorm(CustomOp):
255
self.register_buffer("weight", (rescale_factor / std).to(device))
256
self.register_buffer("bias", (-mean / std).to(device))
257
258
+ @property
259
+ def input_dtype(self) -> torch.dtype | None:
260
+ return torch.uint8
261
+
262
# ------------------------------------------------------------------
263
# Internal helpers shared by the platform-specific forward_* methods
264
# ------------------------------------------------------------------
vllm/model_executor/models/qwen2_5_vl.py
+7
-3
@@ -1976,7 +1976,8 @@ class Qwen2_5_VLForConditionalGeneration(
1976
for _ in range(max_batch_size)
1977
]
1978
1979
- # Create dummy pixel_values
1979
+ # Create dummy pixel_values; uint8 when normalization is fused
1980
+ # on-device. Contents are overwritten before every replay.
1981
patch_embed = self.visual.patch_embed
1982
in_channels = patch_embed.proj.in_channels
1983
patch_size = patch_embed.patch_size
@@ -1985,8 +1986,11 @@ class Qwen2_5_VLForConditionalGeneration(
1986
flattened_patch_size = (
1987
in_channels * temporal_patch_size * patch_size * patch_size
1988
)
1988
- dummy_pixel_values = torch.randn(
1989
- total_patches, flattened_patch_size, device=device, dtype=dtype
1989
+ dummy_pixel_values = torch.zeros(
1990
+ total_patches,
1991
+ flattened_patch_size,
1992
+ device=device,
1993
+ dtype=self.visual.input_norm.input_dtype or dtype,
1994
)
1995
1996
# Override max_seqlen with a safe upper bound for capture.
vllm/model_executor/models/qwen2_vl.py
+7
-3
@@ -1616,7 +1616,8 @@ class Qwen2VLForConditionalGeneration(
1616
for _ in range(max_batch_size)
1617
]
1618
1619
- # Create dummy pixel_values.
1619
+ # Create dummy pixel_values; uint8 when normalization is fused
1620
+ # on-device. Contents are overwritten before every replay.
1621
patch_embed = self.visual.patch_embed
1622
in_channels = patch_embed.proj.in_channels
1623
patch_size = patch_embed.patch_size
@@ -1625,8 +1626,11 @@ class Qwen2VLForConditionalGeneration(
1626
flattened_patch_size = (
1627
in_channels * temporal_patch_size * patch_size * patch_size
1628
)
1628
- dummy_pixel_values = torch.randn(
1629
- total_patches, flattened_patch_size, device=device, dtype=dtype
1629
+ dummy_pixel_values = torch.zeros(
1630
+ total_patches,
1631
+ flattened_patch_size,
1632
+ device=device,
1633
+ dtype=self.visual.input_norm.input_dtype or dtype,
1634
)
1635
1636
# max_seqlen.item() gets baked into the CUDA graph at capture time.
vllm/v1/attention/ops/vit_attn_wrappers.py
+7
@@ -58,6 +58,13 @@ def flash_attn_maxseqlen_wrapper(
58
max_seqlen = max_seqlen.item()
59
60
q, k, v = (einops.rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])
61
+
62
+ if current_platform.is_cuda() and torch.cuda.is_current_stream_capturing():
63
+ # workaround for encoder CUDA-graph replay paddings with NaNs edge case
64
+ kwargs["out"] = torch.zeros(
65
+ q.shape[0], q.shape[1], v.shape[-1], dtype=q.dtype, device=q.device
66
+ )
67
+
68
output = flash_attn_varlen_func(
69
q,
70
k,