qwen4_exp: load modelopt MIXED_PRECISION (NVFP4 experts + block-FP8 dense) checkpoints

[gberasmus87]

What

Adds support for loading Qwen3.8-Flash-Next checkpoints exported with modelopt MIXED_PRECISION quantization - NVFP4 routed experts + 128x128 block-FP8 dense (attention / GDN) weights. The common community build of this shape is lovedheart/Qwen3.8-Flash-Next-NVFP4-FP8, the ~124GB quant that fits the model on a 24GB GPU via the offload backend.

Before this, such a checkpoint fails to load in two places:

1. Config detection - parse_config only recognises a top-level quant_algo of fp8 or *fp4*. MIXED_PRECISION declares the per-module algo in quantized_layers instead, so expert_quant fell through to "none" and the offload cache tried to load the NVFP4 experts as bf16 banks -> ValueError: Missing MoE expert source layers.
2. Dense weight loading - the dense attn / GDN projections are stored as 128x128 block-FP8 (.weight_scale_inv), but iter_weights assumed all non-expert weights are bf16 and passed them straight to _try_fuse -> RuntimeError: Promotion for Float8 Types is not supported, attempted to promote Float8_e4m3fn and BFloat16 on the GDN in_proj fuse.

Changes

• config.py: detect quant_algo == "mixed_precision", read quantized_layers to set expert_quant = "nvfp4" (experts stay native for the offload cache), and keep every dense module bf16.
• weight.py: _load_maybe_block_fp8 dequantizes a dense .weight to bf16 when a sibling .weight_scale_inv is present (reusing the existing dequant_block_fp8 kernel), and .weight_scale_inv is added to _SCALE_SUFFIXES so the scale tensor is not yielded as a weight. Only block-FP8 is added - it is the only dense quant these checkpoints use.

Testing

Loads and serves lovedheart/Qwen3.8-Flash-Next-NVFP4-FP8 (~124GB) on an RTX PRO 4000 Blackwell (24GB VRAM) backed by 172GB system RAM, with --moe-backend hybrid - the offload backend holds the ~124GB checkpoint resident in host RAM and streams experts to the GPU per token. Output is coherent and correct (reasoning / math / coding verified) and tool-calling works, ~22 tok/s single-stream. The new branches are gated on mixed_precision / a present .weight_scale_inv, so no other checkpoint path changes.

[gdevenyi]

Tested with lovedheart/Qwen3.8-Flash-Next-NVFP4-FP8 (123 GiB: 192 NVFP4 expert shards, 4 block-FP8 dense shards with weight_scale_inv, 10 FP8 PLE shards) on 2 x RTX 6000 Ada at TP=2 (#385), --moe-backend offload --ple-backend pinned --num-tokens 262144 --memory-ratio 0.94 --max-running-requests 16. The patch did not apply onto my TP-modified weight.py, so I merged the mixed_precision config branch and the weight_scale_inv dequant by hand; same logic.

It loads and serves. Because the FP8 dense projections are dequantized to bf16 at load, speed is the same as the modelopt NVFP4 build under the same flags:

checkpoint	single-stream tok/s	8 concurrent tok/s	expert residency	TTFT (1k)	8-question probe
RadixArk/Qwen3.8-Flash-Next-NVFP4	90.8	326.0	94.5%	0.81 s	6/8
lovedheart NVFP4-FP8 via this PR	90.2	325.6	94.6%	0.82 s	7/8

Greedy 256-token continuations of three prompts differ from the modelopt build after 7 to 29 words (different dense weights, so expected); the probe difference is within what eight questions can resolve. The obvious follow-up on top of this is keeping the dense weights in FP8 (fp8_block linears) instead of dequantizing: the decode step here is bound by the 8 GiB bf16 dense read per token at TP=1.

[gberasmus87]

@gdevenyi Thanks for testing it on hardware and a checkpoint I don't have. That's a more useful signal than anything I could produce on my own.

On the patch not applying: #320 is mergeable/clean against current main (2 files, +39/-2, both under models/qwen4_exp/). I think the conflict is #385 rather than this branch - it's open, unmerged, and changes models/qwen4_exp/weight.py too (+64/-7). Whichever of the two lands second will need the merge; happy for that to be this one if #385 is closer to ready.

Agreed the divergence is expected - different dense weights, and greedy decoding amplifies the first differing logit. I wouldn't read much into 7/8 vs 6/8 either way at that sample size.

On keeping the dense weights in FP8 rather than dequantizing at load: agreed, and I'll open it as a separate PR rather than growing this one, since it changes the compute path rather than the loading path and deserves its own review and bisect point.

A second data point on the sizing, since my checkpoint is the modelopt MIXED_PRECISION build rather than lovedheart's: there the block-FP8 dense is 156 tensors, all F8_E4M3, 2.67G elements - 2.49 GiB stored, 4.98 GiB once expanded to bf16 at load. Keeping them quantized would halve that read and hand back ~2.5 GiB, which on a 24 GB card is worth having by itself. I'm at TP=1 on a single RTX PRO 4000 Blackwell with --moe-backend hybrid, so the dense read isn't split the way it is in your TP=2 run. I'll benchmark it there and post numbers with the PR.

[jason-fxz]

The config half is now covered on main by #418/#426 (the expert kind comes from the checkpoint QuantConfig), and the model now allocates fp8 weight + block-scale buffers for the FP8_PB_WO dense projections, so a bf16 dequantized weight no longer matches what those layers expect and nothing would fill the scale.