diff --git a/vllm/model_executor/models/step3p5_mtp.py b/vllm/model_executor/models/step3p5_mtp.py index d17b579c1a..9af8b26ce8 100644 --- a/vllm/model_executor/models/step3p5_mtp.py +++ b/vllm/model_executor/models/step3p5_mtp.py @@ -22,7 +22,11 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.sequence import IntermediateTensors from .step3p5 import Step3p5DecoderLayer -from .utils import get_spec_layer_idx_from_weight_name, maybe_prefix +from .utils import ( + get_draft_quant_config, + get_spec_layer_idx_from_weight_name, + maybe_prefix, +) logger = init_logger(__name__) @@ -152,8 +156,19 @@ class Step3p5AMultiTokenPredictor(nn.Module): class Step3p5MTP(nn.Module): + # Each MTP layer's shared_head carries its own LM head weights. + has_own_lm_head = True + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() + # The target config may be a multimodal wrapper (Step3.7); build the + # draft from its own text and quant configs. + assert vllm_config.speculative_config is not None + draft_quant_config = get_draft_quant_config(vllm_config) + vllm_config = vllm_config.with_hf_config( + vllm_config.speculative_config.draft_model_config.hf_config + ) + vllm_config.quant_config = draft_quant_config self.config = vllm_config.model_config.hf_config self.vllm_config = vllm_config self.model = Step3p5AMultiTokenPredictor( diff --git a/vllm/models/inkling/amd/mtp.py b/vllm/models/inkling/amd/mtp.py index f9cba37f3c..f2ed51c075 100644 --- a/vllm/models/inkling/amd/mtp.py +++ b/vllm/models/inkling/amd/mtp.py @@ -323,7 +323,11 @@ class InklingMTP(nn.Module, SupportsMultiModalEmbeddings): logits = logits[..., : self.logits_processor.org_vocab_size] return logits - def get_top_tokens(self, hidden_states: torch.Tensor) -> torch.Tensor: + def get_top_tokens( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor: """Greedy draft tokens via rank-local argmax + tiny (value, index) reduction — no full-vocab logits all-gather. The muP divisor is a positive scalar, so the argmax is invariant and the scaling is diff --git a/vllm/models/inkling/nvidia/mtp.py b/vllm/models/inkling/nvidia/mtp.py index f9cba37f3c..f2ed51c075 100644 --- a/vllm/models/inkling/nvidia/mtp.py +++ b/vllm/models/inkling/nvidia/mtp.py @@ -323,7 +323,11 @@ class InklingMTP(nn.Module, SupportsMultiModalEmbeddings): logits = logits[..., : self.logits_processor.org_vocab_size] return logits - def get_top_tokens(self, hidden_states: torch.Tensor) -> torch.Tensor: + def get_top_tokens( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor: """Greedy draft tokens via rank-local argmax + tiny (value, index) reduction — no full-vocab logits all-gather. The muP divisor is a positive scalar, so the argmax is invariant and the scaling is diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py index 06bed9391e..65289bb18e 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py @@ -13,10 +13,14 @@ from vllm.v1.worker.gpu.spec_decode.utils import get_pp_safe_draft_load_config def _should_share(eagle: nn.Module, flag: str, draft, target) -> bool: - """Share when the draft has no own copy, or its copy matches the target.""" - if not getattr(eagle, flag, False) or draft is None: + """Share unless the draft declares its own copy that differs from the target. + + A draft that declares its own copy but has no top-level one (e.g. MTP + heads stored per layer) keeps it. + """ + if not getattr(eagle, flag, False): return True - if target is None: + if draft is None or target is None: return False # torch.equal on GPU allocates a bool mask the size of the input. # Use the faster GPU path when there is plenty of headroom; diff --git a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py index 928498edf1..c7026acdb8 100644 --- a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py @@ -253,6 +253,17 @@ class MultiModuleMTPSpeculator(DraftModelSpeculator): ) return self.draft_tokens[:num_reqs] + # Each MTP module may carry its own LM head, selected by spec_step_idx. + def compute_draft_logits( + self, hidden_states: torch.Tensor, spec_step_idx: int + ) -> torch.Tensor: + return self.model.compute_logits(hidden_states, spec_step_idx=spec_step_idx) + + def get_draft_top_tokens( + self, hidden_states: torch.Tensor, spec_step_idx: int + ) -> torch.Tensor: + return self.model.get_top_tokens(hidden_states, spec_step_idx=spec_step_idx) + @torch.inference_mode() def _run_model( self, @@ -427,6 +438,7 @@ class MultiModuleMTPSpeculator(DraftModelSpeculator): self.seeds, self.current_draft_step, self.draft_logits, + spec_step_idx=step, ) self.draft_tokens[:num_reqs, step] = draft_tokens diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 17baddf0a8..20a3b29eab 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -403,6 +403,16 @@ class DraftModelSpeculator(BaseSpeculator): "(communication: O(2*tp_size) vs O(vocab_size))." ) + def compute_draft_logits( + self, hidden_states: torch.Tensor, spec_step_idx: int + ) -> torch.Tensor: + return self.model.compute_logits(hidden_states) + + def get_draft_top_tokens( + self, hidden_states: torch.Tensor, spec_step_idx: int + ) -> torch.Tensor: + return self.model.get_top_tokens(hidden_states) + def sample_draft( self, hidden_states: torch.Tensor, @@ -412,9 +422,13 @@ class DraftModelSpeculator(BaseSpeculator): seeds: torch.Tensor, draft_step: torch.Tensor, draft_logits: torch.Tensor | None, + spec_step_idx: int = 0, ) -> torch.Tensor: + if draft_logits is None and self.use_local_argmax_reduction: + return self.get_draft_top_tokens(hidden_states, spec_step_idx) + + logits = self.compute_draft_logits(hidden_states, spec_step_idx) if draft_logits is not None: - logits = self.model.compute_logits(hidden_states) sampler = ( gumbel_sample if self.draft_watermarker is None @@ -432,10 +446,7 @@ class DraftModelSpeculator(BaseSpeculator): logits_cache_col=draft_step, use_fp64=self.use_fp64_gumbel, ) - elif self.use_local_argmax_reduction: - return self.model.get_top_tokens(hidden_states) else: - logits = self.model.compute_logits(hidden_states) sampled = logits.argmax(dim=-1) self._maybe_predict_acceptance(logits, idx_mapping, draft_step) return sampled