[Multimodal] Avoid extra d2d for encoder cudagraph with fused input norm (#56711)

Signed-off-by: Isotr0py <Isotr0py@outlook.com>

vllm/model_executor/layers/fusion/mm_input_norm.py

@@ -195,6 +195,10 @@ class IdentityInputNorm(nn.Module):
     ``visual_dtype`` here.
     """
 
+    @property
+    def input_dtype(self) -> torch.dtype | None:
+        return None
+
     def forward(
         self, pixel_values: torch.Tensor, visual_dtype: torch.dtype
     ) -> torch.Tensor:
@@ -251,6 +255,10 @@ class FusedMMInputNorm(CustomOp):
         self.register_buffer("weight", (rescale_factor / std).to(device))
         self.register_buffer("bias", (-mean / std).to(device))
 
+    @property
+    def input_dtype(self) -> torch.dtype | None:
+        return torch.uint8
+
     # ------------------------------------------------------------------
     # Internal helpers shared by the platform-specific forward_* methods
     # ------------------------------------------------------------------

vllm/model_executor/models/qwen2_5_vl.py

@@ -1976,7 +1976,8 @@ class Qwen2_5_VLForConditionalGeneration(
                 for _ in range(max_batch_size)
             ]
 
-        # Create dummy pixel_values
+        # Create dummy pixel_values; uint8 when normalization is fused
+        # on-device. Contents are overwritten before every replay.
         patch_embed = self.visual.patch_embed
         in_channels = patch_embed.proj.in_channels
         patch_size = patch_embed.patch_size
@@ -1985,8 +1986,11 @@ class Qwen2_5_VLForConditionalGeneration(
         flattened_patch_size = (
             in_channels * temporal_patch_size * patch_size * patch_size
         )
-        dummy_pixel_values = torch.randn(
-            total_patches, flattened_patch_size, device=device, dtype=dtype
+        dummy_pixel_values = torch.zeros(
+            total_patches,
+            flattened_patch_size,
+            device=device,
+            dtype=self.visual.input_norm.input_dtype or dtype,
         )
 
         # Override max_seqlen with a safe upper bound for capture.

vllm/model_executor/models/qwen2_vl.py

@@ -1616,7 +1616,8 @@ class Qwen2VLForConditionalGeneration(
                 for _ in range(max_batch_size)
             ]
 
-        # Create dummy pixel_values.
+        # Create dummy pixel_values; uint8 when normalization is fused
+        # on-device. Contents are overwritten before every replay.
         patch_embed = self.visual.patch_embed
         in_channels = patch_embed.proj.in_channels
         patch_size = patch_embed.patch_size
@@ -1625,8 +1626,11 @@ class Qwen2VLForConditionalGeneration(
         flattened_patch_size = (
             in_channels * temporal_patch_size * patch_size * patch_size
         )
-        dummy_pixel_values = torch.randn(
-            total_patches, flattened_patch_size, device=device, dtype=dtype
+        dummy_pixel_values = torch.zeros(
+            total_patches,
+            flattened_patch_size,
+            device=device,
+            dtype=self.visual.input_norm.input_dtype or dtype,
         )
 
         # max_seqlen.item() gets baked into the CUDA graph at capture time.

vllm/v1/attention/ops/vit_attn_wrappers.py

@@ -58,6 +58,13 @@ def flash_attn_maxseqlen_wrapper(
             max_seqlen = max_seqlen.item()
 
     q, k, v = (einops.rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])
+
+    if current_platform.is_cuda() and torch.cuda.is_current_stream_capturing():
+        # workaround for encoder CUDA-graph replay paddings with NaNs edge case
+        kwargs["out"] = torch.zeros(
+            q.shape[0], q.shape[1], v.shape[-1], dtype=q.dtype, device=q.device
+        )
+
     output = flash_attn_varlen_func(
         q,
         k,