diff --git a/tests/config/test_deepseek_v4_dspark_config.py b/tests/config/test_deepseek_v4_dspark_config.py new file mode 100644 index 000000000000..21511518329f --- /dev/null +++ b/tests/config/test_deepseek_v4_dspark_config.py @@ -0,0 +1,287 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek-V4 checkpoints that ship a DSpark drafter must not be routed to MTP. + +``deepseek-ai/DeepSeek-V4-Flash-0731`` advertises ``num_nextn_predict_layers: 1`` +like every other DeepSeek-V4 config, but the ``mtp.*`` tensors behind it are a +three-stage DSpark drafter with no ``enorm``/``hnorm``/``e_proj``/``h_proj``. +Routing it to ``DeepSeekV4MTPModel`` fails deep inside the weight loader with +``KeyError: model.layers.43.mtp_block.main_norm.weight`` (vllm-project/vllm#52111). + +Which DSpark drafter a config describes was previously re-derived from +architecture strings at four sites; ``DSparkVariant`` resolves it once, so these +tests also pin the variant mapping those sites now share. +""" + +import json + +import pytest +from transformers import PretrainedConfig + +from vllm.config.model import ModelConfig +from vllm.config.parallel import ParallelConfig +from vllm.config.speculative import ( + DSparkVariant, + SpeculativeConfig, + _is_deepseek_v4_dspark, + _is_dspark_draft, +) + +# Trimmed from the published config.json. Both DeepSeek-V4-Flash variants +# declare num_hidden_layers=43 and num_nextn_predict_layers=1, so only the +# dspark_* keys below tell the DSpark drafter apart from a real MTP head. +_DEEPSEEK_V4 = { + "architectures": ["DeepseekV4ForCausalLM"], + "model_type": "deepseek_v4", + "num_hidden_layers": 43, + "num_nextn_predict_layers": 1, + "hidden_size": 512, + "intermediate_size": 1024, + "num_attention_heads": 8, + "num_key_value_heads": 8, + "vocab_size": 129280, + "max_position_embeddings": 4096, + "torch_dtype": "bfloat16", +} + +_DSPARK_KEYS = { + "dspark_block_size": 5, + "dspark_target_layer_ids": [40, 41, 42], + "dspark_markov_rank": 256, + "dspark_noise_token_id": 128799, +} + +# DeepSeek-V4-Pro-0813, the other shipped DSpark checkpoint: deeper, with +# different target layers and markov rank, and the same +# num_nextn_predict_layers=1. Detection must depend on neither the layer count +# nor the particular ids. +_PRO_DSPARK = { + "num_hidden_layers": 61, + "dspark_block_size": 5, + "dspark_target_layer_ids": [58, 59, 60], + "dspark_markov_rank": 512, + "dspark_noise_token_id": 128799, +} + + +def _hf_config(**kwargs) -> PretrainedConfig: + config = PretrainedConfig(**{**_DEEPSEEK_V4, **kwargs}) + config.model_type = kwargs.get("model_type", "deepseek_v4") + return config + + +def _checkpoint(tmp_path, name: str, **extra) -> ModelConfig: + """A real ``ModelConfig`` over a synthetic checkpoint directory. + + Only ``config.json`` is written: routing is decided from the config alone, + so no weights and no network access are needed. + """ + path = tmp_path / name + path.mkdir() + (path / "config.json").write_text(json.dumps({**_DEEPSEEK_V4, **extra})) + return ModelConfig( + model=str(path), + tokenizer_mode="skip", + skip_tokenizer_init=True, + max_model_len=4096, + ) + + +@pytest.mark.cpu_test +@pytest.mark.parametrize( + "name,keys", + [("DeepSeek-V4-Flash-0731", _DSPARK_KEYS), ("DeepSeek-V4-Pro-0813", _PRO_DSPARK)], +) +def test_explicit_mtp_on_dspark_checkpoint_is_rejected(tmp_path, name, keys): + """The reported invocation must fail at config time, not in the workers. + + Both checkpoints crash in the weight loader without this guard, at + ``model.layers.{43,61}.mtp_block.main_norm.weight`` respectively. + """ + target = _checkpoint(tmp_path, name, **keys) + + with pytest.raises(ValueError, match="ships a DSpark drafter"): + SpeculativeConfig( + method="mtp", + num_speculative_tokens=1, + target_model_config=target, + target_parallel_config=ParallelConfig(), + ) + + +@pytest.mark.cpu_test +def test_explicit_mtp_is_rejected_for_a_named_draft_model(tmp_path): + """``hf_config_override`` hides the model_type, so the draft path needs its + own check; without it the explicit request is silently rewritten to DSpark.""" + target = _checkpoint(tmp_path, "DeepSeek-V4-Flash-0731", **_DSPARK_KEYS) + + with pytest.raises(ValueError, match="ships a DSpark drafter"): + SpeculativeConfig( + method="mtp", + model=target.model, + num_speculative_tokens=1, + target_model_config=target, + target_parallel_config=ParallelConfig(), + ) + + +@pytest.mark.cpu_test +def test_rejection_names_the_token_count_dspark_needs(tmp_path): + """The suggested method has its own minimum; say so in one error, not two.""" + target = _checkpoint(tmp_path, "DeepSeek-V4-Flash-0731", **_DSPARK_KEYS) + + with pytest.raises(ValueError, match=r"dspark_block_size \(5\)"): + SpeculativeConfig( + method="mtp", + num_speculative_tokens=1, + target_model_config=target, + target_parallel_config=ParallelConfig(), + ) + + +@pytest.mark.cpu_test +def test_plain_mtp_checkpoint_still_routes_to_mtp(tmp_path): + """DeepSeek-V4-Flash has a real MTP head and must be left alone.""" + target = _checkpoint(tmp_path, "DeepSeek-V4-Flash") + + spec = SpeculativeConfig( + method="mtp", + num_speculative_tokens=1, + target_model_config=target, + target_parallel_config=ParallelConfig(), + ) + + assert spec.method == "mtp" + assert spec.draft_model_config.hf_config.architectures == ["DeepSeekV4MTPModel"] + + +@pytest.mark.cpu_test +def test_omitted_method_auto_detects_dspark(tmp_path): + """Detection must reach DSpark from the config, not from the repo name.""" + target = _checkpoint(tmp_path, "DeepSeek-V4-Flash-0731", **_DSPARK_KEYS) + + spec = SpeculativeConfig( + model=target.model, + num_speculative_tokens=5, + target_model_config=target, + target_parallel_config=ParallelConfig(), + ) + + assert spec.method == "dspark" + assert spec.draft_model_config.hf_config.architectures == ["DSparkDraftModel"] + + +@pytest.mark.cpu_test +def test_explicit_dspark_is_accepted(tmp_path): + target = _checkpoint(tmp_path, "DeepSeek-V4-Flash-0731", **_DSPARK_KEYS) + + spec = SpeculativeConfig( + method="dspark", + num_speculative_tokens=5, + target_model_config=target, + target_parallel_config=ParallelConfig(), + ) + + assert spec.method == "dspark" + assert spec.draft_model_config.hf_config.architectures == ["DSparkDraftModel"] + + +@pytest.mark.cpu_test +@pytest.mark.parametrize("keys", [_DSPARK_KEYS, _PRO_DSPARK]) +def test_dspark_drafter_detected_from_config_keys(keys): + """Both shipped DSpark checkpoints, whose layer counts and target layer ids + differ while num_nextn_predict_layers does not.""" + assert _is_deepseek_v4_dspark(_hf_config(**keys)) + assert not _is_deepseek_v4_dspark(_hf_config()) + + +@pytest.mark.cpu_test +def test_empty_target_layer_ids_is_not_a_drafter(): + """A drafter with no target layers cannot be loaded; do not claim it.""" + assert not _is_deepseek_v4_dspark(_hf_config(dspark_target_layer_ids=[])) + + +@pytest.mark.cpu_test +@pytest.mark.parametrize("model_type", ["deepseek_v3", "deepseek_v32", "qwen3_next"]) +def test_other_model_types_are_untouched(model_type): + config = _hf_config(**_DSPARK_KEYS) + config.model_type = model_type + + assert not _is_deepseek_v4_dspark(config) + + +@pytest.mark.cpu_test +def test_detected_after_hf_config_override(): + """The draft path sees the config after ``hf_config_override`` has rewritten + model_type to ``deepseek_mtp``; the dspark_* keys survive, so detection must + survive with them.""" + overridden = SpeculativeConfig.hf_config_override(_hf_config(**_DSPARK_KEYS)) + + assert overridden.model_type == "deepseek_mtp" + assert overridden.architectures == ["DeepSeekV4MTPModel"] + assert _is_deepseek_v4_dspark(overridden) + + plain = SpeculativeConfig.hf_config_override(_hf_config()) + assert plain.architectures == ["DeepSeekV4MTPModel"] + assert not _is_deepseek_v4_dspark(plain) + + +@pytest.mark.cpu_test +def test_dspark_draft_detected_without_dspark_in_the_name(): + """Auto-detection must not depend on the repo being named ``*dspark*``.""" + assert _is_dspark_draft( + "deepseek-ai/DeepSeek-V4-Flash-0731", _hf_config(**_DSPARK_KEYS) + ) + assert not _is_dspark_draft("deepseek-ai/DeepSeek-V4-Flash", _hf_config()) + + +@pytest.mark.cpu_test +def test_dspark_draft_still_detected_by_name(): + """The name remains a fallback for checkpoints that declare nothing.""" + assert _is_dspark_draft("deepseek-ai/dspark_qwen3_8b_block7", PretrainedConfig()) + + +@pytest.mark.cpu_test +@pytest.mark.parametrize( + "architecture,expected", + [ + ("Qwen3DSparkModel", DSparkVariant.QWEN3), + ("Gemma4DSparkModel", DSparkVariant.GEMMA4), + ("K3DSparkModel", DSparkVariant.K3), + ], +) +def test_variant_resolved_from_declared_architecture(architecture, expected): + config = PretrainedConfig(architectures=[architecture]) + + assert DSparkVariant.from_config(config) is expected + + +@pytest.mark.cpu_test +def test_synthesised_architecture_with_qwen3_resolves_to_qwen3(): + """A Qwen3 DSpark draft may declare the synthesised `DSparkDraftModel` + name (#52197). Without the `model_type` pairing it would fall through to + DEEPSEEK_V4 and have its `model_type` rewritten to `deepseek_v4`.""" + config = PretrainedConfig(architectures=["DSparkDraftModel"], model_type="qwen3") + + assert DSparkVariant.from_config(config) is DSparkVariant.QWEN3 + assert _is_dspark_draft("some/qwen3-draft", config) + + +@pytest.mark.cpu_test +def test_k3_draft_is_not_auto_routed_to_dspark(): + """K3 declares a DSpark architecture but upstream leaves it a plain draft + model unless the method is explicit; detection must not widen that.""" + config = PretrainedConfig(architectures=["K3DSparkModel"], model_type="k3_dspark") + + assert not _is_dspark_draft("Inferact/Kimi-K3", config) + assert DSparkVariant.from_config(config) is DSparkVariant.K3 + + +@pytest.mark.cpu_test +def test_deepseek_v4_is_the_variant_without_its_own_architecture(): + """DeepSeek-V4 DSpark reuses the target's config, so it declares no draft + architecture of its own and is what remains once the others are excluded.""" + assert DSparkVariant.from_config(_hf_config(**_DSPARK_KEYS)) is ( + DSparkVariant.DEEPSEEK_V4 + ) + assert DSparkVariant.DEEPSEEK_V4.value == "DSparkDraftModel" diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index e378906535ca..34eeedb2db9f 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -4,6 +4,7 @@ import copy import functools from collections.abc import Callable +from enum import Enum from typing import TYPE_CHECKING, Any, Literal, get_args from pydantic import Field, SkipValidation, field_validator, model_validator @@ -81,6 +82,93 @@ DraftSampleMethod = Literal["greedy", "probabilistic"] +def _architectures(hf_config: PretrainedConfig) -> list[str]: + return getattr(hf_config, "architectures", None) or [] + + +def _is_deepseek_v4_dspark(hf_config: PretrainedConfig) -> bool: + """Whether a DeepSeek-V4 config's `mtp.*` weights are a DSpark drafter. + + Every DeepSeek-V4 config declares `num_nextn_predict_layers`, but the + `mtp.*` tensors behind it are an MTP head in some checkpoints and a + multi-stage DSpark drafter in others. Only the latter carry the `dspark_*` + keys, and their `mtp.*` weights have no `enorm`/`hnorm`/`e_proj`/`h_proj`, + so `DeepSeekV4MTPModel` cannot load them. + """ + if not getattr(hf_config, "dspark_target_layer_ids", None): + return False + # ``hf_config_override`` rewrites model_type to ``deepseek_mtp`` and pins the + # MTP architecture, so accept the raw and the already-overridden shape alike. + return getattr(hf_config, "model_type", None) == "deepseek_v4" or ( + "DeepSeekV4MTPModel" in _architectures(hf_config) + ) + + +class DSparkVariant(Enum): + """Which DSpark drafter a draft config describes. + + The value is the draft architecture vLLM loads. Qwen3, Gemma4 and K3 declare + theirs in the checkpoint; DeepSeek-V4 DSpark reuses the target's own + DeepSeek-V4 config and has no draft architecture of its own, so vLLM + synthesises `DSparkDraftModel` for it in `_verify_args`. + """ + + QWEN3 = "Qwen3DSparkModel" + GEMMA4 = "Gemma4DSparkModel" + K3 = "K3DSparkModel" + DEEPSEEK_V4 = "DSparkDraftModel" + + @classmethod + def declared(cls, hf_config: PretrainedConfig) -> "DSparkVariant | None": + """The variant a config names, or None if it names no DSpark drafter. + + A Qwen3 DSpark draft may ship the synthesised `DSparkDraftModel` name + rather than its own, so that name only counts paired with `qwen3`. + """ + declared = _architectures(hf_config) + for variant in cls: + if variant is not cls.DEEPSEEK_V4 and variant.value in declared: + return variant + if ( + cls.DEEPSEEK_V4.value in declared + and getattr(hf_config, "model_type", None) == "qwen3" + ): + return cls.QWEN3 + return None + + @classmethod + def detected(cls, hf_config: PretrainedConfig) -> "DSparkVariant | None": + """The variant auto-detection claims, or None if it claims nothing. + + K3 declares a DSpark architecture but is deliberately not auto-routed: + without an explicit method a K3 draft stays a plain draft model. + """ + variant = cls.declared(hf_config) + return None if variant is cls.K3 else variant + + @classmethod + def from_config(cls, hf_config: PretrainedConfig) -> "DSparkVariant": + """Resolve the variant of a config already known to be DSpark. + + Must be called before the branches below rewrite `architectures`. + """ + return cls.declared(hf_config) or cls.DEEPSEEK_V4 + + +def _is_dspark_draft(model: str, hf_config: PretrainedConfig) -> bool: + """Whether a draft checkpoint ships a DSpark drafter. + + The repo name is only a fallback: DeepSeek-V4 DSpark checkpoints are not + reliably named `*dspark*` (e.g. DeepSeek-V4-Flash-0731, which is otherwise + routed to MTP and dies in the weight loader), so ask the config first. + """ + if DSparkVariant.detected(hf_config) is not None: + return True + if _is_deepseek_v4_dspark(hf_config): + return True + return "dspark" in model.lower() + + @config class SpeculativeConfig: """Configuration for speculative decoding.""" @@ -738,6 +826,21 @@ def _is_custom_proposer_path(model: str | None) -> bool: parts = model.split(".") return len(parts) >= 2 and all(part.isidentifier() for part in parts) + def _reject_mtp_on_dspark(self, hf_config: PretrainedConfig, model: str) -> None: + """Refuse an explicit ``method="mtp"`` on a DeepSeek-V4 DSpark checkpoint. + + Reached from both entry points: the drafter taken from the target + checkpoint, and a draft model named on the command line. + """ + if self.method == "mtp" and _is_deepseek_v4_dspark(hf_config): + block_size = getattr(hf_config, "dspark_block_size", None) + raise ValueError( + f"{model} ships a DSpark drafter rather than an MTP head, so " + "method='mtp' cannot load its weights. Use method='dspark' " + f"with num_speculative_tokens >= dspark_block_size " + f"({block_size})." + ) + def __post_init__(self): # Note: "method" is a new parameter that helps to extend the # configuration of non-model-based proposers, and the "model" parameter @@ -768,6 +871,10 @@ def __post_init__(self): if self.method == "mtp": if self.target_model_config is None: raise ValueError("target_model_config must be present for mtp") + self._reject_mtp_on_dspark( + self.target_model_config.hf_text_config, + self.target_model_config.model, + ) # use the draft model from the same model: self.model = self.target_model_config.model # Align the quantization of draft model for cases such as @@ -936,6 +1043,10 @@ def __post_init__(self): draft_hf.vocab_size = target_vocab draft_hf.truncated_vocab_size = target_vocab + self._reject_mtp_on_dspark( + self.draft_model_config.hf_config, self.draft_model_config.model + ) + # Automatically detect the method if self.method in ("eagle", "eagle3", "dflash", "dspark"): pass @@ -954,14 +1065,8 @@ def __post_init__(self): in self.draft_model_config.architectures ): self.method = "dflash" - elif ( - "dspark" in self.draft_model_config.model.lower() - or "Qwen3DSparkModel" in self.draft_model_config.architectures - or "Gemma4DSparkModel" in self.draft_model_config.architectures - or ( - "DSparkDraftModel" in self.draft_model_config.architectures - and self.draft_model_config.hf_config.model_type == "qwen3" - ) + elif _is_dspark_draft( + self.draft_model_config.model, self.draft_model_config.hf_config ): self.method = "dspark" elif self.draft_model_config.hf_config.model_type == "medusa": @@ -1019,34 +1124,35 @@ def __post_init__(self): self.draft_model_config.hf_config = eagle_config self.update_arch_() + # Resolved before the branches below rewrite ``architectures``, + # so every later branch reads one answer. + dspark_variant = ( + DSparkVariant.from_config(self.draft_model_config.hf_config) + if self.method == "dspark" + else None + ) + if ( - self.method == "dspark" - and "DSparkDraftModel" in self.draft_model_config.architectures - and self.draft_model_config.hf_config.model_type == "qwen3" + dspark_variant is DSparkVariant.QWEN3 + and DSparkVariant.DEEPSEEK_V4.value + in self.draft_model_config.architectures ): self.draft_model_config.hf_config.architectures = [ - "Qwen3DSparkModel" + DSparkVariant.QWEN3.value ] self.update_arch_() - elif self.method == "dspark" and ( - "Qwen3DSparkModel" not in self.draft_model_config.architectures - and "Gemma4DSparkModel" not in self.draft_model_config.architectures - and "K3DSparkModel" not in self.draft_model_config.architectures - ): + elif dspark_variant is DSparkVariant.DEEPSEEK_V4: # DeepSeek-V4 DSpark reuses the full DeepSeek-V4 config # and its weights ship in the target checkpoint. self.draft_model_config.hf_config.model_type = "deepseek_v4" self.draft_model_config.hf_config.architectures = [ - "DSparkDraftModel" + dspark_variant.value ] self.draft_model_config.quantization = ( self.target_model_config.quantization ) self.update_arch_() - elif ( - self.method == "dspark" - and "Gemma4DSparkModel" in self.draft_model_config.architectures - ): + elif dspark_variant is DSparkVariant.GEMMA4: # Normalize the self-contained Gemma4 draft's config keys to # the DSpark conventions. hf = self.draft_model_config.hf_config @@ -1115,10 +1221,7 @@ def __post_init__(self): "dspark_draft_topk must be between 1 and the " f"draft vocabulary size ({draft_vocab_size})" ) - if ( - "Qwen3DSparkModel" - not in self.draft_model_config.architectures - ): + if dspark_variant is not DSparkVariant.QWEN3: raise ValueError( "dspark_draft_topk is only supported by " "Qwen3DSparkModel"