@tduncan / Nvidia-Optimization-Stack-VLLM / commits / ab66606993

[bugfix]Indexer init skip and MTP TopK share for iteration (#45895)

Signed-off-by: JaredforReal <w13431838023@gmail.com>

Jared Wen committed Jun 19, 2026 at 09:57 UTC ab666069935c1f23e8ef56038b4659ac9e8f19f8

9 files changed +69 -30

vllm/model_executor/layers/attention/mla_attention.py

@@ -349,6 +349,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):

349 attn_backend: type[AttentionBackend] | None = None,

350 use_sparse: bool = False,

351 indexer: object | None = None,

352 + topk_indices_buffer: torch.Tensor | None = None,

353 **extra_impl_args,

354 ):

355 super().__init__()

@@ -437,6 +438,11 @@ class MLAAttention(nn.Module, AttentionLayerBase):

438 )

439 cache_config.enable_prefix_caching = False

440

441 + # Sparse MLA reads top-k indices from a shared buffer. Pass it

442 + # explicitly so backbone "skip" layers (indexer=None) still find it.

443 + if use_sparse:

444 + extra_impl_args["topk_indices_buffer"] = topk_indices_buffer

445 +

446 impl_cls = cast(type[MLAAttentionImpl], self.attn_backend.get_impl_cls())

447 self.impl = impl_cls( # type: ignore[assignment] # impl_cls always returns an MLAAttentionImpl subclass

448 num_heads=self.num_heads,

vllm/model_executor/layers/mla.py

@@ -112,6 +112,7 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):

112 kv_b_proj=self.kv_b_proj,

113 use_sparse=self.is_sparse,

114 indexer=self.indexer,

115 + topk_indices_buffer=mla_modules.topk_indices_buffer,

116 )

117

118 self.prefix = prefix

vllm/model_executor/models/deepseek_mtp.py

+6 -2

@@ -119,8 +119,12 @@ class DeepSeekMultiTokenPredictorLayer(nn.Module):

119 hidden_states=hidden_states,

120 residual=None,

121 )

122 - hidden_states = residual + hidden_states

123 - return hidden_states

122 + hidden_states = residual + hidden_states # pre-final-norm (logits hidden)

123 + # Recycle the post-final-norm hidden into the next draft step.

124 + # compute_logits applies shared_head (== final norm) to the pre-norm

125 + # element, so logits and the recycle each get exactly one final-norm.

126 + # Matches SGLang's deepseek_nextn.

127 + return hidden_states, self.shared_head(hidden_states)

128

129

130 class DeepSeekMultiTokenPredictor(nn.Module):

vllm/model_executor/models/deepseek_v2.py

+22 -17

@@ -998,8 +998,29 @@ class DeepseekV2MLAAttention(nn.Module):

998

999 self.is_v32 = hasattr(config, "index_topk")

1000

1001 + # IndexCache config

1002 + # Refer: https://arxiv.org/abs/2603.12201 for more details.

1003 _skip_topk = False

1002 - if self.is_v32:

1004 + _index_topk_freq = getattr(config, "index_topk_freq", 1)

1005 + _index_topk_pattern = getattr(config, "index_topk_pattern", None)

1006 + _index_skip_topk_offset = getattr(config, "index_skip_topk_offset", 2)

1007 + layer_id = extract_layer_index(prefix)

1008 +

1009 + if _index_topk_pattern is None:

1010 + _skip_topk = (

1011 + max(layer_id - _index_skip_topk_offset + 1, 0) % _index_topk_freq != 0

1012 + )

1013 + elif 0 <= layer_id < len(_index_topk_pattern):

1014 + _skip_topk = _index_topk_pattern[layer_id] == "S"

1015 +

1016 + # The skip pattern only governs backbone layers. MTP/nextn layers

1017 + # (layer_id >= num_hidden_layers) always build a full indexer: they

1018 + # compute indices at draft step 0 and toggle at runtime via

1019 + # set_skip_topk (index_share_for_mtp_iteration).

1020 + _num_hidden_layers = getattr(config, "num_hidden_layers", None)

1021 + is_mtp_layer = _num_hidden_layers is not None and layer_id >= _num_hidden_layers

1022 +

1023 + if self.is_v32 and (not _skip_topk or is_mtp_layer):

1024 self.indexer_rope_emb = get_rope(

1025 qk_rope_head_dim,

1026 max_position=max_position_embeddings,

@@ -1017,22 +1038,6 @@ class DeepseekV2MLAAttention(nn.Module):

1038 f"{prefix}.indexer",

1039 is_inplace_rope=self.indexer_rope_emb.enabled(),

1040 )

1020 -

1021 - # IndexCache config

1022 - # Refer: https://arxiv.org/abs/2603.12201 for more details.

1023 - _index_topk_freq = getattr(config, "index_topk_freq", 1)

1024 - _index_topk_pattern = getattr(config, "index_topk_pattern", None)

1025 - _index_skip_topk_offset = getattr(config, "index_skip_topk_offset", 2)

1026 - layer_id = extract_layer_index(prefix)

1027 -

1028 - if _index_topk_pattern is None:

1029 - _skip_topk = (

1030 - max(layer_id - _index_skip_topk_offset + 1, 0) % _index_topk_freq

1031 - != 0

1032 - )

1033 - elif 0 <= layer_id < len(_index_topk_pattern):

1034 - _skip_topk = _index_topk_pattern[layer_id] == "S"

1035 -

1041 else:

1042 self.indexer_rope_emb = None

1043 self.indexer = None

vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py

+7 -3

@@ -271,7 +271,7 @@ class FlashInferMLASparseImpl(SparseMLAAttentionImpl[FlashInferMLASparseMetadata

271 attn_type: str,

272 kv_sharing_target_layer_name: str | None,

273 # MLA Specific Arguments

274 - topk_indice_buffer: torch.Tensor | None = None,

274 + topk_indices_buffer: torch.Tensor | None = None,

275 indexer: "Indexer | None" = None,

276 **mla_args,

277 ) -> None:

@@ -301,8 +301,12 @@ class FlashInferMLASparseImpl(SparseMLAAttentionImpl[FlashInferMLASparseMetadata

301 self.qk_nope_head_dim: int = mla_args["qk_nope_head_dim"]

302 self.qk_rope_head_dim: int = mla_args["qk_rope_head_dim"]

303

304 - assert indexer is not None, "Indexer required for sparse MLA"

305 - self.topk_indices_buffer: torch.Tensor | None = indexer.topk_indices_buffer

304 + # The indexer carries the shared buffer for normal layers and tests;

305 + # the explicitly-passed buffer covers backbone skip layers, whose

306 + # indexer is not constructed (see deepseek_v2.py).

307 + self.topk_indices_buffer: torch.Tensor | None = (

308 + indexer.topk_indices_buffer if indexer is not None else topk_indices_buffer

309 + )

310

311 self._workspace_buffer: torch.Tensor | None = None

312 self.bmm1_scale: float | None = None

vllm/v1/attention/backends/mla/flashmla_sparse.py

+6 -2

@@ -568,8 +568,12 @@ class FlashMLASparseImpl(SparseMLAAttentionImpl[FlashMLASparseMetadata]):

568 self.kv_cache_dtype = kv_cache_dtype

569 self.kv_lora_rank: int = mla_args["kv_lora_rank"]

570 self.softmax_scale = scale

571 - assert indexer is not None

572 - self.topk_indices_buffer: torch.Tensor | None = indexer.topk_indices_buffer

571 + # The indexer carries the shared buffer for normal layers and tests;

572 + # the explicitly-passed buffer covers backbone skip layers, whose

573 + # indexer is not constructed (see deepseek_v2.py).

574 + self.topk_indices_buffer: torch.Tensor | None = (

575 + indexer.topk_indices_buffer if indexer is not None else topk_indices_buffer

576 + )

577 # Prefill BF16 kernel requires 64 on Hopper, 128 on Blackwell

578 self.prefill_padding = (

579 128 if current_platform.is_device_capability_family(100) else 64

vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py

+7 -3

@@ -629,7 +629,7 @@ class ROCMAiterMLASparseImpl(SparseMLAAttentionImpl[ROCMAiterMLASparseMetadata])

629 attn_type: str,

630 kv_sharing_target_layer_name: str | None,

631 # MLA Specific Arguments

632 - topk_indice_buffer: torch.Tensor | None = None,

632 + topk_indices_buffer: torch.Tensor | None = None,

633 indexer: "Indexer | None" = None,

634 **mla_args,

635 ) -> None:

@@ -642,8 +642,12 @@ class ROCMAiterMLASparseImpl(SparseMLAAttentionImpl[ROCMAiterMLASparseMetadata])

642 self.kv_cache_dtype = kv_cache_dtype

643 self.kv_lora_rank: int = mla_args["kv_lora_rank"]

644 self.softmax_scale = scale

645 - assert indexer is not None

646 - self.topk_indices_buffer: torch.Tensor | None = indexer.topk_indices_buffer

645 + # The indexer carries the shared buffer for normal layers and tests;

646 + # the explicitly-passed buffer covers backbone skip layers, whose

647 + # indexer is not constructed (see deepseek_v2.py).

648 + self.topk_indices_buffer: torch.Tensor | None = (

649 + indexer.topk_indices_buffer if indexer is not None else topk_indices_buffer

650 + )

651

652 vllm_config = get_current_vllm_config()

653 max_tokens = vllm_config.scheduler_config.max_num_batched_tokens

vllm/v1/attention/backends/mla/xpu_mla_sparse.py

+7 -3

@@ -184,7 +184,7 @@ class XPUMLASparseImpl(SparseMLAAttentionImpl[XPUMLASparseMetadata]):

184 attn_type: str,

185 kv_sharing_target_layer_name: str | None,

186 # MLA Specific Arguments

187 - topk_indice_buffer: torch.Tensor | None = None,

187 + topk_indices_buffer: torch.Tensor | None = None,

188 indexer: Optional["Indexer"] = None,

189 **mla_args,

190 ) -> None:

@@ -195,8 +195,12 @@ class XPUMLASparseImpl(SparseMLAAttentionImpl[XPUMLASparseMetadata]):

195 self.kv_cache_dtype = kv_cache_dtype

196 self.kv_lora_rank: int = mla_args["kv_lora_rank"]

197 self.softmax_scale = scale

198 - assert indexer is not None

199 - self.topk_indices_buffer: torch.Tensor | None = indexer.topk_indices_buffer

198 + # The indexer carries the shared buffer for normal layers and tests;

199 + # the explicitly-passed buffer covers backbone skip layers, whose

200 + # indexer is not constructed (see deepseek_v2.py).

201 + self.topk_indices_buffer: torch.Tensor | None = (

202 + indexer.topk_indices_buffer if indexer is not None else topk_indices_buffer

203 + )

204

205 def _forward_bf16_kv(

206 self,

vllm/v1/spec_decode/llm_base_proposer.py

@@ -918,6 +918,13 @@ class SpecDecodeBaseProposer:

918 return per_group_attn_metadata, per_layer_attn_metadata

919

920 def model_returns_tuple(self) -> bool:

921 + if self.method == "mtp":

922 + # DeepSeek-family MTP (deepseek_mtp.py) recycles the post-final-

923 + # norm hidden, so its forward returns (logit_hidden,

924 + # recycle_hidden). Other MTP families return a single tensor.

925 + return "DeepSeekMTPModel" in (

926 + self.draft_model_config.hf_config.architectures or []

927 + )

928 return self.method not in ("mtp", "draft_model", "dflash")

929

930 def prepare_next_token_ids_cpu(