[Bugfix] Zero-init MLA attention output buffers to prevent NaN from CUDA graph padding

vllm/v1/attention/backends/mla/cutlass_mla.py

@@ -162,6 +162,11 @@ def __init__(
         # Share workspace buffer across all executions
         self._workspace = g_sm100_workspace
 
+        # Pre-allocated output buffer, lazily sized on first call.
+        # Zero-init once to prevent NaN in padding slots (seq_lens=0)
+        # from contaminating downstream per-tensor reductions.
+        self._decode_out: torch.Tensor | None = None
+
     def _sm100_cutlass_mla_decode(
         self,
         q_nope: torch.Tensor,
@@ -218,7 +223,15 @@ def _sm100_cutlass_mla_decode(
             if is_quantized_kv_cache(self.kv_cache_dtype)
             else q_nope.dtype
         )
-        out = q_nope.new_empty((B_q, MAX_HEADS, D_latent), dtype=dtype)
+        # Reuse pre-allocated zero-init output buffer to avoid a memset
+        # kernel on every CUDA graph replay.
+        if (
+            self._decode_out is None
+            or self._decode_out.shape[0] < B_q
+            or self._decode_out.dtype != dtype
+        ):
+            self._decode_out = q_nope.new_zeros((B_q, MAX_HEADS, D_latent), dtype=dtype)
+        out = self._decode_out[:B_q]
         lse = (
             torch.empty((B_q, MAX_HEADS), dtype=torch.float32, device=q_nope.device)
             if self.need_to_return_lse_for_decode

vllm/v1/attention/backends/mla/flashinfer_mla.py

@@ -21,6 +21,7 @@
     AttentionLayer,
     AttentionType,
     MultipleOf,
+    is_quantized_kv_cache,
 )
 from vllm.v1.attention.backends.utils import KVCacheLayoutType
 
@@ -151,6 +152,11 @@ def __init__(
         self.bmm1_scale: float | None = None
         self.bmm2_scale: float | None = None
 
+        # Pre-allocated output buffer, lazily sized on first call.
+        # Zero-init once to prevent NaN in padding slots (seq_lens=0)
+        # from contaminating downstream per-tensor reductions.
+        self._decode_out: torch.Tensor | None = None
+
     def forward_mqa(
         self,
         q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
@@ -181,6 +187,26 @@ def forward_mqa(
         if self.bmm2_scale is None:
             self.bmm2_scale = layer._v_scale_float
 
+        # Reuse pre-allocated zero-init output buffer to avoid a memset
+        # kernel on every CUDA graph replay.
+        B = q.shape[0]
+        dtype = (
+            torch.bfloat16 if is_quantized_kv_cache(self.kv_cache_dtype) else q.dtype
+        )
+        if (
+            self._decode_out is None
+            or self._decode_out.shape[0] < B
+            or self._decode_out.dtype != dtype
+        ):
+            self._decode_out = torch.zeros(
+                B,
+                q.shape[2],
+                self.kv_lora_rank,
+                dtype=dtype,
+                device=q.device,
+            )
+        out = self._decode_out[:B]
+
         o = trtllm_batch_decode_with_kv_cache_mla(
             query=q,
             kv_cache=kv_c_and_k_pe_cache.unsqueeze(1),
@@ -193,6 +219,7 @@ def forward_mqa(
             max_seq_len=attn_metadata.max_seq_len,
             bmm1_scale=self.bmm1_scale,
             bmm2_scale=self.bmm2_scale,
+            out=out,
         )
 
         # Flatten the output for consistent shape