qwen4_exp: serve the block-FP8 dense projections natively (+25% decode) - #392
qwen4_exp: serve the block-FP8 dense projections natively (+25% decode)#392gberasmus87 wants to merge 4 commits into
Conversation
…ense) checkpoints
FlashML-org#320 loads modelopt MIXED_PRECISION checkpoints by dequantizing the block-FP8 dense projections to bf16 at load, which doubles the bytes read on every decode step. Keep them quantized instead and let the existing Fp8Block linears consume them: 4.98 GiB -> 2.49 GiB of dense weights on the modelopt build of Qwen3.8-Flash-Next (156 FP8_PB_WO tensors, 2.67G elements). The checkpoint declares this per module - 48 .mlp.experts as NVFP4 and 156 attn/GDN projections as FP8_PB_WO - so config.py now reports attn_quant="fp8_block" independently of expert_quant, and gdn.py takes the block-fp8 path when either says so. It previously keyed off expert_quant alone, so a checkpoint with NVFP4 experts and block-FP8 dense never reached it. quant_linear.py's factories widen the same way, and qwen4_exp attention builds qkv_proj/o_proj through them instead of hardcoding the bf16 classes. The loader stops dequantizing when the declaration is present, keeps the weight_scale_inv tensors, and swaps in a fusion table matching the modules the model actually builds: the four-way in_proj fusion splits into an fp8 qkv|z GEMM plus a small bf16 b|a GEMM - the split gdn.py already implements for block-fp8, matching sglang/vLLM - and each fp8 group fuses its scale on the same axis as its weight. That split is what unblocks the bulk of this: b|a are bf16, so the old four-way cat mixed dtypes and forced the dequant. Every fp8 part is a whole number of 128-row blocks (10240/6144 and 12288/512/512), so the per-block scales concatenate exactly alongside the rows they describe. A checkpoint carrying weight_scale_inv WITHOUT declaring FP8_PB_WO still takes the FlashML-org#320 dequant path, so builds that quantize the dense side but describe it differently keep working unchanged. Verified on the modelopt checkpoint without a GPU: parse_config yields nvfp4/fp8_block, and "none" when the declaration is removed; iter_weights emits in_proj_qkvz [16384,2560] fp8 + scale [128,20], in_proj_ba [96,2560] bf16, qkv_proj [13312,2560] fp8 + scale [104,20], and fp8 out_proj/o_proj - the exact buffers Fp8BlockLinear declares. Serving numbers to follow.
|
Ran this on 2 x RTX 6000 Ada (sm_89, 48 GB each, PCIe 4 x16, no NVLink) with the lovedheart NVFP4-FP8 (modelopt MIXED_PRECISION) checkpoint, offload backend, KV pool 262,144 tokens, pinned PLE, 8 running requests, one card (TP=1) unless stated. Tree: main af71ba4 + TP (#385) + load-time per-tensor FP8 (#389) + this PR cherry-picked. Same harness as my other PR feedback: single-stream = median of three 256-token generations, aggregate = eight concurrent 256-token requests, TTFT on a ~1k-token prompt, residency = expert slots / 24,576, probe = the same 8-question greedy smoke test (thinking off).
So the win reproduces on Ada at TP=1, and it is the residency effect you describe: +3.9 points of experts on the card. In isolation the Triton block-FP8 GEMV is slower than bf16 cuBLAS on this card (3.78 vs 3.20 ms per 48-layer decode step at M=1, the micro-benchmark in #389), but the freed 2.5 GiB buys more than that costs. Per-tensor At TP=2 (both cards, 16 running, vision tower loaded) the block-FP8 linears have no parallel variant, so I carry this PR on my deploy branch with a precedence rule: Two notes on the diff from that merge: (1) routing the bf16 fallback of |
Both from @gdevenyi's review on 2 x RTX 6000 Ada, where this is carried on a deploy branch alongside FlashML-org#385 (TP). 1. Attention routed its bf16 fallback through the quantized factories too, which swaps in their generic fallback and drops the tensor-parallel classes FlashML-org#385 needs (per-rank local_output_sizes, row-parallel o_proj). That is exactly the path a rank takes under TP>1, since the block-FP8 linears have no parallel variant. The factories are now used only on the fp8_block branch; every other case keeps LinearColParallelMerged / LinearReplicated as before. 2. config.parse_config and weight._dense_is_block_fp8 read the same declaration through two independent code paths, each with its own copy of _FP8_BLOCK_ALGOS. That is safe only while they cannot disagree, and they can: a rank downgrading under TP>1 must have the modules it BUILDS and the buffers it LOADS downgrade together, or the buffers will not match. Both now resolve through one helper, config.dense_quant_mode, which owns the declaration test and the TP downgrade. The duplicate constant is gone. It reads TP through try_get_tp_info, not get_tp_info: Engine.__init__ sets TP info as its first statement so a rank always knows its size by the time this matters, but config parsing also happens with no engine at all (checkpoint conversion, tooling, tests) where get_tp_info raises. Verified on the modelopt checkpoint: parse_config still yields nvfp4/fp8_block; the two sides agree at TP=1 (both fp8_block) and at TP=2 (both downgraded); attention builds LinearColParallelMerged/LinearReplicated under bf16 and Fp8BlockColMerged/Fp8BlockLinear under fp8_block. tests/models/qwen4_exp/test_config.py + test_weight.py: 30 passed. The whole qwen4_exp suite reports 47 failed / 46 passed / 50 skipped both at the merge-base and with these fixes - identical sets, no regressions. Those failures are pre-existing and are an artefact of this box rather than the code: its single 24 GB card is 23.6 GB occupied serving a model, so the GPU-dependent tests cannot allocate. I have not been able to run them on a free card.
|
Thanks for taking both points. The One thing to flag on point 1, not a problem with this PR as it stands but a hazard when it meets #385. The new bf16 branch in self.o_proj = LinearReplicated(self.qo_attn_dim, config.hidden_size, has_bias=False)That is exactly what main does today, so this PR is correct on its own. But it is also the precise path a rank takes under TP>1 — the one your commit message calls out — and there On my deploy branch the same self.qkv_proj = LinearColParallelMerged(
config.hidden_size, qkv_sizes, has_bias=False, local_output_sizes=self._qkv_split
)
self.o_proj = LinearOProj(self.qo_attn_dim, config.hidden_size, has_bias=False)
I am not asking you to change anything here; #385 is mine and the merge is my problem. Flagging it so whichever of the two lands second does not silently keep the Also worth knowing, since your commit message mentions the two cannot be tested together on your box: the combination is running here. 🤖 Generated with Claude Code |
Follow-up to @gdevenyi's note on the previous fix. The bf16 branch kept o_proj as LinearReplicated, which is what main does today and is correct at TP=1, but it is the path a rank falls back to under TP>1 (FlashML-org#385), and there a replicated o_proj is wrong three ways at once: qkv_proj is column-parallel so each rank's attention output is its local head slice, o_proj therefore needs the sharded input dim, and the partial sums need an all-reduce. It also fails quietly, since a missing reduction still decodes to fluent-looking text. LinearOProj does all of that and degenerates to LinearReplicated at TP=1: div_even(x, 1) == x, and the all-reduce is skipped when tp_size == 1. So this is a no-op for main and only changes what FlashML-org#385 finds when the two meet, whichever lands second. It does mean get_tp_info() runs in __init__, but the same branch already does that two lines up through LinearColParallelMerged, so there is no new constraint: this path was engine-only before and still is. Config parsing, the one no-engine path that mattered, stays on try_get_tp_info. The comment above the branch also claimed a row-parallel o_proj that the code did not build; it now describes what is built. Verified on ailab1 (single RTX PRO 4000 Blackwell, TP=1): - tests/models/qwen4_exp/{test_config,test_weight,test_skeleton, test_qsa_backend}.py: 49 passed at 72773b0 and 49 passed with this change. That set includes test_qsa_layer_matches_hf_dense, which builds Qwen4ExpAttention on the card and checks the whole layer against the HF dense reference, so the new o_proj is exercised through a real forward. - A direct check on CPU and CUDA: the built o_proj is a LinearOProj with weight [hidden, qo_attn_dim] and local_input_size == qo_attn_dim, and its forward is bit-identical (max |diff| = 0.0) to a LinearReplicated carrying the same weight. Correction to the previous message: the 47 failures I attributed to VRAM contention were, for these four files at least, a missing ninja on PATH in my throwaway test venv. With the serving venv's bin on PATH the files pass in full.
|
Thanks, that's a fair flag and cheaper to fix now than to remember at merge time. Swapped in 2b2a96e: the bf16 branch now builds On the Verified at TP=1 on the Blackwell box: the four I'll take you up on the TP=2 run. The case I can't exercise is exactly this one, the bf16 fallback under TP, so: deploy/chatdnp with this commit, |
2 participants
Stacks on #320, so this branch contains that commit too. The diff to review here is the second commit (
da06515).What
#320 loads modelopt MIXED_PRECISION checkpoints by dequantizing the block-FP8 dense projections to bf16 at load, which doubles the dense bytes read on every decode step. This keeps them quantized and lets the existing
Fp8Blocklinears consume them directly.On the modelopt build of Qwen3.8-Flash-Next that is 156
FP8_PB_WOtensors, 2.67G elements: 4.98 GiB as bf16 becomes 2.49 GiB.Why this needed more than a config flag
The four-way
linear_attn.in_projfusion concatenatesin_proj_qkvandin_proj_z(bothFP8_PB_WO) within_proj_bandin_proj_a(bf16). You cannottorch.catfp8 with bf16, which is exactly why #320 dequantizes.gdn.pyalready implements the right answer for block-fp8 checkpoints: an fp8in_proj_qkvzGEMM plus a small bf16in_proj_baGEMM, matching sglang/vLLM. It was simply unreachable for this checkpoint, because it keyed offexpert_quant == "fp8_block"and here the experts are NVFP4 while only the dense side is block-FP8. So:config.pyreportsattn_quant="fp8_block"from the checkpoint's own per-moduleFP8_PB_WOdeclaration, independently ofexpert_quantgdn.pytakes the block-fp8 path when either says soquant_linear.py's factories widen the same wayattention.pybuildsqkv_proj/o_projthrough those factories instead of hardcoded bf16 classesweight.pykeeps theweight_scale_invtensors and fuses on a table matching the modules the model actually buildsEvery fp8 fusion group is a whole number of 128-row blocks (10240/6144, and 12288/512/512), so the per-block scales concatenate exactly alongside the rows they describe.
Measured
Single RTX PRO 4000 Blackwell 24 GB, TP=1,
--moe-backend hybrid, PCIe gen 3 x16, modelopt Qwen3.8-Flash-Next. Greedy, warmed, and counting every streamed delta: this model emits most tokens asreasoning_content, so counting onlycontentunderstates the rate badly.+24.7%, run-to-run spread 0.2%, identical output hashes across runs. TTFT unchanged at ~4.87s, as expected since this does not touch prefill. VRAM usage is also unchanged, because
memory_ratiohands the freed 2.49 GiB straight to the expert cache, which is most likely where the gain actually comes from on a PCIe-3 box.That also explains the flat result @gdevenyi measured on 2x RTX 6000 Ada at TP=2 for the loading PR: with 96 GB and the dense read split across two GPUs, neither the bandwidth nor the residency effect has room to matter. The win is regime-dependent, and largest where VRAM is tight and the interconnect is slow.
Compatibility
A checkpoint carrying
weight_scale_invwithout declaringFP8_PB_WOstill takes #320's dequant path unchanged, so builds that quantize the dense side but describe it differently keep working.