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 5, 16:00 UTC
[Bugfix] Revert "Zero-init MLA attention output buffers to prevent NaN from CUDA graph padding" (#38359)
Signed-off-by: Elvir Crncevic <elvircrn@gmail.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com>
Elvir Crnčević committed
Apr 1, 2026 at 15:11 UTC
5e30e9b9a9b0dc91927f2ec4ffbd92c4af3825c0
2 files changed
+1
-57
vllm/v1/attention/backends/mla/cutlass_mla.py
+1
-14
@@ -162,11 +162,6 @@ 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
-
165
def _sm100_cutlass_mla_decode(
166
self,
167
q_nope: torch.Tensor,
@@ -223,15 +218,7 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):
218
if is_quantized_kv_cache(self.kv_cache_dtype)
219
else q_nope.dtype
220
)
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]
221
+ out = q_nope.new_empty((B_q, MAX_HEADS, D_latent), dtype=dtype)
222
lse = (
223
torch.empty((B_q, MAX_HEADS), dtype=torch.float32, device=q_nope.device)
224
if self.need_to_return_lse_for_decode
vllm/v1/attention/backends/mla/flashinfer_mla.py
-43
@@ -152,11 +152,6 @@ 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
-
155
def forward_mqa(
156
self,
157
q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
@@ -192,37 +187,6 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
187
if is_quantized_kv_cache(self.kv_cache_dtype):
188
self.bmm2_scale *= layer._k_scale_float
189
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
-
190
o = trtllm_batch_decode_with_kv_cache_mla(
191
query=q,
192
kv_cache=kv_c_and_k_pe_cache.unsqueeze(1),
@@ -235,15 +199,8 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
199
max_seq_len=attn_metadata.max_seq_len,
200
bmm1_scale=self.bmm1_scale,
201
bmm2_scale=self.bmm2_scale,
238
- **out_kwargs,
202
)
203
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
-
204
# Flatten the output for consistent shape
205
o = o.view(-1, o.shape[-2], o.shape[-1])
206