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, 23:00 UTC
[DSV3.2][MTP] Optimize Indexer MTP handling (#36723)
Signed-off-by: Benjamin Chislett <bchislett@nvidia.com>
Benjamin Chislett committed
Mar 11, 2026 at 00:16 UTC
9040cd40af6bacfd20d1db8637a189f966a2fcc4
1 file changed
+7
-3
vllm/v1/attention/backends/mla/indexer.py
+7
-3
@@ -384,12 +384,14 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
384
385
# [7, 6, 8, 0] -> [7, 7, 7, 6, 8, 8, 8, 8]
386
expanded_base = torch.repeat_interleave(
387
- seq_lens - decode_lens, decode_lens
387
+ seq_lens - decode_lens, decode_lens, output_size=actual_expanded
388
)
389
390
# [0, 3, 4, 8] -> [0, 0, 0, 3, 4, 4, 4, 4]
391
expanded_starts = torch.repeat_interleave(
392
- common_attn_metadata.query_start_loc[:num_decodes], decode_lens
392
+ common_attn_metadata.query_start_loc[:num_decodes],
393
+ decode_lens,
394
+ output_size=actual_expanded,
395
)
396
397
# [0, 1, 2, 0, 0, 1, 2, 3]
@@ -407,7 +409,9 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
409
# Give each of the flattened entries the same block table row as the
410
# original request.
411
self.expanded_block_table_buffer[:actual_expanded] = (
410
- torch.repeat_interleave(block_table, decode_lens, dim=0)
412
+ torch.repeat_interleave(
413
+ block_table, decode_lens, dim=0, output_size=actual_expanded
414
+ )
415
)
416
if actual_expanded < num_decode_tokens:
417
self.expanded_block_table_buffer[