[Model Runner V2][MTP] Share topk index buffer between draft steps (#47352)
Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai> Co-authored-by: Woosuk Kwon <woosuk@inferact.ai>
Giancarlo Delfin committed
Aug 10, 2026 at 15:46 UTC
3e174bb73c70b677dd339c59c39d7460abc5d02d
2 files changed
+68
-6
vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py
+27
-5
@@ -54,6 +54,18 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
54
device=self.device,
55
)
56
57
+ # Lifecycle hooks for model-specific optimizations. Subclasses override
58
+ # the ones they need. These fire in both `capture` and `propose` so that
59
+ # any state they toggle (e.g. attention flags baked into a CUDA graph) is
60
+ # identical at capture time and replay time.
61
+ def on_prefill_begin(self, num_reqs: int) -> None: ...
62
+
63
+ def on_prefill_end(self, num_reqs: int) -> None: ...
64
+
65
+ def on_multi_step_decode_begin(self, num_reqs: int) -> None: ...
66
+
67
+ def on_multi_step_decode_end(self, num_reqs: int) -> None: ...
68
+
69
@property
70
def advance_draft_positions(self) -> bool:
71
"""
@@ -104,6 +116,8 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
116
assert self.prefill_cudagraph_manager is not None
117
if self.prefill_cudagraph_manager.use_breakable_cg:
118
self.prefill_cudagraph_manager.init_breakable_cg_runner(self.model)
119
+
120
+ self.on_prefill_begin(self.max_num_reqs)
121
self.prefill_cudagraph_manager.capture(
122
self._prefill,
123
self.model_state,
@@ -113,10 +127,12 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
127
self.kv_cache_config,
128
progress_bar_desc="Capturing prefill CUDA graphs",
129
)
130
+ self.on_prefill_end(self.max_num_reqs)
131
132
if self.num_speculative_steps == 1:
133
return
134
135
+ self.on_multi_step_decode_begin(self.max_num_reqs)
136
# Capture the decode draft generation routine (model forward +
137
# sample + update_draft_inputs) for a single
138
# step.
@@ -130,6 +146,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
146
self.kv_cache_config,
147
progress_bar_desc="Capturing decode CUDA graphs",
148
)
149
+ self.on_multi_step_decode_end(self.max_num_reqs)
150
151
@torch.inference_mode()
152
def propose(
@@ -159,7 +176,8 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
176
mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
177
is_profile: bool = False,
178
) -> torch.Tensor:
162
- num_tokens = input_batch.num_tokens_after_padding
179
+ num_tokens = input_batch.num_tokens
180
+ num_tokens_padded = input_batch.num_tokens_after_padding
181
num_reqs = input_batch.num_reqs
182
max_query_len = input_batch.num_scheduled_tokens.max()
183
max_seq_len = input_batch.seq_lens_cpu_upper_bound[:num_reqs].max().item()
@@ -180,7 +198,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
198
)
199
else:
200
hidden_states = last_hidden_states
183
- self.hidden_states[:num_tokens].copy_(hidden_states)
201
+ self.hidden_states[:num_tokens_padded].copy_(hidden_states)
202
203
self._copy_request_inputs(
204
num_reqs,
@@ -208,21 +226,22 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
226
num_reqs,
227
# Use the actual number of tokens without padding added by
228
# the target model during FULL cudagraph.
211
- input_batch.num_tokens,
229
+ num_tokens,
230
max_query_len,
231
)
232
prefill_batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp(
233
self.prefill_cudagraph_manager,
234
num_reqs,
217
- num_tokens,
235
+ num_tokens_padded,
236
uniform_token_count,
237
dp_size=self.dp_size,
238
dp_rank=self.dp_rank,
239
need_eager=is_profile,
240
)
241
224
- self._prepare_eplb_forward(input_batch.num_tokens)
242
+ self._prepare_eplb_forward(num_tokens)
243
244
+ self.on_prefill_begin(num_reqs)
245
if prefill_batch_desc.cg_mode == CUDAGraphMode.FULL:
246
# Replay the full graph for draft prefill.
247
assert self.prefill_cudagraph_manager is not None
@@ -240,6 +259,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
259
cudagraph_runtime_mode=prefill_batch_desc.cg_mode,
260
mm_inputs=mm_inputs,
261
)
262
+ self.on_prefill_end(num_reqs)
263
264
if self.num_speculative_steps == 1:
265
# Early exit.
@@ -268,6 +288,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
288
need_eager=is_profile,
289
)
290
291
+ self.on_multi_step_decode_begin(num_reqs)
292
# Generate the remaining num_speculative_steps - 1 draft tokens.
293
self._multi_step_decode(
294
num_reqs,
@@ -276,6 +297,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
297
num_tokens_across_dp,
298
input_batch.seq_lens_cpu_upper_bound,
299
)
300
+ self.on_multi_step_decode_end(num_reqs)
301
302
return self.draft_tokens[:num_reqs]
303
vllm/v1/worker/gpu/spec_decode/mtp/speculator.py
+41
-1
@@ -10,9 +10,49 @@ from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model
10
11
12
class MTPSpeculator(AutoRegressiveSpeculator):
13
+ share_mtp_topk_indices: bool = False
14
+
15
def load_draft_model(
16
self,
17
target_model: nn.Module,
18
target_attn_layer_names: set[str],
19
) -> nn.Module:
18
- return load_eagle_model(target_model, self.vllm_config)
20
+ draft_model = load_eagle_model(target_model, self.vllm_config)
21
+ spec_config = self.vllm_config.speculative_config
22
+ draft_hf_config = (
23
+ spec_config.draft_model_config.hf_config
24
+ if spec_config is not None
25
+ else None
26
+ )
27
+ # Detect index_share_for_mtp_iteration. When True, the proposer
28
+ # toggles skip_topk so step 0 computes MTP's own indices and
29
+ # steps 1+ reuse them.
30
+ self.share_mtp_topk_indices = (
31
+ getattr(draft_hf_config, "index_share_for_mtp_iteration", False)
32
+ and hasattr(draft_model.model, "set_skip_topk")
33
+ and hasattr(draft_model.model, "compact_topk_indices")
34
+ )
35
+ return draft_model
36
+
37
+ def on_prefill_begin(self, num_reqs: int) -> None:
38
+ # Step 0 computes its own top-k. Unconditional, so a step that died
39
+ # midway cannot leave reuse mode on.
40
+ if self.share_mtp_topk_indices:
41
+ self.model.model.set_skip_topk(False)
42
+
43
+ def on_prefill_end(self, num_reqs: int) -> None:
44
+ # Step 0 (prefill) wrote topk indices for every query token in the
45
+ # multi-token batch. Compact them down to each request's last token so
46
+ # steps 1+ can reuse them from the shared buffer.
47
+ if self.share_mtp_topk_indices and self.num_speculative_steps > 1:
48
+ self.model.model.compact_topk_indices(self.last_token_indices[:num_reqs])
49
+
50
+ def on_multi_step_decode_begin(self, num_reqs: int) -> None:
51
+ # Switch to reuse mode so draft steps 1+ skip the indexer op and read
52
+ # the indices that step 0 wrote into the shared buffer.
53
+ if self.share_mtp_topk_indices:
54
+ self.model.model.set_skip_topk(True)
55
+
56
+ def on_multi_step_decode_end(self, num_reqs: int) -> None:
57
+ if self.share_mtp_topk_indices:
58
+ self.model.model.set_skip_topk(False)