NVFP4 checkpoints: DSpark draft experts (MXFP4) routed through NVFP4 path — 0% acceptance; prefix-aware routing fix + measurements

[auroter]

Defect

NVFP4 conversions of DeepSeek-V4-Flash (NVIDIA's modelopt cast_mxfp4_to_nvfp4 flow) quantize only the main 43-layer stack; the DSpark draft module keeps its original MXFP4 experts. DeepseekV4FP8Config.get_quant_method (vllm/models/deepseek_v4/quant_config.py) routes every RoutedExperts layer through ModelOptNvFp4FusedMoE whenever the checkpoint declares moe_quant_algo=NVFP4 — the prefix argument is unused — so the draft module's MXFP4 weights are decoded as NVFP4 noise. The engine boots and serves normally, but draft acceptance is 0% and speculative decode collapses below the non-speculative rate (42.3 tok/s vs 124.5 without speculation on our hardware). Teacher-forced perplexity is unaffected, which makes the failure easy to miss.

Fix (verified)

Make the routing prefix-aware: draft-module layers (runtime layer index >= num_hidden_layers, per DeepSeekV4DSparkLayer's runtime_layer_idx = config.num_hidden_layers + dspark_layer_idx) load through Mxfp4MoEMethod; is_mxfp4_quant gives the mirrored answer. Python-only, no recompile.

--- a/vllm/models/deepseek_v4/quant_config.py
+++ b/vllm/models/deepseek_v4/quant_config.py
@@ -5,6 +5,7 @@
 from __future__ import annotations
 
+import re
 from typing import TYPE_CHECKING, cast
 
@@ get_quant_method @@
             if self.expert_dtype == "fp4":
-                if self.moe_quant_algo == "NVFP4":
+                if self.moe_quant_algo == "NVFP4" and not self._is_draft_layer(prefix):
                     from vllm.model_executor.layers.quantization.modelopt import (
                         ModelOptNvFp4FusedMoE,
                     )
 
                     return ModelOptNvFp4FusedMoE(
                         quant_config=self._get_nvfp4_config(),
                         moe_config=layer.moe_config,
                     )
                 return Mxfp4MoEMethod(layer.moe_config)
@@ class DeepseekV4FP8Config @@
+    def _is_draft_layer(self, prefix: str) -> bool:
+        # NVFP4 conversions quantize only the main stack; the speculative
+        # module (runtime layer index >= num_hidden_layers) keeps its
+        # original MXFP4 experts and must load through the MXFP4 path.
+        m = re.search(r"layers\.(\d+)\.", prefix)
+        if m is None:
+            return False
+        try:
+            hf_config = get_current_vllm_config().model_config.hf_config
+        except Exception:
+            return False
+        return int(m.group(1)) >= hf_config.num_hidden_layers
+
     def is_mxfp4_quant(self, prefix, layer):
         if not isinstance(layer, RoutedExperts) or self.expert_dtype != "fp4":
             return False
-        return self.moe_quant_algo != "NVFP4"
+        return self.moe_quant_algo != "NVFP4" or self._is_draft_layer(prefix)

Measurements

4x RTX PRO 6000 Blackwell Workstation (SM120), TP4, tag sm120-pr-41834-stable-preview-20260727d, DeepSeek-V4-Flash-0731 NVFP4, --speculative-config '{"method":"dspark","num_speculative_tokens":5}':

	draft acceptance	single-stream decode
before	0.0%	42.3 tok/s
after	34–39%	168.7 tok/s
after + VLLM_USE_V2_MODEL_RUNNER=1	41.4%	208.4 tok/s
native MXFP4 checkpoint, same flags (reference)	—	210.5 tok/s

Perplexity pinned at each backend's baseline in all cells (speculation lossless). Artifact, full gate results, and this patch as a file: https://huggingface.co/auroter/DeepSeek-V4-Flash-0731-NVFP4

Happy to turn this into a PR — the patch applies clean at d64074e.

🤖 Generated with Claude Code

[jasl]

Verified still-valid on current HEAD: get_quant_method in vllm/models/deepseek_v4/quant_config.py remains prefix-blind — the NVFP4/MXFP4 decision reads only the memoized expert_dtype/moe_quant_algo from the target config, and the draft reuses the target's quant-config instance (dspark/utils.py deliberately doesn't replace it), so draft RoutedExperts take the NVFP4 branch exactly as you describe.

Two notes for the fix, from tracing the actual prefixes:

• Match on the extracted layer index, not string literals: the routing prefixes are layers.43/44/45.ffn.experts (NVIDIA DSpark, root prefix empty) but model.layers.43.... for MTP — there is no mtp. segment in the routing prefix (that indexing only exists in checkpoint parameter names). extract_layer_index(prefix) >= num_hidden_layers is the robust predicate, and there's an in-tree precedent: _is_dspark_runtime_layer in nvidia/model.py, which also gates on speculative_config.method == "dspark" so MTP layers aren't caught by accident.
• The prefix routing needs to be checked independently of the is_layer_skipped branch: a conversion that lists the draft module in ignored_layers would currently return UnquantizedFusedMoEMethod — also wrong for MXFP4 weights, just a different failure.

Your measurements (0% acceptance, ppl-invisible) match the mechanism. Happy to take this as a PR given you already have it verified on hardware — otherwise I'll pick it up.

[jasl]

Fixed in 9aaf122, following your design: routing keyed on the extracted layer index (layers.N with N >= num_hidden_layers), which covers both prefix roots in-tree — string-literal matching was a trap since no mtp. segment exists in routing prefixes. is_mxfp4_quant gives the mirrored answer, and the check runs independently of the ignored_layers branch. Unit tests cover main-stack, all draft-range prefixes, and the unresolvable cases.

One honest caveat on end-to-end validation: our NVFP4+spec serve attempt is currently blocked by an unrelated dependency deadlock — the flashinfer-jit-cache 0.6.16 binaries need TVMFFIGetCustomAllocator (apache-tvm-ffi >= 0.1.13) while no tilelang release tolerates tvm-ffi 0.1.13 at runtime, so the wheel's fused_moe_120 cannot load in any environment that also runs our TileLang mHC path. FlashInfer's next release upgrades its tvm-ffi pairing and should clear this; we'll validate acceptance end-to-end then. Your hardware-verified measurements plus the unit coverage are what this close rides on — reopen if anything disagrees on your cluster once the FI release lands. Thanks for an exemplary report.