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, 15:45 UTC
[MTP][Sparse MLA] Take advantage of native MTP support in indexer when possible (#36982)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Matthew Bonanni committed
Mar 16, 2026 at 13:51 UTC
c88ea8338b9ad2f01bfb24c7bbf8ae6140866afd
2 files changed
+24
-13
csrc/sampler.cu
+1
-1
@@ -575,7 +575,7 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode(
575
// The range of logits within the row.
576
int rowStart = 0;
577
int seq_len = seqLens[rowIdx / next_n];
578
- int rowEnd = seq_len - next_n + (rowIdx % next_n) + 1;
578
+ int rowEnd = max(0, seq_len - next_n + (rowIdx % next_n) + 1);
579
580
// Local pointers to this block
581
if constexpr (!multipleBlocksPerRow && !mergeBlocks) {
vllm/v1/attention/backends/mla/indexer.py
+23
-12
@@ -206,6 +206,8 @@ def get_max_prefill_buffer_size(vllm_config: VllmConfig):
206
207
class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
208
reorder_batch_threshold: int = 1
209
+ natively_supported_next_n: list[int] = [1, 2]
210
+ # TODO (matt): integrate kernel with next_n = 4 support
211
212
@classmethod
213
def get_cudagraph_support(
@@ -231,7 +233,9 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
233
if self.vllm_config.speculative_config
234
else 0
235
)
236
+ next_n = self.num_speculative_tokens + 1
237
self.reorder_batch_threshold += self.num_speculative_tokens
238
+ self.use_flattening = next_n not in self.natively_supported_next_n
239
240
sm_count = num_compute_units(self.device.index)
241
self.num_sms = sm_count
@@ -241,10 +245,11 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
245
dtype=torch.int32,
246
device=self.device,
247
)
244
-
245
- # Pre-allocated buffers for flattening (spec decode).
248
+ self.offsets_buffer = torch.arange(
249
+ next_n, device=self.device, dtype=torch.int32
250
+ )
251
self.arange_buffer = torch.arange(
247
- scheduler_config.max_num_seqs * (1 + self.num_speculative_tokens),
252
+ scheduler_config.max_num_seqs * next_n,
253
dtype=torch.int32,
254
device=self.device,
255
)
@@ -323,7 +328,9 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
328
query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
329
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
330
split_decodes_and_prefills(
326
- common_attn_metadata, decode_threshold=self.reorder_batch_threshold
331
+ common_attn_metadata,
332
+ decode_threshold=self.reorder_batch_threshold,
333
+ require_uniform=not self.use_flattening,
334
)
335
)
336
@@ -372,11 +379,21 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
379
block_table.clamp_(min=0)
380
381
max_decode_len = int(decode_lens_cpu.max().item())
375
- if max_decode_len > 1:
382
+ next_n = 1 + self.num_speculative_tokens
383
+ use_native = not self.use_flattening and max_decode_len == next_n
384
+
385
+ if use_native and next_n > 1:
386
+ offsets = self.offsets_buffer
387
+ batch_size = num_decodes
388
+ elif max_decode_len > 1:
389
# Flatten multi-token decode requests into single-token
390
# batch entries, expanding seq_lens and block tables so
391
# the kernel always sees next_n=1.
392
393
+ # Also handles the edge case where use_flattening=False
394
+ # but max_decode_len != next_n (e.g. a batch containing some
395
+ # short prefills (q_len < next_n) and no true decodes).
396
+
397
# Assume 4 requests with seq_lens [10, 7, 12, 0] (the final req is
398
# padding) and decode_lens [3, 1, 4, 0] in the below example comments.
399
# The context lengths are therefore
@@ -428,13 +445,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
445
offsets = None
446
batch_size = num_decode_tokens
447
else:
431
- next_n = 1 + self.num_speculative_tokens
432
- if next_n > 1:
433
- offsets = torch.arange(
434
- next_n, device=self.device, dtype=torch.int32
435
- )
436
- else:
437
- offsets = None
448
+ offsets = None
449
batch_size = num_decodes
450
451
# DeepGEMM is required for the paged MQA logits on CUDA devices