From 5bec0b0ba3d313ac8b36e29517bdc938417d80ca Mon Sep 17 00:00:00 2001 From: Benjamin Chislett Date: Wed, 11 Mar 2026 00:16:56 -0400 Subject: [PATCH] [DSV3.2][MTP] Optimize Indexer MTP handling (#36723) Signed-off-by: Benjamin Chislett (cherry picked from commit 9040cd40af6bacfd20d1db8637a189f966a2fcc4) --- vllm/v1/attention/backends/mla/indexer.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index e84312970..a0ab157f5 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -372,12 +372,14 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): # [7, 6, 8, 0] -> [7, 7, 7, 6, 8, 8, 8, 8] expanded_base = torch.repeat_interleave( - seq_lens - decode_lens, decode_lens + seq_lens - decode_lens, decode_lens, output_size=actual_expanded ) # [0, 3, 4, 8] -> [0, 0, 0, 3, 4, 4, 4, 4] expanded_starts = torch.repeat_interleave( - common_attn_metadata.query_start_loc[:num_decodes], decode_lens + common_attn_metadata.query_start_loc[:num_decodes], + decode_lens, + output_size=actual_expanded, ) # [0, 1, 2, 0, 0, 1, 2, 3] @@ -395,7 +397,9 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): # Give each of the flattened entries the same block table row as the # original request. self.expanded_block_table_buffer[:actual_expanded] = ( - torch.repeat_interleave(block_table, decode_lens, dim=0) + torch.repeat_interleave( + block_table, decode_lens, dim=0, output_size=actual_expanded + ) ) if actual_expanded < num_decode_tokens: self.expanded_block_table_buffer[