diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py index 7dbd81125325..5ceb6c7558d6 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py @@ -54,6 +54,18 @@ def load_model(self, target_model: nn.Module) -> None: device=self.device, ) + # Lifecycle hooks for model-specific optimizations. Subclasses override + # the ones they need. These fire in both `capture` and `propose` so that + # any state they toggle (e.g. attention flags baked into a CUDA graph) is + # identical at capture time and replay time. + def on_prefill_begin(self, num_reqs: int) -> None: ... + + def on_prefill_end(self, num_reqs: int) -> None: ... + + def on_multi_step_decode_begin(self, num_reqs: int) -> None: ... + + def on_multi_step_decode_end(self, num_reqs: int) -> None: ... + @property def advance_draft_positions(self) -> bool: """ @@ -104,6 +116,8 @@ def capture(self) -> None: assert self.prefill_cudagraph_manager is not None if self.prefill_cudagraph_manager.use_breakable_cg: self.prefill_cudagraph_manager.init_breakable_cg_runner(self.model) + + self.on_prefill_begin(self.max_num_reqs) self.prefill_cudagraph_manager.capture( self._prefill, self.model_state, @@ -113,10 +127,12 @@ def capture(self) -> None: self.kv_cache_config, progress_bar_desc="Capturing prefill CUDA graphs", ) + self.on_prefill_end(self.max_num_reqs) if self.num_speculative_steps == 1: return + self.on_multi_step_decode_begin(self.max_num_reqs) # Capture the decode draft generation routine (model forward + # sample + update_draft_inputs) for a single # step. @@ -130,6 +146,7 @@ def capture(self) -> None: self.kv_cache_config, progress_bar_desc="Capturing decode CUDA graphs", ) + self.on_multi_step_decode_end(self.max_num_reqs) @torch.inference_mode() def propose( @@ -159,7 +176,8 @@ def propose( mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, is_profile: bool = False, ) -> torch.Tensor: - num_tokens = input_batch.num_tokens_after_padding + num_tokens = input_batch.num_tokens + num_tokens_padded = input_batch.num_tokens_after_padding num_reqs = input_batch.num_reqs max_query_len = input_batch.num_scheduled_tokens.max() max_seq_len = input_batch.seq_lens_cpu_upper_bound[:num_reqs].max().item() @@ -180,7 +198,7 @@ def propose( ) else: hidden_states = last_hidden_states - self.hidden_states[:num_tokens].copy_(hidden_states) + self.hidden_states[:num_tokens_padded].copy_(hidden_states) self._copy_request_inputs( num_reqs, @@ -208,21 +226,22 @@ def propose( num_reqs, # Use the actual number of tokens without padding added by # the target model during FULL cudagraph. - input_batch.num_tokens, + num_tokens, max_query_len, ) prefill_batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp( self.prefill_cudagraph_manager, num_reqs, - num_tokens, + num_tokens_padded, uniform_token_count, dp_size=self.dp_size, dp_rank=self.dp_rank, need_eager=is_profile, ) - self._prepare_eplb_forward(input_batch.num_tokens) + self._prepare_eplb_forward(num_tokens) + self.on_prefill_begin(num_reqs) if prefill_batch_desc.cg_mode == CUDAGraphMode.FULL: # Replay the full graph for draft prefill. assert self.prefill_cudagraph_manager is not None @@ -240,6 +259,7 @@ def propose( cudagraph_runtime_mode=prefill_batch_desc.cg_mode, mm_inputs=mm_inputs, ) + self.on_prefill_end(num_reqs) if self.num_speculative_steps == 1: # Early exit. @@ -268,6 +288,7 @@ def propose( need_eager=is_profile, ) + self.on_multi_step_decode_begin(num_reqs) # Generate the remaining num_speculative_steps - 1 draft tokens. self._multi_step_decode( num_reqs, @@ -276,6 +297,7 @@ def propose( num_tokens_across_dp, input_batch.seq_lens_cpu_upper_bound, ) + self.on_multi_step_decode_end(num_reqs) return self.draft_tokens[:num_reqs] diff --git a/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py b/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py index 4b9354f23e70..4889e9a750f8 100644 --- a/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py @@ -10,9 +10,49 @@ class MTPSpeculator(AutoRegressiveSpeculator): + share_mtp_topk_indices: bool = False + def load_draft_model( self, target_model: nn.Module, target_attn_layer_names: set[str], ) -> nn.Module: - return load_eagle_model(target_model, self.vllm_config) + draft_model = load_eagle_model(target_model, self.vllm_config) + spec_config = self.vllm_config.speculative_config + draft_hf_config = ( + spec_config.draft_model_config.hf_config + if spec_config is not None + else None + ) + # Detect index_share_for_mtp_iteration. When True, the proposer + # toggles skip_topk so step 0 computes MTP's own indices and + # steps 1+ reuse them. + self.share_mtp_topk_indices = ( + getattr(draft_hf_config, "index_share_for_mtp_iteration", False) + and hasattr(draft_model.model, "set_skip_topk") + and hasattr(draft_model.model, "compact_topk_indices") + ) + return draft_model + + def on_prefill_begin(self, num_reqs: int) -> None: + # Step 0 computes its own top-k. Unconditional, so a step that died + # midway cannot leave reuse mode on. + if self.share_mtp_topk_indices: + self.model.model.set_skip_topk(False) + + def on_prefill_end(self, num_reqs: int) -> None: + # Step 0 (prefill) wrote topk indices for every query token in the + # multi-token batch. Compact them down to each request's last token so + # steps 1+ can reuse them from the shared buffer. + if self.share_mtp_topk_indices and self.num_speculative_steps > 1: + self.model.model.compact_topk_indices(self.last_token_indices[:num_reqs]) + + def on_multi_step_decode_begin(self, num_reqs: int) -> None: + # Switch to reuse mode so draft steps 1+ skip the indexer op and read + # the indices that step 0 wrote into the shared buffer. + if self.share_mtp_topk_indices: + self.model.model.set_skip_topk(True) + + def on_multi_step_decode_end(self, num_reqs: int) -> None: + if self.share_mtp_topk_indices: + self.model.model.set_skip_topk(False)