From be441ac536181ee54546cb3ed2f3f505ebe6ff58 Mon Sep 17 00:00:00 2001
From: varun <vrh>
Date: Tue, 17 Mar 2026 17:24:29 +0000
Subject: [PATCH] zero out attn output

fixes and zero out input to fi kernel
---
 vllm/model_executor/layers/attention/mla_attention.py | 11 ++++++++++-
 vllm/v1/attention/backends/mla/flashinfer_mla.py      |  5 +++++
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py
index b613f3ba983..5ceabce5ede 100644
--- a/vllm/model_executor/layers/attention/mla_attention.py
+++ b/vllm/model_executor/layers/attention/mla_attention.py
@@ -518,7 +518,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
                 self._k_scale,
             )
             if self.attn_backend.accept_output_buffer:
-                output = torch.empty(output_shape, dtype=q.dtype, device=q.device)
+                output = torch.zeros(output_shape, dtype=q.dtype, device=q.device)
                 torch.ops.vllm.unified_mla_attention_with_output(
                     q,
                     kv_c_normed,
@@ -692,6 +692,15 @@ class MLAAttention(nn.Module, AttentionLayerBase):
                 assert attn_metadata.decode is not None
             attn_out, lse = self.impl.forward_mqa(mqa_q, kv_cache, attn_metadata, self)
 
+            # Zero out padded region for CUDA graphs. Padded requests have
+            # seq_lens=0 in the decode metadata (device tensor updated in-place
+            # before each graph replay), so masked_fill_ always launches a kernel
+            # during capture yet only zeros padding slots during replay.
+            if attn_metadata.decode is not None:
+                decode_seq_lens = attn_metadata.decode.seq_lens
+                pad_mask = (decode_seq_lens == 0).view(-1, 1, 1)
+                attn_out[:decode_seq_lens.shape[0]].masked_fill_(pad_mask, 0)
+
             # correct dcp attn_out with lse.
             if self.impl.dcp_world_size > 1:
                 if self.dcp_a2a:
diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla.py b/vllm/v1/attention/backends/mla/flashinfer_mla.py
index ec8f4e6400b..7ac37e6dadd 100644
--- a/vllm/v1/attention/backends/mla/flashinfer_mla.py
+++ b/vllm/v1/attention/backends/mla/flashinfer_mla.py
@@ -181,7 +181,12 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
         if self.bmm2_scale is None:
             self.bmm2_scale = layer._v_scale_float
 
+        out = torch.zeros(q.shape[0], q.shape[2], self.kv_lora_rank,
+                  dtype=torch.bfloat16, device=q.device)
+
+        self._workspace_buffer.fill_(0)
         o = trtllm_batch_decode_with_kv_cache_mla(
+            out=out,
             query=q,
             kv_cache=kv_c_and_k_pe_cache.unsqueeze(1),
             workspace_buffer=self._workspace_buffer,