feat(qwen4_exp): load block-fp8 dense projections natively

[jason-fxz]

Loads modelopt MIXED_PRECISION Qwen3.8-Flash-Next checkpoints whose attention / GDN projections are 128x128 block-fp8 (FP8_PB_WO), e.g. lovedheart/Qwen3.8-Flash-Next-NVFP4-FP8. The model already builds the fp8 buffers (#418); the reader still fused GDN in_proj_{qkv,z,b,a} into one bf16 tensor and died in torch.cat when in_proj_qkv was fp8.

• _DenseFuser replaces _FUSIONS / _try_fuse: parts come from packed_modules_mapping, .weight and .weight_scale_inv fuse per kind, GDN splits into in_proj_qkvz (fp8) + in_proj_ba (bf16) when the QuantConfig says so.
• Each part is checked against its scheme (dtype, 128-row multiple, scale only where declared); a mismatch raises with the module name.
• Released NVFP4 / FP8 checkpoints keep bf16 dense projections; their output is unchanged.

Tests: synthetic bf16, NVFP4 and block-fp8 checkpoints in tests/models/qwen4_exp/test_weight.py (emitted keys == model state dict, fp8 slice-back, four rejection cases); scheme resolution in test_config.py.

Tested on H100 80GB / Xeon 8480+, driver 580.95.05, torch 2.11.0+cu130, TP=1, offload + --moe-cache-auto: lovedheart/Qwen3.8-Flash-Next-NVFP4-FP8 passes tests/e2e/test_aime.py, all 926 emitted tensors match the model's buffers, 18 fused fp8 tensors are bit-exact against the raw shards; RadixArk and nvidia NVFP4 unchanged. FREETOKEN_TEST_MODEL=<dir> FREETOKEN_TEST_MOE_CACHE_AUTO=1 pytest -x -s tests/e2e/test_aime.py.

Depends on #427. Supersedes #392 / #320 by @gberasmus87, which reached the same reader design before #418. @gberasmus87 @gdevenyi: a run on your cards would be a useful check.

[gberasmus87]

Happy to see this land properly — _DenseFuser deriving the parts from packed_modules_mapping is better than the fusion tables I had, and #427 is the right home for handing readers the QuantConfig. I'd built around that gap with a function-level import of engine.config, which I wasn't happy about.

One behaviour question, not an objection.

check() has:

if scheme is None:
    if is_fp8:
        raise ValueError(f"{name} is {tensor.dtype} but the checkpoint's quant config declares {module} unquantized")

So a checkpoint that stores block-fp8 dense weights without declaring them per module now fails to load. #320 dequantized that case to bf16 at load and served it, which is how I first got the lovedheart build running — my early notes on it say the block-fp8 attention/GDN tensors were not flagged in quantized_layers, and I carried a dequant fallback specifically for that. The copy I have today declares all 156 as FP8_PB_WO, so either it was re-uploaded or I misread it back then, and I no longer have a checkpoint that exercises the undeclared path.

Raising it only because the failure mode changes from "loads, dequantized" to "raises", and community requants are the population most likely to store quantized weights without describing them. If that's a deliberate call — an explicit error beats a silent divergence between the buffers the model built and what the reader emits — then it's the right one and worth a line in the docstring so the next person doesn't re-add a fallback. If it isn't, a scheme is None and is_fp8 branch that dequantizes rather than raises would preserve it.

On the run you asked for: yes. My box is the regime this feature is actually for — a single 24 GB RTX PRO 4000 Blackwell on a PCIe gen 3 board, offload backend, where the model is nowhere near resident and expert-cache residency is the binding constraint. That's the opposite end from your H100 80GB and @gdevenyi's 2 x 48 GB Ada, where the freed dense bytes have much less to buy. On the earlier form of this reader I measured 27.4 -> 34.2 tok/s there, and the gain tracked moe_cache_auto going 3057 -> 4041 entries rather than the halved dense read. I'll run this branch on the lovedheart checkpoint and post single-stream decode, TTFT, and the resulting cache entry count against main.

Also flagging, since #392's reader half is superseded by this: the LinearOProj fix from that PR is unrelated to the reader and still applies. I've split it out as #429 so it isn't lost when this closes #392. @gdevenyi raised it there originally as a merge hazard against #385.

[jason-fxz]

Deliberate. Since #418 a module the config does not list is built bf16, so an fp8 tensor there means the checkpoint's declaration and its tensors disagree. We treat that as a broken checkpoint, not a layout to support; the error names the module to fix. Workaround: add the module to quantization_config.quantized_layers ({"quant_algo": "FP8_PB_WO", "group_size": 128}) in config.json, and in hf_quant_config.json if present.

Thanks for the run and for splitting out #429.

[gberasmus87]

Ran it. The reader works — lovedheart/Qwen3.8-Flash-Next-NVFP4-FP8 loads cleanly on a single 24 GB RTX PRO 4000 Blackwell (sm_120, PCIe gen 3 x16, EPYC 7302, driver on CUDA 13.0, torch 2.11.0+cu130), offload/hybrid, TP=1: expert banks build, --moe-cache-auto resolves moe_cache_size=3741 num_pages=1025, CUDA graphs capture at [1,2,4], 2.18 GiB free after init, API server is ready to serve. Generation produces correct output. No loader complaints at all.

I can't give you decode numbers, though, because this tree runs 40-60x slower than the build I'm coming from on this configuration, and I think that matters more than the numbers I owed you.

What I measured

Same box, same checkpoint, same flags, back-to-back:

build	engine-reported decode throughput
v0.1.2 + my local patches (what I serve today)	33.4 / 35.2 / 36.3 / 37.5 / 38.3 / 38.9 tok/s
this PR's tree (b16877b = main + #427 + #428)	0.17 / 0.43 / 0.55 / 0.56 / 0.74 / 1.02 tok/s

End-to-end, counting every streamed delta (this model emits most tokens as reasoning_content): 33.5 tok/s mean on the old build across three prompts; on the new tree a 64-token completion takes minutes. The slow figures come from three independent launches of the new tree, so it is not a cold-start artifact.

What it looks like while it is slow

• GPU utilisation 0% during generation (the old build sits at 74-95%). Weights are resident — 21,962 MiB.
• One CPU thread pinned at ~100%, state=R, wchan=0 — a userspace spin on the engine's main thread. The other 76 threads are idle, and the 15-thread pinned CPU MoE pool is not busy, so neither compute path is doing the work.
• Not JIT compilation: ~/.triton/cache stopped growing (4015 files, 145 MB, unchanged over a 30 s sample) while that thread burned ~1,034 s of CPU.
• Not expert-cache starvation: moe_cache_size=3741 here versus 4041 on the old build — 7% down, nowhere near enough to explain it.
• Not the deprecated flag: I re-ran with --nvfp4-backend triton removed entirely (quant_backend=None). It still resolves MoE experts: nvfp4 via triton and is still 0.55 tok/s.

What I can't tell you

I cannot attribute this to #428, and I do not think it is #428. My baseline is v0.1.2 plus local patches rather than upstream main, because main cannot load this checkpoint at all — that is the thing this PR fixes. So the comparison spans v0.1.2..b16877b, which includes #418. #428's own diff is confined to the reader and the tests, and the reader demonstrably does its job.

I also could not get a stack. py-spy needs ptrace and the box is ptrace_scope=1 with only a NOPASSWD systemctl restart; running the engine under py-spy record --subprocesses attached and sampled but never flushed its output before the window closed.

Why this configuration is probably not covered

You tested on an H100 80 GB and @gdevenyi is on 2 x RTX 6000 Ada 48 GB. On both, this model is close to resident and the offload expert path barely engages. Here it carries essentially the whole model — 123 GiB of weights against 24 GB of VRAM on a gen-3 board, which is the regime the offload backend exists for, and the one where the dense-side savings in this PR are worth the most. If the slowdown is in that path, it would be invisible on either of your boxes and fatal on mine.

Happy to bisect v0.1.2..main if that is useful — it is roughly 34 commits and each cycle is a ~5 minute install plus load here, so it is a few hours rather than a few minutes, but I would rather do that than have you chase it blind. Equally happy to run any specific commit, flag combination, or instrumented build you want to point at. Say which and I will get you numbers.

[gberasmus87]

Narrowed it, and it is not this PR — moved to #436 so it does not sit on your thread.

It is --moe-strategy hybrid. Same install, same checkpoint, only the strategy flag changed:

--moe-strategy	decode	GPU util
hybrid	0.17 - 1.02 tok/s	0%
offload	23.07 tok/s	85%

So this PR's reader is fine on my hardware — the offload run is the same install serving the same block-fp8 checkpoint correctly at a sane rate. Consider the run you asked for delivered: #428 loads and serves lovedheart/Qwen3.8-Flash-Next-NVFP4-FP8 correctly on a single 24 GB Blackwell, and the only reason I could not hand you a decode number the first time was the unrelated hybrid problem.

The tell is that under hybrid one thread spins at ~100% while the 15-thread CPU MoE pool sits idle and the GPU does nothing; under offload that same thread reads 3%. Details, ruled-out causes and the offer to bisect are all in #436.

[gdevenyi]

I cannot give you the lovedheart/Qwen3.8-Flash-Next-NVFP4-FP8 run you asked for — our box is a production server on 2 x RTX 6000 Ada and a second checkpoint does not fit beside the resident one. What I can give you is the thing your setup cannot reach: how _DenseFuser meets TP=2, and it is better news than I expected.

_DenseFuser is the right home for a bug I hit this week. Our deployment carries load-time synthetic FP8 for the dense projections (#389). The model side gates it per module — attention.py tests config.quant.scheme_for(f"{prefix}.qkv_proj") is None, gdn.py the same for in_proj_qkvz — but the reader side gates it on one model-wide boolean with no scheme test. Invisible on every released checkpoint, because the skip lists cover all dense projections. On a checkpoint like yours it emits weight_scale for a module the model built from a declared scheme, and you get

RuntimeError: Unexpected keys in state_dict: [...]

which is exactly the failure I spent a GPU window on in the lm_head case. Your fuser.scheme(module) makes the fix one condition instead of a name-mapping guess, so I would rather land it on top of this PR than patch the flat _FUSIONS table underneath it. I flagged the gap on #389 as unfixed; this is where it should be fixed.

Merge state against our tree. Merging pr/428 onto our deploy line (upstream main fb7f732 + #385 TP + #389 FP8 dense) gives 3 conflict hunks, all in weight.py, all at the same seam. Nothing structural — the resolution is:

1. Keep your _DenseFuser as the fusion mechanism; drop our _FUSIONS / _try_fuse entirely.
2. Fused output then passes through _shard(name, tensor, config, tp.rank, tp.size) before it is yielded. This PR has no sharding hook — it yields the fused tensor straight out — so under TP>1 the shard has to happen after the fuser, not inside it.
3. _fp8_dense(...) gates on fuser.scheme(module) is None, which is the fix above.
4. Re-add the torch.cuda.empty_cache() after the load. That one is ours and specific to synthetic FP8: the bf16 originals sit in the caching allocator, and our expert-cache planner sizes itself from free VRAM after load, so without it the planner sees slack instead of the halved dense footprint.

The one thing to decide before this lands. Line 205:

if get_tp_info().size > 1:
    raise NotImplementedError("qwen4_exp weight loading supports TP=1 only")

That is correct for main today and it is a hard stop for #385. Whichever of the two lands second owns removing it. I am happy for that to be my problem — I mention it only so the guard is a known hand-off and not a surprise.

Also note #429 (o_proj row-parallel) has the matching issue on the model side: the block-fp8 branch inherited from #392 still builds o_proj replicated via make_replicated_quant. Consistent while #392 declares itself TP=1-only; a live bug the moment this PR removes that restriction. I left the same note there.

Not adopting it into our deployment yet, to be explicit about why: our candidate branch has four post-merge fixes on it and has not completed a GPU startup since. A 429-line reader rewrite goes in after that is validated, not before. Our checkpoint has bf16 dense projections, so this PR changes nothing for us functionally — it is the design I want under #389, not a capability we are missing.

🤖 Generated with Claude Code

https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt