[Misc][Spec Decode] Detect DeepSeek-V4 DSpark checkpoints from config

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"