diff --git a/vllm/model_executor/layers/fusion/mm_input_norm.py b/vllm/model_executor/layers/fusion/mm_input_norm.py index a21828033d..0017a48d6d 100644 --- a/vllm/model_executor/layers/fusion/mm_input_norm.py +++ b/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 # ------------------------------------------------------------------ diff --git a/vllm/model_executor/models/qwen2_5_vl.py b/vllm/model_executor/models/qwen2_5_vl.py index 864a77f095..581ff77303 100644 --- a/vllm/model_executor/models/qwen2_5_vl.py +++ b/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. diff --git a/vllm/model_executor/models/qwen2_vl.py b/vllm/model_executor/models/qwen2_vl.py index 3dae3ffd2a..5a7ccacc48 100644 --- a/vllm/model_executor/models/qwen2_vl.py +++ b/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. diff --git a/vllm/v1/attention/ops/vit_attn_wrappers.py b/vllm/v1/attention/ops/vit_attn_wrappers.py index dccc670ce5..36f7de393c 100644 --- a/vllm/v1/attention/ops/vit_attn_wrappers.py +++ b/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,