Skip to content

DS4 MTP drafter loaded with wrong quant method (ignored_layers 'mtp.*' never matches) + uninitialized per-tensor scales: 0% MTP acceptance under W2, allocation-dependent degradation on stock FP4 #20

Description

@macdad222

Summary

The DeepSeek-V4-Flash NVFP4 checkpoint ships its MTP drafter experts in MXFP4 format (FP4 weights + E8M0 block-32 scales, no per-tensor scales), but the loader routes them to the NVFP4 quant method (ModelOptNvFp4FusedMoE). Two compounding problems result:

  1. The checkpoint's quantization_config.ignored_layers contains "mtp.*", which never matches the drafter's real module path model.layers.43.ffn.experts, so the MXFP4 → NVFP4 misroute happens silently. E8M0 scale bytes are reinterpreted as E4M3.
  2. The NVFP4 method expects w13_weight_scale_2, w2_weight_scale_2, w13_input_scale, w2_input_scale — none exist in the checkpoint for the drafter. They are allocated with torch.empty and never filled, and nothing validates this after load. Whatever garbage is resident in those allocations becomes the drafter's scales.

Because the failure depends on uninitialized memory contents, symptoms are allocation-dependent and nondeterministic across configurations:

  • Benign garbage (tiny denormals ~1e-43): drafter logits underflow but argmax roughly tracks → ~79% pos-0 acceptance that looks healthy. We believe the published MTP baselines were measured in this state.
  • Hostile garbage (~1e31, e.g. after the W2 planes builder churns the allocator): inf alphas → NaN hidden states → drafter emits token 0 forever → 0% acceptance. This makes every W2 recipe lose all speculation benefit (~10 tok/s instead of 100+), while appearing perfectly coherent otherwise.
  • Toggling PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True on stock FP4 TP2 moves acceptance between 79% and 57% with no legitimate config change — a one-line demonstration that model behavior depends on allocator state.

With the drafter correctly loaded as MXFP4, true healthy acceptance is ~90% pos-0 / ~61% pos-1 — about 10 points above the luck-degraded baseline, worth ~35% single-stream tok/s in our measurements.

Environment

  • 2x RTX PRO 6000 Blackwell (96 GB; 1x 600W Workstation + 1x 300W Max-Q), TP2, PCIe Gen4
  • vllm-moet-recipes:v024 image built from this repo (vLLM 0.24.0 base)
  • Model: DeepSeek-V4-Flash NVFP4 (149 GB), recipes deepseek-v4-flash/pro6000x2-tp2, pro6000x1-24k, pro6000x1-512k

Evidence trail

Instrumented rejection_sampler, the MTP proposer, DeepSeekV4MultiTokenPredictorLayer, and ModelOptNvFp4FusedMoE.process_weights_after_loading (dump-over-mount, no rebuild):

  • Under W2=1: drafter input target_hidden_states healthy; NaN first appears inside mtp_block's FFN; MoERunner output NaN.
  • Expert scale params at PWL entry, layer 43 (drafter), W2=1 boot: w13_weight_scale_2 absmax≈1e31, w13_input_scale contains inf — pure torch.empty residue. Same params on a W2=0 boot: denormals ~1e-43. Main-stack layers (0–42) load real values (~2.9e-4 / ~5.9e-3).
  • moe_w2_cubit.is_w2_layer("model.layers.43...") correctly returns False; the drafter takes the stock NVFP4 branch — confirming the misroute is in get_quant_method, not W2 code. W2 merely changes which garbage the empty tensors inherit.
  • Checkpoint inspection: MTP expert tensors are F8_E8M0-scaled block-32 MXFP4; quantization_config.moe_quant_algo == "NVFP4" applies globally because ignored_layers: ["mtp.*"] matches nothing.

Fix we validated (two parts)

1. Route the drafter to MXFP4 in vllm/models/deepseek_v4/quant_config.py::get_quant_method — treat experts whose layer index >= num_hidden_layers as MTP and return Mxfp4MoEMethod (the format they are stored in) instead of ModelOptNvFp4FusedMoE:

_is_mtp = False
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)
if self.moe_quant_algo == "NVFP4" and not _is_mtp:
    return ModelOptNvFp4FusedMoE(...)
return Mxfp4MoEMethod(layer.moe_config)

2. Defense in depth (optional once #1 lands): in the MTP load_weights, fill the drafter's per-tensor scale params from main-layer means if they were never loaded, so torch.empty residue can never become live scales. More generally: a post-load assertion that every parameter the quant method registered was actually initialized would have caught this class of bug loudly.

Results after fix (same hardware, same checkpoints)

Config MTP pos-0 acceptance single-stream tok/s
Any W2 recipe, before 0% ~10
pro6000x1-24k + fix 82% 93–121 (matches your inline ~102–111 notes)
pro6000x2-tp2 W2 + fix 83% 79–103
Stock FP4 TP2 + fix 90% pos-0 / 61% pos-1 131–136 (was 94–101)

Happy to open a PR with the routing fix and/or the load-validation guard, and to share full instrumentation diffs and logs. Also flagging: the checkpoint publisher's ignored_layers: ["mtp.*"] pattern appears wrong for any vLLM-style module naming, so an upstream report there may also be warranted — but the loader failing silently on missing quant params is the more dangerous half.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions