Fix vanilla DSpark checkpoint exports for vLLM

CHANGELOG.rst

@@ -32,6 +32,7 @@ Changelog
 
 **Bug Fixes**
 
+- Fix vanilla DSpark exports to use the Qwen3 DSpark architecture, head names and configuration expected by vLLM. Re-export existing training checkpoints to use this format.
 - Fix ONNX quantization placing a per-tensor Q/DQ pair on a rank-0 (scalar) constant input, which TensorRT's ONNX parser rejects (``convertAxis`` asserts the quantization axis lies in ``[0, nbDims]`` but ``nbDims == 0``). Rank-0 constant operands are now kept unquantized in the int8, fp8, and autotune export paths; the existing single-element Mul/Div/Pad stripping stays gated under ``target_dla``.
 - Fix PTQ of native Transformers FP8 experts to apply source weight scales during calibration and export. Re-quantize affected checkpoints from the original FP8 source.
 - Fix ONNX preprocessing with custom TensorRT plugin libraries on TensorRT 11.

modelopt/torch/export/plugins/hf_spec_export.py

@@ -582,14 +582,23 @@ def _export_config(self):
 class DSparkExporter(DFlashExporter):
     """Draft model exporter for DSpark (DFlash backbone + sequential Markov head).
 
-    Same z-lab-compatible format as DFlash, plus the DSpark head weights
-    (``markov_w1.*`` / ``markov_w2.*`` / ``gate_proj.*`` / ``joint_proj.*`` /
-    ``confidence_proj.*``, already captured by the inherited ``dflash_module.``
-    stripping) and the extra config fields the loader needs to rebuild the head
-    (``projector_type``, ``markov_rank``, ``markov_head_type``,
-    ``use_confidence_head``, ``shift_label``).
+    Vanilla heads use the vLLM Qwen3 DSpark layout. Gated and recurrent heads keep
+    the native layout because that runtime does not implement those variants.
     """
 
+    def _extract_state_dict(self, full_state_dict: dict):
+        export_sd = super()._extract_state_dict(full_state_dict)
+        if self.model.dflash_module.markov_head_type == "vanilla":
+            prefixes = {
+                "markov_w1.": "markov_head.markov_w1.",
+                "markov_w2.": "markov_head.markov_w2.",
+                "confidence_proj.": "confidence_head.proj.",
+            }
+            for old_prefix, new_prefix in prefixes.items():
+                for key in [key for key in export_sd if key.startswith(old_prefix)]:
+                    export_sd[new_prefix + key[len(old_prefix) :]] = export_sd.pop(key)
+        return export_sd
+
     def _export_config(self):
         """Extend the DFlash config with the DSpark head fields."""
         config = super()._export_config()
@@ -600,10 +609,22 @@ def _export_config(self):
                 "projector_type": getattr(draft_config, "projector_type", "dspark"),
                 "shift_label": getattr(draft_config, "shift_label", True),
                 "markov_rank": draft_config.markov_rank,
-                "markov_head_type": getattr(draft_config, "markov_head_type", "vanilla"),
+                "markov_head_type": self.model.dflash_module.markov_head_type,
                 "use_confidence_head": bool(getattr(draft_config, "use_confidence_head", False)),
             }
         )
+        if config["dflash_config"]["markov_head_type"] == "vanilla":
+            config["architectures"] = ["Qwen3DSparkModel"]
+            config.update(
+                {
+                    "mask_token_id": config["dflash_config"]["mask_token_id"],
+                    "target_layer_ids": config["dflash_config"]["target_layer_ids"],
+                    "markov_rank": draft_config.markov_rank,
+                    "sample_from_anchor": True,
+                    "enable_confidence_head": config["dflash_config"]["use_confidence_head"],
+                    "confidence_head_with_markov": config["dflash_config"]["use_confidence_head"],
+                }
+            )
         return config
 
 

modelopt/torch/speculative/plugins/modeling_draft_checkpoint_compat.py

@@ -13,13 +13,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Per-released-checkpoint key remaps for loading published drafters.
+"""Checkpoint key remaps for loading exported and published drafters.
 
-ModelOpt's draft modules define one canonical parameter layout, which is also the export
-format. Some published drafters ship the same tensors under different key names; loading
-those requires rewriting keys, which is a property of a specific release rather than of the
-architecture. Keeping those remaps here — instead of in the module that defines the
-architecture — means a new checkpoint quirk adds an entry to this file only.
+ModelOpt's draft modules define one canonical parameter layout. Serving exports and
+published drafters can use different key names; the load hooks here translate those
+names back to the module layout.
 
 Each entry is a ``load_state_dict`` pre-hook body: it mutates ``state_dict`` in place to
 move released-layout keys onto the canonical names.
@@ -29,17 +27,20 @@
 
 logger = logging.getLogger(__name__)
 
-# Head submodules that some published DSpark checkpoints nest under a `markov_head.` parent.
-# ModelOpt keeps them flat (``markov_w1`` / ``markov_w2`` / ...), matching the upstream
-# DeepSpec layout; ``nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16-DSpark`` nests them.
+# DSpark modules keep Markov heads flat; serving checkpoints can nest them.
 _DSPARK_NESTED_HEAD_PREFIX = "markov_head."
 
 
 def remap_dspark_nested_head_keys(state_dict, prefix, *args, **kwargs):
-    """Accept DSpark head weights nested under ``markov_head.`` as well as flat."""
+    """Accept DSpark Markov and confidence heads in serving or native layouts."""
     nested = prefix + _DSPARK_NESTED_HEAD_PREFIX
-    for key in [k for k in state_dict if k.startswith(nested)]:
-        flat = prefix + key[len(nested) :]
+    confidence = prefix + "confidence_head.proj."
+    for key in [k for k in state_dict if k.startswith((nested, confidence))]:
+        flat = (
+            prefix + key[len(nested) :]
+            if key.startswith(nested)
+            else prefix + "confidence_proj." + key[len(confidence) :]
+        )
         nested_value = state_dict.pop(key)
         # A flat key already present wins. The Markov tables are ~14% of the draft's
         # parameters and loading the wrong copy is invisible at runtime, so say which was

modelopt/torch/speculative/plugins/modeling_dspark.py

@@ -62,8 +62,8 @@
 The head uses its OWN embedding table (``markov_w1``), not the base model's, so
 the bias computation is fully self-contained on this module. Submodule names
 (``markov_w1`` / ``markov_w2`` / ``gate_proj`` / ``joint_proj`` /
-``confidence_proj``) match the upstream DeepSpec checkpoint layout so exported
-checkpoints stay portable.
+``confidence_proj``) match the upstream DeepSpec layout; ``DSparkExporter`` adapts
+the serving checkpoint format.
 """
 
 import logging

tests/unit/torch/speculative/plugins/test_hf_dspark.py

@@ -18,9 +18,7 @@
 DSpark reuses the DFlash mode/pipeline and adds a lightweight sequential (Markov)
 head plus an optional confidence head. These tests cover conversion routing for
 the three head variants, the three-term training forward (CE + TVD + confidence
-BCE), and the export format (head weights + config) against the z-lab-compatible
-layout (``markov_w1.*`` / ``markov_w2.*`` / ``gate_proj.*`` / ``joint_proj.*`` /
-``confidence_proj.*``).
+BCE), vLLM vanilla exports, and native gated/recurrent exports.
 """
 
 import json
@@ -267,7 +265,7 @@ def spy(*args, **kwargs):
 
 
 class TestDSparkExporter:
-    """Test the DSpark checkpoint export format (z-lab-compatible layout)."""
+    """Test vanilla serving exports and the native gated/recurrent formats."""
 
     def _export(self, tmp_path, head_type="vanilla", use_confidence_head=False):
         model = get_tiny_llama(num_hidden_layers=4)
@@ -286,38 +284,57 @@ def _export(self, tmp_path, head_type="vanilla", use_confidence_head=False):
         model.get_exporter().export(export_dir)
         return export_dir
 
-    @pytest.mark.parametrize("head_type", HEAD_TYPES)
-    def test_export_weight_keys_match_reference(self, tmp_path, head_type):
-        """Exported weights carry the head tensors under reference names, no prefix."""
-        sd = load_file(str(self._export(tmp_path, head_type=head_type) / "model.safetensors"))
+    @pytest.mark.parametrize("head_type", [*HEAD_TYPES, "VANILLA"])
+    @pytest.mark.parametrize("use_confidence_head", [False, True])
+    def test_export_weight_keys_match_reference(self, tmp_path, head_type, use_confidence_head):
+        """Only vanilla heads use the Qwen3 DSpark serving tensor names."""
+        export_dir = self._export(
+            tmp_path, head_type=head_type, use_confidence_head=use_confidence_head
+        )
+        sd = load_file(str(export_dir / "model.safetensors"))
         for key in sd:
             assert "dflash_module." not in key
             assert "rotary_emb" not in key
-        assert "markov_w1.weight" in sd
-        assert "markov_w2.weight" in sd
+        nested = head_type.lower() == "vanilla"
+        for name in ("markov_w1.weight", "markov_w2.weight"):
+            assert (name in sd) == (not nested)
+            assert (f"markov_head.{name}" in sd) == nested
         assert ("gate_proj.weight" in sd) == (head_type == "gated")
         assert ("joint_proj.weight" in sd) == (head_type == "rnn")
-
-    def test_export_includes_confidence_weights(self, tmp_path):
-        """The confidence head weights are exported when enabled."""
-        sd = load_file(str(self._export(tmp_path, use_confidence_head=True) / "model.safetensors"))
-        assert "confidence_proj.weight" in sd
-
-    def test_export_config_has_dspark_fields(self, tmp_path):
-        """config.json carries the dflash_config DSpark head fields."""
-        export_dir = self._export(tmp_path, head_type="gated")
+        for suffix in ("weight", "bias"):
+            assert (f"confidence_head.proj.{suffix}" in sd) == (nested and use_confidence_head)
+            assert (f"confidence_proj.{suffix}" in sd) == (not nested and use_confidence_head)
+
+    @pytest.mark.parametrize("head_type", [*HEAD_TYPES, "VANILLA"])
+    @pytest.mark.parametrize("use_confidence_head", [False, True])
+    def test_export_config_has_dspark_fields(self, tmp_path, head_type, use_confidence_head):
+        """Vanilla exports select the runtime and preserve the trained head settings."""
+        export_dir = self._export(
+            tmp_path, head_type=head_type, use_confidence_head=use_confidence_head
+        )
         with open(export_dir / "config.json") as f:
             cfg = json.load(f)
 
-        assert cfg["architectures"] == ["DFlashDraftModel"]
+        vanilla = head_type.lower() == "vanilla"
+        assert cfg["architectures"] == ["Qwen3DSparkModel" if vanilla else "DFlashDraftModel"]
         dc = cfg["dflash_config"]
         assert dc["projector_type"] == "dspark"
         assert dc["markov_rank"] == MARKOV_RANK
-        assert dc["markov_head_type"] == "gated"
-        assert dc["use_confidence_head"] is False
+        assert dc["markov_head_type"] == head_type.lower()
+        assert dc["use_confidence_head"] is use_confidence_head
         assert dc["shift_label"] is True
         assert "mask_token_id" in dc
         assert "target_layer_ids" in dc
+        if vanilla:
+            assert cfg["model_type"] == "qwen3"
+            assert cfg["markov_rank"] == MARKOV_RANK
+            assert cfg["sample_from_anchor"] is True
+            assert cfg["enable_confidence_head"] is use_confidence_head
+            assert cfg["confidence_head_with_markov"] is use_confidence_head
+            assert cfg["mask_token_id"] == dc["mask_token_id"]
+            assert cfg["target_layer_ids"] == dc["target_layer_ids"]
+        else:
+            assert "enable_confidence_head" not in cfg
 
 
 class TestDraftAttentionPattern:
@@ -533,15 +550,17 @@ def test_export_omits_sink_when_disabled(self, tmp_path):
 class TestMarkovHeadKeyRemap:
     """Head weights load from either the flat or the nested `markov_head.` layout.
 
-    ModelOpt (and the upstream DeepSpec reference) keeps the head tensors flat on the
-    module and exports them that way. Some released drafters — notably
+    ModelOpt (and the upstream DeepSpec reference) keeps head tensors flat on the
+    module. Serving exports and some released drafters — notably
     nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16-DSpark — nest them under a
     `markov_head.` parent, so loading accepts both.
     """
 
-    def _module(self):
+    def _module(self, use_confidence_head=False):
         model = get_tiny_llama(num_hidden_layers=4)
-        mtsp.convert(model, [("dflash", _get_dspark_config())])
+        mtsp.convert(
+            model, [("dflash", _get_dspark_config(use_confidence_head=use_confidence_head))]
+        )
         return model.dflash_module
 
     def test_nested_keys_load(self):
@@ -560,8 +579,8 @@ def test_nested_keys_load(self):
         )
 
     def test_flat_keys_still_load(self):
-        """The exported (flat) layout keeps working unchanged."""
-        module = self._module()
+        """The legacy flat layout keeps working unchanged."""
+        module = self._module(use_confidence_head=True)
         flat = dict(module.state_dict())
         flat["markov_w1.weight"] = torch.full_like(flat["markov_w1.weight"], 0.25)
         res = module.load_state_dict(flat, strict=True)
@@ -570,16 +589,21 @@ def test_flat_keys_still_load(self):
             module.markov_w1.weight, torch.full_like(module.markov_w1.weight, 0.25)
         )
 
-    def test_flat_key_wins_over_nested(self):
+    @pytest.mark.parametrize(
+        ("flat_key", "nested_key"),
+        [
+            ("markov_w1.weight", "markov_head.markov_w1.weight"),
+            ("confidence_proj.weight", "confidence_head.proj.weight"),
+        ],
+    )
+    def test_flat_key_wins_over_nested(self, flat_key, nested_key):
         """An explicit flat key is never clobbered by a nested duplicate."""
-        module = self._module()
+        module = self._module(use_confidence_head=True)
         sd = dict(module.state_dict())
-        sd["markov_w1.weight"] = torch.full_like(sd["markov_w1.weight"], 1.0)
-        sd["markov_head.markov_w1.weight"] = torch.full_like(sd["markov_w1.weight"], 9.0)
+        sd[flat_key] = torch.full_like(sd[flat_key], 1.0)
+        sd[nested_key] = torch.full_like(sd[flat_key], 9.0)
         module.load_state_dict(sd, strict=True)
-        assert torch.allclose(
-            module.markov_w1.weight, torch.full_like(module.markov_w1.weight, 1.0)
-        )
+        assert torch.allclose(module.state_dict()[flat_key], torch.full_like(sd[flat_key], 1.0))
 
 
 class TestInitCheckpoint:
@@ -605,10 +629,41 @@ def _export(self, tmp_path, **overrides):
         model.get_exporter().export(export_dir)
         return export_dir
 
-    def test_warm_start_loads_exported_weights(self, tmp_path):
-        """Every draft parameter comes from the checkpoint, not a fresh init."""
-        export_dir = self._export(tmp_path)
-        model = self._make_model(init_checkpoint=export_dir)
+    @pytest.mark.parametrize("use_confidence_head", [False, True])
+    @pytest.mark.parametrize("legacy_layout", [False, True])
+    def test_warm_start_loads_exported_weights(self, tmp_path, use_confidence_head, legacy_layout):
+        """New serving exports and legacy exports warm-start every draft parameter."""
+        architecture = {
+            **_get_dspark_config()["dflash_architecture_config"],
+            "use_confidence_head": use_confidence_head,
+        }
+        export_dir = self._export(tmp_path, dflash_architecture_config=architecture)
+        if legacy_layout:
+            path = export_dir / "model.safetensors"
+            sd = load_file(str(path))
+            sd = {
+                key.removeprefix("markov_head.").replace(
+                    "confidence_head.proj.", "confidence_proj."
+                ): value
+                for key, value in sd.items()
+            }
+            save_file(sd, str(path))
+            config_path = export_dir / "config.json"
+            config = json.loads(config_path.read_text())
+            config["architectures"] = ["DFlashDraftModel"]
+            for field in (
+                "mask_token_id",
+                "target_layer_ids",
+                "markov_rank",
+                "sample_from_anchor",
+                "enable_confidence_head",
+                "confidence_head_with_markov",
+            ):
+                config.pop(field)
+            config_path.write_text(json.dumps(config))
+        model = self._make_model(
+            init_checkpoint=export_dir, dflash_architecture_config=architecture
+        )
         for name, p in model.dflash_module.named_parameters():
             assert torch.allclose(p, torch.full_like(p, 0.125)), f"{name} was not warm-started"
 
@@ -695,20 +750,26 @@ def test_restore_does_not_replay_warm_start(self, tmp_path):
             # than returning False. What matters here is the values, not the load dtype.
             assert torch.allclose(p, ref.to(p.dtype)), f"{name} differs after restore"
 
-    def test_nested_head_shape_mismatch_reported(self, tmp_path):
-        """A wrong-shaped tensor is caught even under the nested `markov_head.` layout.
+    @pytest.mark.parametrize(
+        "head_key", ["markov_head.markov_w1.weight", "confidence_head.proj.weight"]
+    )
+    def test_nested_head_shape_mismatch_reported(self, tmp_path, head_key):
+        """A wrong-shaped tensor is caught even under the nested serving layout.
 
         The shape check has to resolve the module's load hooks first; otherwise a remapped
         key skips it and fails later with a far less obvious error.
         """
-        export_dir = self._export(tmp_path)
+        architecture = {
+            **_get_dspark_config()["dflash_architecture_config"],
+            "use_confidence_head": True,
+        }
+        export_dir = self._export(tmp_path, dflash_architecture_config=architecture)
         path = export_dir / "model.safetensors"
         sd = load_file(str(path))
-        sd["markov_head.markov_w1.weight"] = torch.zeros(3, MARKOV_RANK)
-        del sd["markov_w1.weight"]
+        sd[head_key] = torch.zeros(3, sd[head_key].shape[1])
         save_file(sd, str(path))
         with pytest.raises(ValueError, match="shape mismatch"):
-            self._make_model(init_checkpoint=export_dir)
+            self._make_model(init_checkpoint=export_dir, dflash_architecture_config=architecture)
 
 
 class TestExplicitTargetLayerIds: