[#16196][fix] Support MIXED_PRECISION experts-only NVFP4 DeepSeek-V4 checkpoints

[waynehacking8]

Summary by CodeRabbit

• Bug Fixes
  • Improved mixed-precision model loading and runtime behavior for DeepSeek V4.
  • Fixed handling of fused attention weights when scale metadata is unavailable.
  • Improved compatibility with FP8-block-scaled checkpoints and NVFP4 expert repacking.
  • Ensured attention, shared-expert, and adapter projections use consistent quantization settings.

Description

nvidia/DeepSeek-V4-Pro-NVFP4 (the ModelOpt experts-only NVFP4 repack) cannot be loaded (#16196). The checkpoint is MIXED_PRECISION: only the routed experts are re-quantized to NVFP4, and the modules in exclude_modules (attention, shared experts, MTP e_proj/h_proj) are still stored at the FP8 base model's precision (fp8_e4m3 weight + ue8m0 block scale) -- not BF16. Two failures follow on main:

1. Construction dies at DeepseekV4DecoderLayer's "MIXED_PRECISION is ambiguous" assert.
2. Even past construction (as on the rc15.post1/rc20 containers), the excluded attention modules are built unquantized, so the fused-A loader dereferences a missing module.weight_scale (AttributeError, the reported crash), and at runtime the MLA fused-FP8-q-quant contract breaks (_q_branch assert: q_b_proj BF16 while _is_fused_q_fp8_quant_enabled() is True).

This PR resolves both:

• _mixed_precision_base_quant_config(): under MIXED_PRECISION, derive an FP8_BLOCK_SCALES config from the checkpoint's embedded base quantization_config (quant_method=fp8 + weight_block_size) and build attention, shared experts, and MTP e_proj/h_proj with it -- matching how the native FP8 checkpoint flow builds them.
• _get_decoder_layer_quant_config(): under MIXED_PRECISION, resolve the decoder-layer algo from the per-layer quantized_layers table (routed experts, NVFP4), removing the ambiguity assert failure while keeping is_nvfp4 semantics for the MoE.
• _get_experts_quant_config(): also accept the checkpoint-native key form layers.N.ffn.experts emitted by these repacks (quant_config_dict keys are stored verbatim, so the internal-form lookup was silently missing and returning the global config).
• Fused-A loader: guard on hasattr(module, "weight_scale") and dequantize with the existing weight_dequant kernel as a defensive fallback for any excluded-but-FP8-stored module that is still built unquantized.
• _set_deepseek_v4_routed_moe_quant_config(): detect the MTP routed-expert layout separately (commit authored by @d3nb, taken from Detect MTP routed-expert layout separately (dense NVFP4 vs MTP MXFP4) waynehacking8/TensorRT-LLM#1). These repacks re-quantize only the dense routed experts to NVFP4 (U8) and leave the MTP routed experts at the base model's MXFP4 (I8); deriving a single layout from layers.0 assigned NVFP4 to the MTP layers too and crashed in fused_moe load_quant_scales. The MTP layer indices now get their own config (MXFP4 via get_mxfp4_quant_algo, group_size 32), falling back to the dense config when the layouts match or the probe tensor is absent.

MTP speculative decoding is now in scope: the last commit is @d3nb's per-region layout detection, merged from waynehacking8#1 with his authorship intact.

Fixes #16196.

Test Coverage

• Real weights, end-to-end (4x B300, TP4, moe_config.backend=TRTLLM): validated by the issue reporter @d3nb on the actual 851 GB nvidia/DeepSeek-V4-Pro-NVFP4: weights load to 100%, server starts, chat completions generate correctly; q_b_proj/shared experts end up FP8 with weight_scale present and the previous _q_branch assert no longer fires. Details: [DeepSeek-V4] Loading nvidia/DeepSeek-V4-Pro-NVFP4 crashes: 'Linear' object has no attribute 'weight_scale' in fused_a loader (MIXED_PRECISION attn excluded → built BF16 but stored as FP8) #16196 (comment) (and the follow-up validation comment).
• Synthetic tiny replica (RTX PRO 6000, SM120): a faithful shrunk replica of the checkpoint (same tensor name set, same dtypes including ue8m0 scales, same hf_quant_config.json shape): construction passes, all weight groups load, the fused-A fp8 payload is bit-exact vs the checkpoint and the fused scales match exactly; decoder layers report is_nvfp4=True with experts packed NVFP4.
• The dequant fallback path (weight_dequant on fp8_e4m3 weights with ue8m0 scales) matches a blockwise reference exactly, including non-multiple-of-128 edge blocks.
• MTP (real weights, 4x B300, TP4): @d3nb validated the layout detection end-to-end (rc15.post1 container carrying the equivalent change, plus the construction fix from this branch): loads 100%, serves, generates with MTP=1 and MTP=3 (mtp_eagle_one_model builds one shared MTP layer); MTP accept_len ~2.86.

PR Checklist

• PR title and description follow the template
• Commits are signed off (DCO)
• Verified on real weights (reporter) and a synthetic replica (author)

[coderabbitai]

📝 Walkthrough

Walkthrough

DeepSeek V4 mixed-precision handling now derives FP8 block-scale configurations, dequantizes fused-A weights for unquantized destinations, and applies consistent quantization settings to attention, shared-expert, and MTP projections.

Changes

DeepSeek V4 mixed-precision quantization

Layer / File(s)	Summary
Mixed-precision configuration resolution tensorrt_llm/_torch/models/modeling_deepseekv4.py	Adds helpers for FP8 block-scale configuration and updates per-layer expert quantization resolution.
Fused-A weight loading tensorrt_llm/_torch/models/modeling_deepseekv4.py	Dequantizes fused-A checkpoint weights when the destination module lacks weight_scale.
Projection quantization wiring tensorrt_llm/_torch/models/modeling_deepseekv4.py	Applies derived mixed-precision configuration to attention, shared-expert, and MTP projections.

Estimated code review effort: 4 (Complex) | ~45 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 41.67% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	The changes match #16196 by fixing excluded-module FP8 loading, per-layer NVFP4 selection, and checkpoint-native expert key handling.
Out of Scope Changes check	✅ Passed	The modified loader and quantization paths all support the DeepSeek-V4 NVFP4 loading fix; no unrelated code changes are evident.
Title check	✅ Passed	The title is concise, specific, and accurately summarizes the main fix for MIXED_PRECISION DeepSeek-V4 checkpoint loading.
Description check	✅ Passed	The description follows the template and clearly covers the problem, solution, test coverage, and checklist items.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}