Repository navigation
perf(cuda/gemma4): decode is bound by 8-bit MLP weight bytes plus a near-saturated dual CPU/GPU pipeline; qmm kernels are healthy #683
Description
Activity
Step 2 result: uniform 4-bit requantization measured at 1.59-1.62x decode, confirming the byte-count attribution
Built a uniform 4-bit gemma-4-12b checkpoint by dequantizing the 144 8-bit MLP modules (group 64, affine, u32 low-first layout) and requantizing them to 4-bit min-max affine with scales/biases rounded to bf16 before computing q (self-consistent triples; max per-group reconstruction error 0.053). Checkpoint size 10.99 -> 6.74 GB. The requantization script is a standalone numpy tool (raw safetensors parse/write, no GPU); it can be PR'd into scripts/ if wanted.
Measured on GB10 (mlxcel-bench-decode, greedy 64 tokens, chunked prefill defaults), uniform vs the shipped mixed checkpoint:
context mixed decode tok/s uniform decode tok/s speedup peak GB (mixed -> uniform) 512 14.63 23.68 1.62x 11.60 -> 7.45 2048 14.39 23.11 1.61x 13.00 -> 8.98 8192 14.28 22.86 1.60x 14.51 -> 10.92 This lands ABOVE the 1.35x single-lever projection: at ~43 ms/step the non-qmm overhead visible in the step budget shrank from ~19 ms to ~10 ms, i.e. part of the overhead previously serialized against the longer 8-bit qmm stream now overlaps. Prefill throughput is unchanged within noise (950-1045 tok/s). The remaining gap to the original 1.8x bar (25.8 tok/s = 38.8 ms/step) is ~4.5 ms/step, which is the block-scale fusion territory of step 3.
Quality: greedy smoke prompts behave identically between mixed and uniform (both exhibit this checkpoint's pre-existing temp-0 pad-collapse noted in #682's validation; leading tokens are equally sensible, no relative regression). A real quality gate (perplexity or task eval) should be run before adopting uniform 4-bit as the default serving quantization for gemma-4-12b; the 8-bit MLP override in the shipped checkpoint presumably exists for quality headroom.
Conclusion: the checkpoint byte count is confirmed as the dominant lever exactly as sized by the microbench (PR #684). Step 3 (fusion for the last ~4.5 ms) remains optional on top.
- added a commit that references this issue
on Jul 7, 2026 Quality gate outcome: BLOCKED by #686 (gemma-4 forward is broken); speed numbers stand
Steps 2-3 status update. Tools landed: PR #684 (qmm microbench) and PR #685 (scripts/requantize_mlp.py generalized with --match and --target-group-size, plus examples/perplexity.rs).
Built the planned quality/throughput frontier variants: B uniform 4-bit gs64 (6.74 GB), C down_proj-kept-8-bit (8.16 GB), D MLP 4-bit group_size 32 (7.27 GB). Two loader limitations surfaced: mixed bits WITHIN one MLP (variant C) fails to load because the gemma-4 MLP loader infers bits at the global group_size for the whole projection set, and per-module group_size 32 overrides (variant D) are misread as bits=2. Exact errors are recorded in #686; supporting these is a separate loader improvement if the frontier is to include C/D.
The perplexity gate then invalidated itself in the most useful way possible: gemma-4-12b scores ppl ~12,000-29,000 on a corpus where gemma-3-4b scores 15.3 and llama-8b 24.8, flat-bad from 32-token contexts, meaning the gemma-4 dense forward is catastrophically degraded in mlxcel (filed as #686, priority critical, likely the same root as the temp-0 pad-collapse). Every quality conclusion about quantization variants is therefore VOID until #686 is fixed: comparing 4-bit vs 8-bit MLP quality on a broken forward measures noise.
What remains valid: all throughput measurements (uniform 4-bit B = 23.1 tok/s @2048 = 1.61x, peak -3.6 GB) and the step-1 microbench attribution. The answer to "can we keep the 8-bit-MLP quality condition AND regain speed" is deferred: once #686 is fixed, re-run the frontier with the landed tools (B is built; C/D need the loader improvement or can be approximated by requantizing at the whole-MLP granularity).
Recommended order: fix #686 first (it may also change the perf baseline), then re-run this quality gate, then decide between uniform 4-bit, selective 8-bit, or gs32 on real quality data, and only then invest in step-3 fusion.
Summary
Follow-up investigation to #680 (closed by PR #682 with corrected attribution). A per-shape quantized-matmul GEMV microbench on GB10 (new tool:
examples/qmm_gemv_microbench.rs, results CSVbenchmarks/cuda_gb10_qmm_gemv_micro_2026-07-07.csv, both in an accompanying PR) shows the qmm GEMV kernels are HEALTHY, sustaining 201 to 214 GB/s (74 to 78% of the 273 GB/s peak) on every gemma/llama layer shape at both 4-bit and 8-bit, with no shape cliff (smaller attention shapes reach 170 to 190 GB/s from fixed launch latency, same on llama). The earlier "gemma reaches only ~52% effective bandwidth" figure was decode wall time diluted by non-qmm work, not kernel inefficiency. Kernel-side optimization of qmm GEMV is NOT a lever; do not pursue it.Measured decode decomposition (GB10, greedy, 2048-token context)
Using measured per-shape qmm times (microbench) and CUDA GRAPH_TRACE gap analysis of the #680 nsys captures:
The model closes on all three (predicted vs measured within a few percent). gemma-4-12b's per-step GPU gaps are ~1.6 large gaps (>0.5 ms) totaling ~3.7 ms at step boundaries (sampling plus next-step graph construction before the first launch), which llama does not have at all; the CPU lazy-graph pipeline (~2,400 primitives/step, see #680) therefore runs within ~10% of the GPU and becomes the binding constraint as soon as GPU time is reduced.
Why gemma-4-12b is slow: the checkpoint streams 2.3x llama's bytes
The "4bit" gemma-4-12b-it checkpoint is mixed precision: 144 modules (48 layers x gate/up/down) are 8-bit, which is ~9 GB of the ~10.3 GB streamed per token (MLP dominates parameter count). llama-3.1-8b-4bit streams ~4.4 GB. At healthy kernel bandwidth the byte ratio alone explains most of the 14.3 vs 47 tok/s gap; the rest is the ~19 ms/step non-qmm overhead (small-kernel GPU time + gaps) tied to gemma-4's ~50 primitives/layer (#680).
Levers, sized with measured numbers (gemma-4-12b @2048)
Suggested plan
Related