A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 6, 10:45 UTC
[Bugfix] Zero-init MLA attention output buffers to prevent NaN from CUDA graph padding (#37442)
Signed-off-by: Elvir Crncevic <elvircrn@gmail.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
Elvir Crnčević committed
Mar 19, 2026 at 01:28 UTC
ef2c4f778df5aa07a44e663330e2dfdc16927d2a
2 files changed
+58
-1
vllm/v1/attention/backends/mla/cutlass_mla.py
+14
-1
@@ -162,6 +162,11 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):
162
# Share workspace buffer across all executions
163
self._workspace = g_sm100_workspace
164
165
+ # Pre-allocated output buffer, lazily sized on first call.
166
+ # Zero-init once to prevent NaN in padding slots (seq_lens=0)
167
+ # from contaminating downstream per-tensor reductions.
168
+ self._decode_out: torch.Tensor | None = None
169
+
170
def _sm100_cutlass_mla_decode(
171
self,
172
q_nope: torch.Tensor,
@@ -218,7 +223,15 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):
223
if is_quantized_kv_cache(self.kv_cache_dtype)
224
else q_nope.dtype
225
)
221
- out = q_nope.new_empty((B_q, MAX_HEADS, D_latent), dtype=dtype)
226
+ # Reuse pre-allocated zero-init output buffer to avoid a memset
227
+ # kernel on every CUDA graph replay.
228
+ if (
229
+ self._decode_out is None
230
+ or self._decode_out.shape[0] < B_q
231
+ or self._decode_out.dtype != dtype
232
+ ):
233
+ self._decode_out = q_nope.new_zeros((B_q, MAX_HEADS, D_latent), dtype=dtype)
234
+ out = self._decode_out[:B_q]
235
lse = (
236
torch.empty((B_q, MAX_HEADS), dtype=torch.float32, device=q_nope.device)
237
if self.need_to_return_lse_for_decode
vllm/v1/attention/backends/mla/flashinfer_mla.py
+44
@@ -21,6 +21,7 @@ from vllm.v1.attention.backend import (
21
AttentionLayer,
22
AttentionType,
23
MultipleOf,
24
+ is_quantized_kv_cache,
25
)
26
from vllm.v1.attention.backends.utils import KVCacheLayoutType
27
@@ -151,6 +152,11 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
152
self.bmm1_scale: float | None = None
153
self.bmm2_scale: float | None = None
154
155
+ # Pre-allocated output buffer, lazily sized on first call.
156
+ # Zero-init once to prevent NaN in padding slots (seq_lens=0)
157
+ # from contaminating downstream per-tensor reductions.
158
+ self._decode_out: torch.Tensor | None = None
159
+
160
def forward_mqa(
161
self,
162
q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
@@ -186,6 +192,37 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
192
if self.kv_cache_dtype.startswith("fp8"):
193
self.bmm2_scale *= layer._k_scale_float
194
195
+ # Reuse pre-allocated zero-init output buffer to avoid a memset
196
+ # kernel on every CUDA graph replay.
197
+ # q is 4D: (batch, q_len_per_req, num_heads, head_dim)
198
+ # FlashInfer has a bug where out= validation hardcodes 3D shape
199
+ # (batch, num_heads, kv_lora_rank), but the kernel writes 4D
200
+ # (batch, q_len, num_heads, kv_lora_rank) when q_len > 1.
201
+ # So we can only pass out= for single-token decode (q_len == 1).
202
+ # For q_len > 1, we zero padding slots after the kernel returns.
203
+ # TODO: upstream fix to FlashInfer
204
+ B, q_len_per_req = q.shape[0], q.shape[1]
205
+ out_kwargs: dict[str, torch.Tensor] = {}
206
+ if q_len_per_req == 1:
207
+ dtype = (
208
+ torch.bfloat16
209
+ if is_quantized_kv_cache(self.kv_cache_dtype)
210
+ else q.dtype
211
+ )
212
+ if (
213
+ self._decode_out is None
214
+ or self._decode_out.shape[0] < B
215
+ or self._decode_out.dtype != dtype
216
+ ):
217
+ self._decode_out = torch.zeros(
218
+ B,
219
+ q.shape[2],
220
+ self.kv_lora_rank,
221
+ dtype=dtype,
222
+ device=q.device,
223
+ )
224
+ out_kwargs["out"] = self._decode_out[:B]
225
+
226
o = trtllm_batch_decode_with_kv_cache_mla(
227
query=q,
228
kv_cache=kv_c_and_k_pe_cache.unsqueeze(1),
@@ -198,8 +235,15 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
235
max_seq_len=attn_metadata.max_seq_len,
236
bmm1_scale=self.bmm1_scale,
237
bmm2_scale=self.bmm2_scale,
238
+ **out_kwargs,
239
)
240
241
+ # For q_len > 1, we can't pass out= so we work around by zeroing padding slots
242
+ if not out_kwargs:
243
+ num_real = attn_metadata.num_decodes
244
+ if num_real < o.shape[0]:
245
+ o[num_real:] = 0
246
+
247
# Flatten the output for consistent shape
248
o = o.view(-1, o.shape[-2], o.shape[-1])
249