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.

[gdevenyi]

Two follow-up commits pushed onto this branch, both measured end to end on 2 x RTX 6000 Ada (sm_89, TP=2, nvidia/Qwen3.8-Flash-Next-NVFP4, torch 2.11).

perf(fp8): quantize activations with a block-parallel amax

quant_per_tensor walked the whole tensor in a single Triton program, so its cost was linear in the input: 9.8 us for an [8, 10240] activation, 68.6 us at [64, 10240]. Above 16384 elements it now splits into a block-parallel partial amax plus a reduce+cast. Same arithmetic, so a given tensor quantizes bit for bit as before (there is a test for that); the old three-launch torch reduction path above 65536 elements goes away.

n	before	after
20,480	2.59 us	2.22 us
81,920	9.83 us	2.44 us
655,360	68.60 us	3.06 us

The partial count reaches the kernels as a runtime argument, not a constexpr, and pass 1 is grid-strided and capped, so both kernels have exactly one compiled variant. Passing it as a constexpr instead costs a fresh compilation per distinct input length, which in a server means unbounded variant growth and a compile stall mid-generation — that cost ~7 tok/s on the first benchmark run before it was fixed.

feat(layers): optional per-tensor FP8 lm_head (FREETOKEN_FP8_LMHEAD=1)

The vocab-parallel head is the last large bf16 read on the decode path: 0.64 GB per step per rank after the TP=2 vocab split, 0.80 ms of a 9.9 ms step in an nsys trace. ParallelLMHead.forward gains a _logits() seam and Fp8ParallelLMHead overrides only that, so the all_gather above it is untouched. Untied embeddings only.

Its own env flag rather than riding on FREETOKEN_FP8_DENSE, because this one moves the logits directly rather than feeding a norm or a sigmoid, so it should carry its own quality gate.

Measured against this branch as the baseline, production flags, warm single-stream (the benchmark's first run is always cold):

	single-stream	8 concurrent	TTFT 1k	GSM8K-300	8-question probe
baseline	100.95 tok/s	323.3	0.90 s	97.00%	7/8
FREETOKEN_FP8_LMHEAD=1	103.39	324.5	0.87 s	97.00%	7/8

The aggregate barely moves, and that is expected: at 8 concurrent the step is ~24.7 ms rather than ~10 ms because the routed-expert read scales with the batch, so a fixed dense-side saving is a much smaller fraction of it. This is a single-stream lever.

Tests: tests/models/qwen4_exp + tests/layers give 120 passed / 50 skipped on this branch, the same 3 test_qsa_backend.py::test_chunked_prefill_matches_one_shot failures as the branch point (those are what #388 fixes) and no new ones.

One thing I tried on top and am not proposing: the same treatment for the hyper-connection mixers, which look like the obvious next target (LinearReplicated, so every rank re-reads all 1.30 GB of them every step). It benchmarks flat end to end — 100.95 -> 100.56 tok/s. The FP8 GEMMs do save 5.7 us per mix block cold, but the dynamic per-tensor quant hands 4.6 us of that straight back in three extra launches. The quant cost is per call, not per byte, which is why it pays on one 636 MB GEMM and not on 97 GEMMs of 6.9 MB. Worth knowing before someone repeats it; making it land needs the activation quant fused into the norm that already writes the tensor.

🤖 Generated with Claude Code

https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt