Skip to content

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

Open
jason-fxz wants to merge 1 commit into
mainfrom
feat/qwen4-exp-fp8-dense
Open

feat(qwen4_exp): load block-fp8 dense projections natively#428
jason-fxz wants to merge 1 commit into
mainfrom
feat/qwen4-exp-fp8-dense

Conversation

@jason-fxz

Copy link
Copy Markdown
Collaborator

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

Copy link
Copy Markdown

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
jason-fxz added this pull request to stack #432 September 10, 2026 08:44
Base automatically changed from refactor/quant-config-handoff to main September 10, 2026 08:46
@jason-fxz
jason-fxz force-pushed the feat/qwen4-exp-fp8-dense branch from b16877b to 540b2ce Compare September 10, 2026 08:46
@jason-fxz

Copy link
Copy Markdown
Collaborator Author

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

Copy link
Copy Markdown

Ran it. The reader workslovedheart/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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants