diff --git a/tests/v1/spec_decode/test_dsv4_mtp_local_argmax.py b/tests/v1/spec_decode/test_dsv4_mtp_local_argmax.py new file mode 100644 index 0000000000..2c508864c6 --- /dev/null +++ b/tests/v1/spec_decode/test_dsv4_mtp_local_argmax.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""``get_top_tokens`` on the DeepSeek V4 MTP drafter. + +The speculator opts into ``use_local_argmax_reduction`` by probing the draft +model for ``get_top_tokens``, and expects it to select the same greedy draft +token the full-vocab argmax would. V4 reaches its LM head through hc_head and +the shared head's RMSNorm, so both methods have to share that transform. +""" + +from types import SimpleNamespace +from unittest import mock + +import pytest +import torch +from transformers import PreTrainedConfig + +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.mhc import HCHeadOp +from vllm.model_executor.models.deepseek_mtp import SharedHead +from vllm.platforms import current_platform + +if not current_platform.is_rocm(): + pytest.skip( + "get_top_tokens is implemented on the ROCm DeepSeek V4 MTP drafter only.", + allow_module_level=True, + ) + +from vllm.models.deepseek_v4.amd.mtp import ( # noqa: E402 + DeepSeekV4MTP, + DeepSeekV4MultiTokenPredictor, +) + +DEVICE_TYPE = current_platform.device_type +START_LAYER_IDX = 61 +HIDDEN_SIZE = 4096 +HC_MULT = 4 +RMS_NORM_EPS = 1e-6 +HC_EPS = 1e-6 +_EMBEDDING = "vllm.model_executor.layers.vocab_parallel_embedding" + + +def _make_predictor(vocab_size: int) -> DeepSeekV4MultiTokenPredictor: + config = PreTrainedConfig() + config.hidden_size = HIDDEN_SIZE + config.vocab_size = vocab_size + config.rms_norm_eps = RMS_NORM_EPS + + # A single vocab shard, so the reduction returns the local argmax without + # touching the process group. + with ( + mock.patch(f"{_EMBEDDING}.get_tensor_model_parallel_rank", return_value=0), + mock.patch( + f"{_EMBEDDING}.get_tensor_model_parallel_world_size", return_value=1 + ), + ): + shared_head = SharedHead(config=config, prefix="shared_head") + + generator = torch.Generator().manual_seed(7) + with torch.no_grad(): + shared_head.head.weight.normal_(0.0, 0.02, generator=generator) + # The default norm weight is all ones, and the argmax is invariant to + # the norm's per-row rescale, so a uniform gain would hide a skipped + # shared_head. Vary it so the comparison below can see one. + shared_head.norm.weight.normal_(1.0, 0.5, generator=generator) + shared_head = shared_head.to(device=DEVICE_TYPE, dtype=torch.bfloat16) + + hc_head_fn = torch.randn( + HC_MULT, HC_MULT * HIDDEN_SIZE, dtype=torch.float32, generator=generator + ) + # Non-trivial gate scale and bias: hc_head mixes the hc_mult streams with + # these, and defaults of 0 would make a skipped hc_head harder to see. + hc_head_scale = torch.randn(1, dtype=torch.float32, generator=generator) * 0.1 + hc_head_base = torch.randn(HC_MULT, dtype=torch.float32, generator=generator) * 0.1 + + mtp_layer = SimpleNamespace( + config=config, + hc_mult=HC_MULT, + hc_eps=HC_EPS, + rms_norm_eps=RMS_NORM_EPS, + hc_head_op=HCHeadOp(), + hc_head_fn=(hc_head_fn * 1e-4).to(DEVICE_TYPE), + hc_head_scale=hc_head_scale.to(DEVICE_TYPE), + hc_head_base=hc_head_base.to(DEVICE_TYPE), + shared_head=shared_head, + ) + + # Only the attributes the two methods read; a real predictor would also + # build the MLA/MoE draft block, which is irrelevant here. + predictor = object.__new__(DeepSeekV4MultiTokenPredictor) + torch.nn.Module.__init__(predictor) + predictor.layers = {str(START_LAYER_IDX): mtp_layer} + predictor.mtp_start_layer_idx = START_LAYER_IDX + predictor.num_mtp_layers = 1 + predictor.logits_processor = LogitsProcessor(vocab_size) + return predictor + + +def _hidden_states(num_tokens: int, seed: int) -> torch.Tensor: + """The flat pre-hc_head residual the MTP layer returns.""" + generator = torch.Generator().manual_seed(seed) + hidden_states = torch.randn( + num_tokens, HC_MULT * HIDDEN_SIZE, dtype=torch.float32, generator=generator + ) + return hidden_states.to(device=DEVICE_TYPE, dtype=torch.bfloat16) + + +def _assert_selects_max_logit(top: torch.Tensor, logits: torch.Tensor) -> None: + """The selected token must carry the row's maximum logit. + + hc_head emits bf16, so rows can tie for the maximum and then no single + index is the right answer; compare the values the two paths select. + """ + assert top.shape == (logits.shape[0],) + selected = logits.gather(1, top.unsqueeze(1)).squeeze(1) + assert torch.equal(selected, logits.max(dim=-1).values) + + +def test_get_top_tokens_matches_full_argmax(default_vllm_config): + """Same draft token as the path taken when the reduction is off. + + ``get_top_tokens`` has to run the same pre-head transform as + ``compute_logits``: hc_head over the (T, hc_mult, D) residual, then the + shared head's RMSNorm. Dropping either one returns a different draft token + instead of failing, which is what this comparison catches. + """ + predictor = _make_predictor(vocab_size=512) + + for num_tokens in (1, 16, 128): + hidden_states = _hidden_states(num_tokens, seed=11) + + logits = predictor.compute_logits(hidden_states) + top = predictor.get_top_tokens(hidden_states) + + assert top.dtype == torch.int64 + _assert_selects_max_logit(top, logits) + + +def test_get_top_tokens_honors_padded_vocab(default_vllm_config): + """Padding entries above ``org_vocab_size`` must never win the argmax.""" + vocab_size = 500 + predictor = _make_predictor(vocab_size) + head = predictor.layers[str(START_LAYER_IDX)].shared_head.head + assert head.shard_indices.num_org_vocab_padding > 0 + + hidden_states = _hidden_states(32, seed=13) + logits = predictor.compute_logits(hidden_states) + top = predictor.get_top_tokens(hidden_states) + + _assert_selects_max_logit(top, logits) + assert (top < vocab_size).all() + + +def test_wrapper_forwards_spec_step_idx(): + """``DeepSeekV4MTP`` is what the speculator probes, so it needs the method.""" + assert hasattr(DeepSeekV4MTP, "get_top_tokens") + + hidden_states = torch.empty(0) + wrapper = object.__new__(DeepSeekV4MTP) + wrapper.model = SimpleNamespace(get_top_tokens=lambda *args: args) + + assert DeepSeekV4MTP.get_top_tokens(wrapper, hidden_states, 2) == (hidden_states, 2) + assert DeepSeekV4MTP.get_top_tokens(wrapper, hidden_states) == (hidden_states, 0) diff --git a/vllm/models/deepseek_v4/amd/mtp.py b/vllm/models/deepseek_v4/amd/mtp.py index 045a731321..48ab036563 100644 --- a/vllm/models/deepseek_v4/amd/mtp.py +++ b/vllm/models/deepseek_v4/amd/mtp.py @@ -240,13 +240,11 @@ class DeepSeekV4MultiTokenPredictor(nn.Module): current_step_idx, ) - def compute_logits( + def _pre_head( self, + mtp_layer: DeepSeekV4MultiTokenPredictorLayer, hidden_states: torch.Tensor, - spec_step_idx: int = 0, ) -> torch.Tensor: - current_step_idx = spec_step_idx % self.num_mtp_layers - mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)] # MTP forward returns the pre-hc_head residual (T, hc_mult * D); apply # hc_head here so logits are computed from the dense hidden state. hidden_states = hidden_states.view( @@ -260,14 +258,38 @@ class DeepSeekV4MultiTokenPredictor(nn.Module): mtp_layer.rms_norm_eps, mtp_layer.hc_eps, ) - hidden_states = _MTP_SHARED_HEAD_RMSNORM_KERNEL( + return _MTP_SHARED_HEAD_RMSNORM_KERNEL( hidden_states, mtp_layer.shared_head.norm.weight.data, mtp_layer.shared_head.norm.variance_epsilon, ) - logits = self.logits_processor(mtp_layer.shared_head.head, hidden_states) + + def compute_logits( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor: + current_step_idx = spec_step_idx % self.num_mtp_layers + mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)] + logits = self.logits_processor( + mtp_layer.shared_head.head, self._pre_head(mtp_layer, hidden_states) + ) return logits + def get_top_tokens( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor: + current_step_idx = spec_step_idx % self.num_mtp_layers + mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)] + # Vocab-parallel argmax for the greedy draft: per-rank head projection + # + local argmax + a [batch, 2*tp] (value, index) reduce, instead of + # all-gathering the full vocab logits. + return self.logits_processor.get_top_tokens( + mtp_layer.shared_head.head, self._pre_head(mtp_layer, hidden_states) + ) + class DeepSeekV4MTP(nn.Module): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): @@ -302,6 +324,15 @@ class DeepSeekV4MTP(nn.Module): ) -> torch.Tensor | None: return self.model.compute_logits(hidden_states, spec_step_idx) + def get_top_tokens( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor: + # Greedy-draft path used when use_local_argmax_reduction is enabled: + # vocab-parallel argmax, no full-vocab logits. + return self.model.get_top_tokens(hidden_states, spec_step_idx) + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # Weight name remapping for checkpoint compatibility. # Maps checkpoint weight paths to model parameter paths.