Repository navigation
ggml-cuda : mxfp4 quantization, KV cache and Blackwell MMA - #14
timothyeburke wants to merge 5 commits into
Conversation
4263123 to
622b320
Compare
e243eb5 to
7d4e822
Compare
|
Hi, thanks for sharing the scaling experiments! Quick question about the earlier UOS result mentioned in the scale-selection section. For MXFP4 E2M1, MXAttention defines UOS using
The current implementation here appears to parameterize the E8M0 scale using round-to-nearest on How was the UOS comparison implemented? Was the E2M1 weight scale computed explicitly with Also, since the current configuration is W4A8, what E4M3 activation scaling was paired with UOS in that comparison? Thanks! |
|
@oupenai I've been polishing this PR here before moving to the main repo. I ended up backing out the UOS scale changes because I got sub-optimal results with them. I derived the e4m3 UOS value at 464. It's still visible in https://github.com/timlikesai/llama.cpp/tree/mxfp68-quant |
- new dense LLAMA_FTYPE_MOSTLY_MXFP4; token embeddings and output stay Q8_0 until mxfp8 lands - MXFP4 KV cache: f32<->mxfp4 cpy/set-rows, f16/bf16/f32 convert, mxfp4-mxfp4 FA vec instance - Blackwell MMQ: activations quantized to e4m3, block-scaled mxf8f6f4 mma (W4A8) instead of mxf4 (W4A4): slightly slower, much better accuracy. Y-side tile load loop generalized, codegen for existing types unchanged - imatrix weight quantization: per-block e8m0 exponent search minimizing weighted error - e2m1 rounding is RNE matching the intrinsics; e8m0 scale rounds the log2 instead of truncating Assisted-by: pi and qwen/qwen3.8-27b
Use the MXAttention Universal Optimal Scaling boundary (e_base = 7.25, arXiv 2607.24377) for mxfp4 KV cache and W4A4 activations. Lower KLD on all tested models; weight path unchanged. Assisted-by: pi and qwen/qwen3.8-27b
Extract ggml_e8m0_scale and ggml_float_to_fp4_e2m1_index to ggml-impl.h and share them across the CPU/CUDA paths (the CUDA fp4 conversion is a wrapper around the shared one). Parameterize the mxfp4 activation quantize kernel by prec instead of duplicating it for W4A4/W4A8, generalize the y-side tile load loop, and guard the Blackwell-only vec-dot kernel so it compiles as a no-op elsewhere. Removes the orphaned compute_e8m0_scale. Assisted-by: pi and qwen/qwen3.8-27b
The e4m3 activation pack uses fp8x4 conversion intrinsics that do not exist or are host-only on HIP, and the MUSA compiler does not define __CUDACC__, so GGML_HOST_DEVICE left the shared fp4 helper host-only. The fp8 path is only reachable where the Blackwell block-scaled mma is available; guard it, and treat GGML_USE_MUSA as device compilation for the helper macro. Assisted-by: pi and qwen/qwen3.8-27b
Overview
End-to-end MXFP4 for CUDA Blackwell, narrowed from ggml-org#20609: dense MXFP4 ftype, W4A8 block-scaled mma, imatrix-driven weight quantization, and an MXFP4 KV cache. Small and mighty at +840/-205 - mostly plumbing to complete the type across ggml quant/dequant, MMQ kernels, FA KV cache, and the llama ftype, plus tests.
LLAMA_FTYPE_MOSTLY_MXFP4=42;--cache-type-k/--cache-type-v mxfp4KV read directly by the FA vec kernelmxf8f6f4mma instead of W4A4: ~2.4x lower KLD at unchanged decode. The same instruction and mostly the same code back the mxfp8/mxfp6 (W8A8/W6A8) weight follow-ups; the standard Q8_1-activation MMQ path remains selectable (GGML_CUDA_MMQ_PREC=q8)--imatrixweight path picks the optimal per-block weight scale using the importance matrixGGML_CUDA_MMQ_PREC=q4). The measured KV-cache quantization effect is lower (KLD) on 0.8B, within run noise on 27B/35B; the weight path is unchangedThe mxfp4 imatrix quantizations are on HuggingFace: 0.8B, 27B, 35B-A3B - each repo includes the imatrix file used, so the quants are reproducible.
Results
Tested using 2x 5060 Ti 16GB throttled to 150W/180W due to a slightly defective power supply.
W4A8 (this PR) vs W4A4 vs Q8 activations
Details
Same mxfp4 files, 2x 5060 Ti,
--n-gpu-layers 999 --split-mode tensor --flash-attn on. Four MMQ activation arms: W4A4 with the old OCP scale (masterfb27a525d, which loads these files with anunknown type mxfp4warning and routes to its existing W4A4 mma), W4A4 with the UOS e2m1 activation scale (this branch,GGML_CUDA_MMQ_PREC=q4), the shipped W4A8mxf8f6f4path (e4m3 activations, default), and the standard Q8_1-activation MMQ path (GGML_CUDA_MMQ_PREC=q8). W4A8 cuts W4A4's KLD by ~2.4-2.7x; decode is unchanged across arms.W4A arms, imx (imatrix) files, 72-chunk KLD, f16 KV, vs the dumped BF16 base:
Same for the plain (no-imatrix) files:
UOS scale variants, same 72-chunk KLD (W4A4 activation scale old OCP vs UOS 7.25; W4A8 e4m3 boundary 256 shipped vs 343/464 UOS candidates), imx files:
UOS helps the coarse e2m1 grid (W4A4 activations: 8-13% KLD on 0.8B/35B, neutral on 27B) like it does the KV cache, but W4A8 stays ~2.4x better; the fine e4m3 grid is flat across 256/343/464 (all within run noise). The Q8_1-activation path measures slightly better than W4A8 on both accuracy and prefill on this hardware.
Throughput (imx files, pp4096 / tg128, -r 5):
Accuracy vs file size
Details
Qwen3.8-27B (dense):
Qwen3.5-0.8B (dense):
Qwen3.6-35B-A3B (MoE):
KL divergence + same top-p vs BF16
Details
Qwen3.8-27B (dense):
Qwen3.5-0.8B (dense):
Qwen3.6-35B-A3B (MoE):
KV cache (
--cache-type-k/--cache-type-v mxfp4)Details
KV cache memory (GiB) at 100k tokens, and throughput by KV type: GPU (2x RTX 5060 Ti, 150W, Q4_1-imx weights, pp4096/tg128, --flash-attn 1) and CPU (9900X, pp512/tg32, 24 threads, --n-gpu-layers 0, MXFP4-imx weights), -r 5:
72-chunk PPL by KV type (mxfp4-imx weights, GPU): 0.8B f16 16.1791 / q5_1 16.2197 / q4_1 16.3278 / mxfp4 16.3649 / q4_0 16.3715; 27B f16 6.3552 / q5_1 6.3578 / q4_1 6.3689 / q4_0 6.3749 / mxfp4 6.4051; 35B f16 5.9302 / q5_1 5.9330 / q4_1 5.9415 / q4_0 5.9442 / mxfp4 5.9488.
KV cache scale
UOS (Universal Optimal Scaling) lowers the mxfp4 KV-cache quantization error on 0.8B; on 27B/35B the two arms are within run noise.
Details
mxfp4-imx weights, 2x RTX 5060 Ti; KV-cache effect = arm - control (f16 KV). KLD / PPL(Q) / top-p %:
The 27B KV effect itself is only +0.0025 to +0.0034 KLD over the f16 control - below the per-run std (+/-0.0019) - so neither scale formula is distinguishable there. UOS is kept as the default on the strength of the 0.8B result (and the W4A4-activation result, where UOS is also never worse).
Additional information
Why W4A8, scale derivation, and controls
Why W4A8
The previous native MXFP4 path quantized activations to e2m1 (W4A4). W4A8 uses e4m3 activations via the native mixed-precision
kind::mxf8f6f4mma - the sm_120-supported form (see ggml-org#19662) - cutting W4A4's KLD by ~2.4x at unchanged decode; the same instruction and mostly the same code back the mxfp8/mxfp6 (W8A8/W6A8) weight follow-ups. The standard Q8_1-activation MMQ path (GGML_CUDA_MMQ_PREC=q8) stays selectable and measures slightly better than W4A8 on both accuracy and prefill on this hardware (Q8_1 activations carry more precision than e4m3); W4A8 remains the default as the format-native instruction path. Prior art: the closed ggml-org#27315 kept e2m1 activations; its data showed W4A4 at 84.24% same-top-P (KLD 0.1316) vs 89.99% for W4A16 without MMQ.Scale selection
The e8m0 block scale is the standard
round_to_pow2(amax / C)with C stepped in from the format's max, so the block's largest value maps inside the representable range instead of onto its edge:GGML_CUDA_MMQ_PREC=q4): UOS 7.25 - the same MXAttention boundary for the e2m1 activation grid. Improves KLD/PPL vs the OCP 4.0 pin (0.8B KLD 0.408 -> 0.355, PPL 21.39 -> 20.06; 35B KLD 0.169 -> 0.150, PPL 6.47 -> 6.34; 27B within noise) - see the W4A8 sectionA single-pass "pick the scale that minimizes output error" (the natural per-tensor optimum) measured worse than the fixed stepped-in scales on an early sweep; the e8m0 scale is a power of two, so the per-tensor optimum does not generalize. The fixed stepped-in scales are the robust choice.
Optional imatrix weight quantization. With calibration data,
llama-quantize --imatrixsearches a band of weight scales around /4.0 and picks, per 32-wide block, the one that minimizes the imatrix-weighted output error. Opt-in; the default is the calibration-free /4.0.Existing GGUFs are unaffected; dequantization is unchanged.
Controls
Before this PR, MXFP4 was only available for MoE expert weights (MXFP4_MOE). Dense MXFP4 is new here. All mxfp4 quants are quantized from the ggml-org BF16 bases (Qwen3.8-27B-GGUF, Qwen3.6-35B-A3B-GGUF, Qwen3.5-0.8B-GGUF) with the measured-optimal scales (weight /4.0, activation /256); the mxfp4 PPL above is the imatrix-weighted variant (27B: 6.355 imx vs 6.449 plain). Results are from this branch (0c666e5); the W4A4-old arm is current master (fb27a52).
How measured
llama-bench--n-gpu-layers 999 --split-mode tensor --flash-attn on --n-prompt 4096 --n-predict 128 --n-repeat 5llama-bench--n-gpu-layers 0 --threads 24 --n-prompt 512 --n-predict 32 --n-repeat 5llama-perplexity--file wikitext-2 --n-gpu-layers 999 --split-mode tensor --flash-attn 1 --context-size 4096 --batch-size 512(72 chunks)llama-perplexity --kl-divergence--file wikitext-2 --context-size 4096against a dumped bf16 base (full corpus)mxf8f6f4;GGML_CUDA_MMQ_PREC=q4= W4A4 (UOS e2m1 scale);GGML_CUDA_MMQ_PREC=q8= standard Q8_1 activations; W4A4-old = master buildllama-quantize --imatrixllama-quantize/llama-benchHardware: 2x RTX 5060 Ti (throttled to 150W); CPU: AMD Ryzen 9 9900X (24 threads). All models are ggml-org; KLD and PPL are hardware-independent.
mxf8f6f4plumbing; the CPU-side MXFP8/MXFP6/E4M3 quantizer work (ggml: add support for MXFP8 CPU ggml-org/llama.cpp#26157, ggml: Add initial MXFP6 CPU implementation ggml-org/llama.cpp#22671, ggml : add E4M3 (fp8) CPU quantization type ggml-org/llama.cpp#25336) is the other half of that pictureggml_prec(extended here withGGML_PREC_MXFP8), CUDA compilation fails on Blackwell (sm_120) with MXFP4: "Instruction 'mma with block scale' not supported" ggml-org/llama.cpp#19662 tracks the sm_120 block-scale mma build issue, CUDA: Add experimental SM120 CUTLASS MoE prefill for MXFP4 and NVFP4 ggml-org/llama.cpp#26704 is adjacent SM120 CUTLASS MoE prefill for MXFP4/NVFP4, [CUDA][Blackwell] Qwen3.8-27B-NVFP4 decode hangs, CPU spin, no GPU work ggml-org/llama.cpp#27329 reports an NVFP4 decode hang on sm_100, and Eval bug: Muse Glimmer NVFP4 fails to load, then emits one repeated token ggml-org/llama.cpp#27178/Eval bug: Windows HIP gfx1201: Qwen3.6-35B-A3B NVFP4 loads, then first MUL_MAT fails with ROCm invalid argument ggml-org/llama.cpp#27670 report NVFP4 load/compute failures on CUDA/HIPRequirements