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 3, 14:04 UTC
[Bugfix] Fix DeepGEMM context lens contiguity in MLA indexer (#42135)
Signed-off-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>
Mohammad Miadh Angkad committed
May 15, 2026 at 23:29 UTC
ee58665aac60f60f46b0f0f153b792327ef50da2
1 file changed
+16
-18
vllm/v1/attention/backends/mla/indexer.py
+16
-18
@@ -288,20 +288,14 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
288
dtype=torch.int32,
289
device=self.device,
290
)
291
- if not self.use_flattening and next_n > 1:
292
- # Native MTP: 2D buffer for per-token seq_lens.
293
- self.decode_seq_lens_buffer = torch.zeros(
294
- (scheduler_config.max_num_seqs, next_n),
295
- dtype=torch.int32,
296
- device=self.device,
297
- )
298
- else:
299
- # Flattening or no MTP: 1D buffer for expanded per-token seq_lens.
300
- self.decode_seq_lens_buffer = torch.zeros(
301
- (scheduler_config.max_num_batched_tokens,),
302
- dtype=torch.int32,
303
- device=self.device,
304
- )
291
+ # Shared workspace for decode seq_lens. Native MTP views this as
292
+ # (B, max_decode_len) at runtime, keeping context_lens contiguous even
293
+ # when max_decode_len is smaller than next_n.
294
+ self.decode_seq_lens_buffer = torch.zeros(
295
+ (scheduler_config.max_num_batched_tokens,),
296
+ dtype=torch.int32,
297
+ device=self.device,
298
+ )
299
self.arange_buffer = torch.arange(
300
max(
301
scheduler_config.max_num_seqs * next_n,
@@ -373,7 +367,8 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
367
Plain decode or spec-decode with 2D per-token context lengths.
368
369
Returns (seq_lens, block_table, decode_lens, batch_size, requires_padding).
376
- seq_lens is 1D (batch_size,) for flatten/plain, 2D (B, next_n) for native MTP.
370
+ seq_lens is 1D (batch_size,) for flatten/plain, 2D (B, max_decode_len)
371
+ for native MTP.
372
"""
373
min_decode_len = int(decode_lens_cpu.min().item())
374
if not use_native and max_decode_len > 1:
@@ -454,16 +449,19 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
449
# (requires_padding) instead.
450
requires_padding = min_decode_len != max_decode_len
451
if use_native and next_n > 1:
457
- assert self.decode_seq_lens_buffer.dim() == 2
452
+ assert self.decode_seq_lens_buffer.dim() == 1
453
# (B, max_decode_len): token j attends to
454
# L - max_decode_len + j + 1 KV tokens.
460
- self.decode_seq_lens_buffer[:num_decodes, :max_decode_len] = (
455
+ seq_lens_buffer = self.decode_seq_lens_buffer[
456
+ : num_decodes * max_decode_len
457
+ ].view(num_decodes, max_decode_len)
458
+ seq_lens_buffer[:] = (
459
seq_lens.unsqueeze(1)
460
- max_decode_len
461
+ 1
462
+ self.offsets_buffer[:max_decode_len]
463
)
466
- seq_lens = self.decode_seq_lens_buffer[:num_decodes, :max_decode_len]
464
+ seq_lens = seq_lens_buffer
465
return seq_lens, block_table, decode_lens, num_decodes, requires_padding
466
467
def build(