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 1, 18:45 UTC
[feature] add index share feature for DSA MTP (#44420)
Signed-off-by: JaredforReal <w13431838023@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Jared Wen committed
Jun 7, 2026 at 13:04 UTC
32f34d393524d18c7ec9c4add0be7731cde372dc
5 files changed
+114
-25
vllm/model_executor/models/deepseek_mtp.py
+24
-2
@@ -115,7 +115,9 @@ class DeepSeekMultiTokenPredictorLayer(nn.Module):
115
)
116
117
hidden_states, residual = self.mtp_block(
118
- positions=positions, hidden_states=hidden_states, residual=None
118
+ positions=positions,
119
+ hidden_states=hidden_states,
120
+ residual=None,
121
)
122
hidden_states = residual + hidden_states
123
return hidden_states
@@ -147,6 +149,22 @@ class DeepSeekMultiTokenPredictor(nn.Module):
149
)
150
self.logits_processor = LogitsProcessor(config.vocab_size)
151
152
+ def set_skip_topk(self, skip: bool):
153
+ """Toggle skip_topk on all MTP layers with sparse attention.
154
+
155
+ Called by the proposer to implement index_share_for_mtp_iteration:
156
+ step 0 sets skip=False (compute own indices), steps 1+ set skip=True
157
+ (reuse step 0's indices).
158
+ """
159
+ for layer in self.layers.values():
160
+ mtp_block = getattr(layer, "mtp_block", None)
161
+ if mtp_block is not None:
162
+ self_attn = getattr(mtp_block, "self_attn", None)
163
+ if self_attn is not None:
164
+ mla_attn = getattr(self_attn, "mla_attn", None)
165
+ if mla_attn is not None and hasattr(mla_attn, "skip_topk"):
166
+ mla_attn.skip_topk = skip
167
+
168
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
169
return self.embed_tokens(input_ids)
170
@@ -225,7 +243,11 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
243
spec_step_idx: int = 0,
244
) -> torch.Tensor:
245
hidden_states = self.model(
228
- input_ids, positions, hidden_states, inputs_embeds, spec_step_idx
246
+ input_ids,
247
+ positions,
248
+ hidden_states,
249
+ inputs_embeds,
250
+ spec_step_idx,
251
)
252
return hidden_states
253
vllm/model_executor/models/deepseek_v2.py
+16
-15
@@ -1018,19 +1018,20 @@ class DeepseekV2MLAAttention(nn.Module):
1018
is_inplace_rope=self.indexer_rope_emb.enabled(),
1019
)
1020
1021
- # Enable IndexCache for DeepSeek models to reduce redundant top-k
1022
- # token selection computations in sparse attention.
1023
- use_index_cache = getattr(config, "use_index_cache", False)
1024
- if use_index_cache:
1025
- # IndexCache config
1026
- # Refer: https://arxiv.org/abs/2603.12201 for more details.
1027
- _index_topk_freq = getattr(config, "index_topk_freq", 1)
1028
- _index_topk_pattern = getattr(config, "index_topk_pattern", None)
1029
- layer_id = extract_layer_index(prefix)
1030
- if _index_topk_pattern is None:
1031
- _skip_topk = max(layer_id - 1, 0) % _index_topk_freq != 0
1032
- elif 0 <= layer_id < len(_index_topk_pattern):
1033
- _skip_topk = _index_topk_pattern[layer_id] == "S"
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
1036
else:
1037
self.indexer_rope_emb = None
@@ -1252,8 +1253,8 @@ class DeepseekV2Model(nn.Module):
1253
self.start_layer, self.end_layer, self.layers = make_layers(
1254
config.num_hidden_layers,
1255
lambda prefix: DeepseekV2DecoderLayer(
1255
- vllm_config,
1256
- prefix,
1256
+ vllm_config=vllm_config,
1257
+ prefix=prefix,
1258
topk_indices_buffer=topk_indices_buffer,
1259
),
1260
prefix=f"{prefix}.layers",
vllm/transformers_utils/model_arch_config_convertor.py
+33
-1
@@ -50,7 +50,7 @@ class ModelArchConfigConvertorBase:
50
# special case for deepseek_v4
51
if hasattr(self.hf_text_config, "compress_ratios"):
52
return self.hf_text_config.head_dim
53
- qk_rope_head_dim = getattr(self.hf_text_config, "qk_rope_head_dim", 0)
53
+ qk_rope_head_dim = self._get_qk_rope_head_dim()
54
if not envs.VLLM_MLA_DISABLE:
55
return self.hf_text_config.kv_lora_rank + qk_rope_head_dim
56
else:
@@ -71,6 +71,38 @@ class ModelArchConfigConvertorBase:
71
# FIXME(woosuk): This may not be true for all models.
72
return self.get_hidden_size() // total_num_attention_heads
73
74
+ def _get_qk_rope_head_dim(self) -> int:
75
+ """Get qk_rope_head_dim, fixing the transformers v5.4+ attribute_map bug."""
76
+ cfg = self.hf_text_config
77
+ qk_rope_head_dim = getattr(cfg, "qk_rope_head_dim", 0)
78
+ qk_nope_head_dim = getattr(cfg, "qk_nope_head_dim", 0)
79
+
80
+ # In valid MLA configs, qk_rope_head_dim != qk_nope_head_dim.
81
+ if qk_rope_head_dim == 0 or qk_rope_head_dim != qk_nope_head_dim:
82
+ return qk_rope_head_dim # not corrupted
83
+
84
+ # Read the correct value from raw config.json.
85
+ from vllm.transformers_utils.repo_utils import get_hf_file_to_dict
86
+
87
+ model_path = self.hf_config.name_or_path
88
+ if not model_path:
89
+ return qk_rope_head_dim
90
+ raw = get_hf_file_to_dict("config.json", model_path)
91
+ if raw and "qk_rope_head_dim" in raw:
92
+ correct = raw["qk_rope_head_dim"]
93
+ if correct != qk_rope_head_dim:
94
+ logger.info(
95
+ "Fixing qk_rope_head_dim: %d -> %d "
96
+ "(transformers v5.4+ attribute_map bug)",
97
+ qk_rope_head_dim,
98
+ correct,
99
+ )
100
+ # Patch the config so downstream model layers also get
101
+ # the correct value.
102
+ cfg.qk_rope_head_dim = correct
103
+ return correct
104
+ return qk_rope_head_dim
105
+
106
def get_total_num_kv_heads(self) -> int:
107
attributes = [
108
# For Falcon:
vllm/v1/spec_decode/llm_base_proposer.py
+32
-3
@@ -70,6 +70,7 @@ class SpecDecodeBaseProposer:
70
self.draft_model_config = self.speculative_config.draft_model_config
71
self.method = self.speculative_config.method
72
self.pass_hidden_states_to_model = pass_hidden_states_to_model
73
+ self._share_mtp_indices = False
74
75
self.device = device
76
self.dtype = vllm_config.model_config.dtype
@@ -490,6 +491,11 @@ class SpecDecodeBaseProposer:
491
model_kwargs, slot_mapping_size = self.build_model_inputs_first_pass(
492
num_tokens, num_input_tokens, mm_embed_inputs
493
)
494
+ # Step 0 of index_share_for_mtp_iteration: let the MTP layer
495
+ # compute its own indices (skip_topk=False) so subsequent steps
496
+ # can reuse them.
497
+ if self._share_mtp_indices and hasattr(self.model.model, "set_skip_topk"):
498
+ self.model.model.set_skip_topk(False)
499
500
with set_forward_context(
501
per_layer_attn_metadata,
@@ -508,6 +514,11 @@ class SpecDecodeBaseProposer:
514
else:
515
last_hidden_states, hidden_states = ret_hidden_states
516
517
+ # After step 0: switch to reuse mode so steps 1+ skip the indexer
518
+ # and read the indices that step 0 just wrote into the shared buffer.
519
+ if self._share_mtp_indices and hasattr(self.model.model, "set_skip_topk"):
520
+ self.model.model.set_skip_topk(True)
521
+
522
sample_hidden_states = last_hidden_states[token_indices_to_sample]
523
524
# Early exit if there is only one draft token to be generated.
@@ -1418,16 +1429,34 @@ class SpecDecodeBaseProposer:
1429
)
1430
1431
if hasattr(target_language_model.model, "topk_indices_buffer"):
1432
+ target_buffer = target_language_model.model.topk_indices_buffer
1433
if hasattr(self.model.model, "topk_indices_buffer"):
1434
del self.model.model.topk_indices_buffer
1423
- self.model.model.topk_indices_buffer = (
1424
- target_language_model.model.topk_indices_buffer
1425
- )
1435
+ self.model.model.topk_indices_buffer = target_buffer
1436
+ # Also share at per-module level so that the indexer and
1437
+ # sparse-attention backends in each MTP layer read from
1438
+ # the target model's buffer.
1439
+ for _, module in self.model.model.named_modules():
1440
+ if hasattr(module, "topk_indices_buffer"):
1441
+ module.topk_indices_buffer = target_buffer
1442
logger.info(
1443
"Detected MTP model with topk_indices_buffer. "
1444
"Sharing target model topk_indices_buffer with the draft model."
1445
)
1446
1447
+ # Detect index_share_for_mtp_iteration: when True, the proposer
1448
+ # toggles skip_topk so step 0 computes MTP's own indices and
1449
+ # steps 1+ reuse them.
1450
+ spec_config = self.vllm_config.speculative_config
1451
+ draft_hf_config = (
1452
+ spec_config.draft_model_config.hf_config
1453
+ if spec_config is not None
1454
+ else None
1455
+ )
1456
+ self._share_mtp_indices = getattr(
1457
+ draft_hf_config, "index_share_for_mtp_iteration", False
1458
+ )
1459
+
1460
if self.use_local_argmax_reduction:
1461
if not hasattr(self.model, "get_top_tokens"):
1462
raise ValueError(
vllm/v1/worker/gpu/spec_decode/eagle/utils.py
+9
-4
@@ -76,10 +76,15 @@ def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mod
76
del sh.head
77
sh.head = target_lm_head
78
79
- # MTP also shares a topk_indices_buffer between target and draft.
79
+ # MTP shares topk_indices_buffer with the target model. We update
80
+ # every module in the draft that holds a buffer reference so that
81
+ # the per-layer indexer and sparse-attention backends all point to
82
+ # the target's buffer.
83
if hasattr(target_inner, "topk_indices_buffer"):
81
- if hasattr(draft_inner, "topk_indices_buffer"):
82
- del draft_inner.topk_indices_buffer
83
- draft_inner.topk_indices_buffer = target_inner.topk_indices_buffer
84
+ target_buffer = target_inner.topk_indices_buffer
85
+ if target_buffer is not None:
86
+ for _, module in draft_inner.named_modules():
87
+ if hasattr(module, "topk_indices_buffer"):
88
+ module.topk_indices_buffer = target_buffer
89
90
return eagle_model