feat(qwen4_exp): load-time per-tensor FP8 dense projections (W8A8 via _scaled_mm)

[gdevenyi]

What this adds

FREETOKEN_FP8_DENSE=1: load-time per-tensor FP8 for the bf16 attention / GDN projections of qwen4_exp, run as cuBLASLt W8A8 GEMMs (torch._scaled_mm). Opt-in, default off, no checkpoint change: the weight reader quantizes qkv_proj, o_proj, GDN in_proj (q|k|v|z; the b|a gate rows stay bf16 as in_proj_ba, as in the block-fp8 checkpoints and in sglang / vLLM) and out_proj to e4m3 with one fp32 scale each, after TP sharding. layers/fp8_dynamic.py holds the op: a dynamic per-tensor activation scale (one fused Triton launch at decode sizes: amax pass, then the cast; a torch reduction plus a cast kernel above 64k elements), the _scaled_mm call, and the all-reduce for the row-parallel case. No host sync anywhere, so the decode path is CUDA-graph safe; the branch between the two quant paths is on the tensor shape, never on its values. Requires sm_89+ (_scaled_mm's floor).

Stacked on #385 (the TP commits): the column-merged / row-parallel classes shard the same way its bf16 ones do.

Why

On 2 x RTX 6000 Ada (sm_89, torch 2.11.0+cu130, flashinfer 0.6.18) these projections are 2.67 GB of the ~4 GB a TP=2 rank reads per decode token. Micro-benchmark at the per-rank shapes (12 x [6656x2560] + [2560x3072], 36 x [8192x2560] + [2560x3072]), weight rotations larger than the 96 MB L2, activation quantization outside the timed region for the raw rows:

route, 48 layers per rank	M=1	M=8	M=16
bf16 cuBLAS	3.20 ms	3.38 ms	3.38 ms
FreeToken block-FP8 Triton (fp8_block_linear.py)	3.78	4.59	4.58
FreeToken per-row W8A16 (fp8_pertensor_linear.py)	3.70	3.07	3.06
FreeToken per-tensor W8A8 wrapper (same file)	3.62	3.59	3.58
raw torch._scaled_mm, per-tensor scales	1.94	1.93	1.92

The FP8 tensor cores are fine here; the existing Triton FP8 kernels reach 20-60% of bf16's bandwidth on this stack and the per-tensor wrapper spends ~24 us per call before cuBLASLt, so the FP8 checkpoints are slower than bf16 at decode on this card. This path takes the direct route.

Measurements

A/B on the machine above (RadixArk/Qwen3.8-Flash-Next-NVFP4, TP=2, --moe-backend offload --ple-backend pinned --num-tokens 262144 --memory-ratio 0.94 --moe-prefill-hit-d2d --max-running-requests 16 --cuda-graph-max-bs 16, vision tower loaded), both runs from the same build in the same session:

	bf16 dense	FREETOKEN_FP8_DENSE=1
single-stream decode, median of 3 x 256 tokens	89.8 tok/s	99.2 tok/s (+10.4%)
8 concurrent, aggregate	323.7 tok/s	331.9 tok/s (+2.5%)
TTFT, 1.8k-token prompt	0.84 s	0.85 s
expert residency (--moe-cache-auto)	91.9% (22,594 slots)	95.8% (23,539 slots) after 3f8f249 (a separate restart, same flags; single-stream on that build 99.6 tok/s), 91.4% before (see below)
8-question probe, thinking off	6/8	6/8, same answers
greedy 256 tokens: essay / code / 1k-prompt summary	reference	coherent; first divergence after 13 / 44 / 7 words

For the greedy rows the bf16 run-to-run floor at TP=2 is 84 words / identical / identical, so FP8 changes the sampling trajectory more than the noise does, as expected from different numerics; the outputs stay on topic (the code case continues the same memoised Fibonacci past where bf16 stopped). The +10% matches the micro-benchmark: ~1.3 ms per step per rank (4 GEMMs per layer, 48 layers) of an ~11 ms step.

Residency first fell (22,594 -> 22,458 slots) although the reader emits 1.25 GiB less per rank: in_proj_ba.weight was produced as t[qkvz:].contiguous(), and .contiguous() on an already-contiguous row slice returns the view, so every GDN layer's 48 bf16 gate rows kept the whole sharded bf16 in_proj alive next to its fp8 copy (36 x 42 MB = 1.5 GiB per rank). 3f8f249 clones the slice (regression assertion in the test); the planner then resolves 23,539 slots, 95.8% residency, with the same 2.6 GiB headroom.

Limits

• Per-tensor dynamic activation scaling across the batch is the coarsest scheme there is (one outlier channel in one row sets the scale for all rows). Row-wise scales were measured on torch 2.12.1+cu130 (where the sm_89 row-wise path is on-stream; pytorch@252bb4a is in v2.12.1), same shapes, GEMM only, 12 weight rotations. Absolute times in this measurement come from that separate torch build and differ from the table above (2.44 vs 1.94 ms for the same per-tensor route; cause not identified), so compare within the row only: M=1 / 8 / 16 per-tensor 2.44 / 2.39 / 2.40 ms per step per rank vs row-wise 2.41 / 2.40 / 2.38 (same), M=2048 18.5 vs 67.4 ms (the CUTLASS row-wise kernel is 3.6x slower than cuBLASLt at prefill sizes), relative error vs bf16 at M=16 with outlier channels 0.034-0.037 for both schemes. So row-wise buys nothing on this card, and it needs torch >= 2.12, which the torch<2.12 pin (sglang-kernel 0.4.5) blocks anyway; per-tensor stays.
• The quality gate is a smoke test: the API has no logprobs, so it is the 8-question probe plus greedy comparisons, not perplexity. Numbers above.
• Only the four projection families; shared expert, hyper-connections, lm_head and PLE stay bf16. A block-FP8 checkpoint through this path is double quantization (block-FP8 dequantized, then re-quantized per tensor). Measured on lovedheart's MIXED_PRECISION build (NVFP4 experts, block-FP8 dense) on the same box, TP=1 on one card, 8 running: dequantize-to-bf16 58.4 tok/s / 135.9 at 8 concurrent / 36.6% residency, qwen4_exp: serve the block-FP8 dense projections natively (+25% decode) #392's native block-FP8 69.9 / 150.2 / 40.5%, this path 71.7 / 153.4 / 40.4%, 8-question probe 7/8 for all three. At TP=2 with the production flags this path gives 99.8 / 336.7 / 95.8%, the same as the RadixArk bf16-dense build (99.2-99.6 / 331.9). Probe 6/8 there. The 6/8 at TP=2 is the prime-count question (primes in 10..30, answered 5): the RadixArk build also misses it at TP=2 in bf16 and in FP8, the bf16-dequant lovedheart run at TP=2 scored 7/8, and every TP=1 run scored 7/8. It is a borderline question that flips with the numeric path (TP=2 reduction order, FP8), and this is a one-question smoke test, not a perplexity measurement.

Testing

• tests/models/qwen4_exp/test_fp8_dense.py: loader round trip within e4m3 tolerance, the in_proj split per rank (fp8 q|k|v|z rows + bf16 b|a rows equal to the source), op state-dict contract (CPU); the op against F.linear on the dequantized weight at M=1 / 16 / 300 and the zero-input scale floor (GPU).
• tests/models/qwen4_exp + tests/scheduler on the box: 153 passed, 104 skipped (CPU).

🤖 Generated with Claude Code

https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt

[gdevenyi]

3f8f249: the expert-cache planner now sees the halved dense bytes. in_proj_ba.weight came out of t[qkvz:].contiguous(), which is a view of the sharded bf16 in_proj (a contiguous row slice is already contiguous), so each GDN layer kept 42 MB of bf16 alive next to its fp8 copy, 1.5 GiB per TP=2 rank. With a clone: 22,458 -> 23,539 slots (91.4% -> 95.8% residency) (a separate restart with the same flags, not the paired A/B runs in the body). Single-stream re-measured on that build: 99.6 tok/s; vision, long-context and text probes pass. Body updated.