Proposed fix for #20: route DS4 MTP drafter experts to Mxfp4MoEMethod; guard against uninitialized drafter scales

[macdad222]

Proposed fix for #20

Ready-to-apply diffs below, validated on 2x RTX PRO 6000 Blackwell across four configs (stock FP4 TP2, W2 TP2, W2 TP1 x2 variants). Filing as an issue rather than a PR because patch/vllm-moet-v0.24.0.patch is generated from the private fork branch (patch/SOURCE.txt: "Never edit by hand") — these diffs are meant to be merged into the moet-v0.24.0 branch and regenerated via tools/check_patch_files.py --update.

Note on ownership of the two halves:

• Fix 1 (quant_config.py) patches a file that is stock vLLM 0.24.0 (not in patch/FILES.txt), so this half is arguably an upstream vLLM bug — but your fork branch is the practical place to carry it for now.
• Fix 2 (mtp.py) patches a file your patch set already owns (vllm/models/deepseek_v4/nvidia/mtp.py).

Fix 1 — route MTP drafter experts to the method their weights are stored in

The drafter's experts are MXFP4 in the checkpoint; ignored_layers: ["mtp.*"] never matches model.layers.43.ffn.experts, so they fall into the NVFP4 branch. Detect drafter layers by index >= num_hidden_layers:

--- a/vllm/models/deepseek_v4/quant_config.py
+++ b/vllm/models/deepseek_v4/quant_config.py
@@ -140,7 +140,26 @@
             ):
                 return UnquantizedFusedMoEMethod(layer.moe_config)
             if self.expert_dtype == "fp4":
-                if self.moe_quant_algo == "NVFP4":
+                # W2FIX: the checkpoint quantization_config ignores "mtp.*" -
+                # the MTP drafter experts stay MXFP4 (ue8m0 block-32 scales,
+                # no per-tensor scales) even when the main stack was
+                # re-exported as NVFP4 (moe_quant_algo=NVFP4). Route the
+                # drafter to the MXFP4 method; NVFP4 would load e8m0 bytes
+                # into e4m3 scale params and read uninitialized
+                # weight_scale_2/input_scale.
+                _is_mtp = False
+                try:
+                    import re as _re
+                    from vllm.config import get_current_vllm_config
+                    _m = _re.search(r"\.layers\.(\d+)\.", prefix or "")
+                    if _m is not None:
+                        _nh = (get_current_vllm_config()
+                               .model_config.hf_config.get_text_config()
+                               .num_hidden_layers)
+                        _is_mtp = int(_m.group(1)) >= int(_nh)
+                except Exception:
+                    _is_mtp = "mtp" in (prefix or "")
+                if self.moe_quant_algo == "NVFP4" and not _is_mtp:
                     from vllm.model_executor.layers.quantization.modelopt import (
                         ModelOptNvFp4FusedMoE,
                     )
@@ -149,6 +168,11 @@
                         quant_config=self._get_nvfp4_config(),
                         moe_config=layer.moe_config,
                     )
+                from vllm.logger import init_logger as _il
+                if _is_mtp:
+                    _il(__name__).info(
+                        "W2FIX: MTP drafter experts (%s) -> Mxfp4MoEMethod",
+                        prefix)
                 return Mxfp4MoEMethod(layer.moe_config)
             # expert_dtype == "fp8": fall through to Fp8Config which
             # returns Fp8MoEMethod with block-wise float32 scales.

Fix 2 — never let torch.empty residue become live scales (defense in depth)

Even with Fix 1, any future checkpoint that legitimately routes MTP experts through ModelOptNvFp4FusedMoE but lacks per-tensor scales would regress silently. Harvest main-layer means during load_weights and fill the drafter's scale params so uninitialized memory can never be read as scales:

--- a/vllm/models/deepseek_v4/nvidia/mtp.py
+++ b/vllm/models/deepseek_v4/nvidia/mtp.py
@@ -66,6 +66,10 @@
 _EXPERT_SCALE_RE = re.compile(r"\.experts\.\d+\.w[123]\.scale$")
 
 
+from vllm.logger import init_logger as _w2fix_init_logger
+_w2fix_logger = _w2fix_init_logger("vllm.w2fix.mtp")
+
+
 class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
     def __init__(
         self,
@@ -331,6 +335,8 @@
         ]
         params_dict = dict(self.named_parameters())
         loaded_params: set[str] = set()
+        _w2fix_in: list[float] = []
+        _w2fix_s2: list[float] = []
 
         # TP for attention
         tp_size = get_tensor_model_parallel_world_size()
@@ -365,6 +371,14 @@
         )
 
         for name, loaded_weight in weights:
+            if ".ffn.experts." in name and name.startswith("layers.42."):
+                try:
+                    if name.endswith(".input_scale"):
+                        _w2fix_in.append(float(loaded_weight.float().mean()))
+                    elif name.endswith(".weight_scale_2"):
+                        _w2fix_s2.append(float(loaded_weight.float().mean()))
+                except Exception:  # noqa: BLE001
+                    pass
             mtp_layer_idx = _find_mtp_layer_idx(name)
             # V4 checkpoints store MTP weights as `mtp.{i}.*`; remap to
             # `model.layers.{num_hidden_layers + i}.*` so that
@@ -481,6 +495,29 @@
                     f"Use a checkpoint that includes MTP layer weights, "
                     f"or disable speculative decoding."
                 )
+        # W2FIX: this checkpoint ships MTP experts with only weight+block
+        # scale; ModelOpt NVFP4 also expects per-tensor weight_scale_2 /
+        # input_scale, which otherwise stay uninitialized (torch.empty
+        # garbage -> inf alphas -> NaN drafter, 0% MTP acceptance).
+        try:
+            _s2 = (sum(_w2fix_s2) / len(_w2fix_s2)) if _w2fix_s2 else 2.0 ** -13
+            _isc = (sum(_w2fix_in) / len(_w2fix_in)) if _w2fix_in else 0.0026
+            for _lyr in self.model.layers.values():
+                _exp = _lyr.mtp_block.ffn.experts
+                _re = getattr(_exp, "routed_experts", _exp)
+                for _pn, _val in (("w13_weight_scale_2", _s2),
+                                  ("w2_weight_scale_2", _s2),
+                                  ("w13_input_scale", _isc),
+                                  ("w2_input_scale", _isc)):
+                    _prm = getattr(_re, _pn, None)
+                    if _prm is not None:
+                        _prm.data.fill_(_val)
+            _w2fix_logger.info(
+                "W2FIX: filled drafter expert scales scale_2=%.6g"
+                " input_scale=%.6g (harvested %d/%d main-layer values)",
+                _s2, _isc, len(_w2fix_s2), len(_w2fix_in))
+        except Exception as _e:  # noqa: BLE001
+            _w2fix_logger.warning("W2FIX failed: %s", _e)
         self.finalize_mega_moe_weights()
         logger.info_once("MTP draft model loaded: %d params", len(loaded_params))
         return loaded_params

A stronger general guard would be a post-load assertion that every parameter registered by a quant method was actually written by the loader — that would have surfaced this bug loudly on day one. Happy to discuss/iterate on either half.

Validation (details in #20)

Config	MTP pos-0 accept	single-stream tok/s
any W2 recipe, before	0%	~10
pro6000x2-tp2 W2 + fixes	83%	79–103
pro6000x1-24k + fixes	82%	93–121
stock FP4 TP2 + fixes	90%	131–136 (was 94–101)