Repository navigation
perf(gemma4/cuda): gemma-4 decode carries ~44 ms/token model-specific fixed overhead on GB10, independent of context length #680
Description
Summary
On the CUDA backend (GB10), gemma-4-12b-it-4bit decodes at 13.7 to 14.6 tok/s (68 to 70 ms/token) essentially FLAT across prompt contexts 512, 2048, 8192, and 32768. The 4-bit weight-streaming floor is ~24.5 ms/token (~6.7 GB at ~273 GB/s), so ~44 ms/token is fixed per-step overhead that does not scale with k_len. Comparable models do not show it, so it is gemma-4 model-implementation-specific, not an mlxcel-wide or MLX-wide cost.
Evidence (measured 2026-07-06/07, GB10, release, mlxcel-bench-decode, 64 generated tokens)
Cross-model comparison at 8192-token context:
| Model | Decode tok/s | ms/token | Weight floor | Overhead | Layers | Overhead/layer |
|---|---|---|---|---|---|---|
| llama-3.1-8b-4bit | ~47.1 | 21.2 | ~16.5 ms | ~5 ms | 32 | ~0.15 ms |
| qwen3.5-9b-4bit (full attention every layer) | 38.55 | 25.9 | ~18 ms | ~8 ms | ~0.2 ms | |
| gemma-4-12b-it-4bit | 14.28 | 70.0 | ~24.5 ms | ~44 ms | 48 | ~0.9 ms |
Context-independence: gemma-4-12b decode is 14.62 tok/s at 512-token context and 13.78 tok/s at 32768 (68.4 to 72.6 ms/token, a delta of only a few ms for 64x more keys), so the overhead is not attention-data-volume-bound.
Attention is ruled out as the cause: the #675 sdpa_vector A/B (fused vector kernel vs the unfused materializing fallback, toggled via MLXCEL_SDPA_VECTOR_LARGE_D with no rebuild) measured 0.993x / 1.001x / 0.994x / 1.007x at 512 / 2048 / 8192 / 32768 on gemma-4-12b, i.e. the SDPA implementation choice is invisible at decode. Full data in benchmarks/cuda_gb10_sdpav_675_2026-07-06.csv.
Why it matters
This overhead, not the unfused SDPA fallback, explains the gemma-4 vs llama decode gap that motivated #675 (gemma-4-31b at 5.75 to 7.3 tok/s, 12b at ~14 tok/s, vs llama-3.1-8b at 43 to 54 tok/s). Eliminating it takes gemma-4-12b from ~70 ms/token toward ~30 ms/token, roughly 2.3x decode throughput, which is the 2x target #675 originally aimed for. The whole gemma-4 family (dense, MoE, unified) plus gemma-3/3n likely share the cost.
Suspects to profile (in the mlxcel gemma-4 implementation, src/models/gemma4.rs and shared layers)
- Per-step attention-mask construction: the iota mask builders and flag-path masks added for fix: gemma-4 long-prompt prefill OOMs on CUDA because attention materializes O(heads*L^2) scores instead of using flash SDPA #672/fix(cuda/attention): bound long-prompt prefill memory for flash-ineligible models #676 prefill, and mask_stale_key_gap, may be running (or allocating) per decode step per layer even when the decoder mask should be None for batch-1.
- Rotating KV cache per-step operations on the 40 sliding-window layers (5 local : 1 global pattern): slice/rotate/copy work per step per layer.
- The Gemma 4 Unified multimodal wrapper on the text decode path: the audio embedder is loaded even for text-only runs ("Loaded Gemma4 Unified audio embedder" in bench stderr), so wrapper-level per-step work may be in the loop.
- MTP target module (gemma4_mtp_target) touching the per-step graph even when speculative decoding is off.
- Small-kernel launch count per block: gemma-4 sandwich norms (pre/post attention and MLP) plus QK-norm roughly double the per-layer op count vs llama; on GB10 the per-launch/graph-build cost may add up across 48 layers.
Fix direction
Profile ONE decode step first (nsys, MLX op-level timers, or bisecting by stubbing suspects) and attribute the 44 ms before changing anything. Then eliminate the dominant contributor(s). Re-run the decode A/B at 512/2048/8192 afterward.
Acceptance criteria
- A per-step profile or bisection attributing the ~44 ms/token overhead is documented in the issue or PR
- The dominant contributor is fixed and decode tok/s before/after lands in a benchmarks CSV
- gemma-4-12b-it-4bit decode reaches at least 1.8x its current tok/s at 2048-token context on GB10
- Token-identity (greedy, chat template) is unchanged by the fix
- No prefill throughput or peak-memory regression (fix: gemma-4 long-prompt prefill OOMs on CUDA because attention materializes O(heads*L^2) scores instead of using flash SDPA #672/fix(cuda/attention): bound long-prompt prefill memory for flash-ineligible models #676 behavior preserved)
Related
- Discovered during perf(cuda): decode for head_dim > 128 models runs the unfused SDPA fallback; extend the sdpa_vector kernels to cover them #675 validation (sdpa_vector head_dim 256/288 extension, which is correct but cannot deliver the 2x on gemma-4 because of this overhead)
- fix: gemma-4 long-prompt prefill OOMs on CUDA because attention materializes O(heads*L^2) scores instead of using flash SDPA #672 / fix(cuda/attention): bound long-prompt prefill memory for flash-ineligible models #676: prefill-side fixes for the same head_dim > 128 model family
- fix(cache): unbreak rotating-cache chunked prefill and default the CLI prefill to chunked #679: chunked prefill default
Activity
perf(cuda): extend sdpa_vector decode kernels to head_dim 256/288 via…
inureyes commentedon Jul 7, 2026
Attribution complete: the overhead is CPU-side per-step graph construction, proportional to MLX primitive count
nsys profile (GB10, release, 128-token greedy decode,
--trace=cuda,nvtx), gemma-4-12b-it-4bit vs llama-3.1-8b-4bit control, normalized per decode step:gpu::evalranges)CommandEncoder::commit)Both models cost a near-constant ~25-31 us per MLX primitive end to end on the GB10 CPU (graph build + dispatch + CUDA API). Decode wall time is
primitives_per_step x ~28usfor both models; the GPU is idle waiting (total GPU kernel time in the trace is negligible; the only 10 ms/step GPU-side blocking wait is the finalcudaEventSynchronizeper token). gemma-4 is 3.4x slower than llama simply because it evaluates 2.5x the primitives through the same CPU-bound pipeline.Per-primitive breakdown (per step, gemma vs llama): Multiply 394 vs <30, RMSNorm 337 vs 66 (7/layer: 4 sandwich norms + q/k/v norms, all already fused to one
fast_rms_normeach), QuantizedMatmul 329 vs 165, Add 194 vs 65, SliceUpdate 95 vs 66, RoPE 96 vs 66, SDPA 39 vs 33, plain Matmul 17 vs ~0 (the 8 full-attention layers take a materializing single-query path plus lm_head), plus ~850/step of unlisted view/elementwise primitives (reshape/transpose/astype/etc).Ruled out by measurement
l == 1decode branch leaves both masksNone(gemma4.rs), and gemma-4-12b hashidden_size_per_layer_input: 0so the per-layer-input machinery is off.ops=20 / mb=25per graph (upstream [CUDA] Tune ops per buffer based on device ml-explore/mlx#2761), which sounded like the smoking gun, but raising the caps viaMLX_MAX_OPS_PER_BUFFER/MLX_MAX_MB_PER_BUFFERmakes decode STRICTLY WORSE, measured at 2048 ctx: gemma 14.33 -> 13.49 (100/1000) -> 12.83 (1000/10000) tok/s; llama 46.37 -> 39.71 (100/1000). Upstream's Spark tuning is correct; do not touch it.MLXCEL_PIPELINE_GRANULARITYdefaults off, no eval boundaries inserted.Fix direction (revised)
Reduce the gemma-4 decode path's primitives per layer from ~50 toward llama's ~30, and/or collapse elementwise chains into compiled subgraphs (1 primitive each). Candidate levers, to be sized by a per-layer subop profile (
MLXCEL_PROFILE_LAYER_SUBOPS=1) and code enumeration: fold the per-layerlayer_scalarmultiplies into adjacent pre-computed norm weights at load time (same trick as GemmaRMSNorm's(1+w)); eliminate redundant reshape/astype/transpose primitives on the per-step path (the ~850/step unlisted bucket); compile the residual+norm+scale elementwise chains (mlx::core::compilevia the existingcompiled_softcap_sdpaFFI precedent); route the 8 full-attention layers' single-query decode to fused SDPA now that head_dim 256 is supported (#681, removes the materializing matmul+softmax composite ops). Napkin math: llama-parity op count projects ~75 -> ~45-50 ms/step (~1.5-1.7x); the 1.8x acceptance bar likely needs the compiled-chain lever on top of op elimination.inureyes commentedon Jul 7, 2026
Re-scope: decode is GPU-weight-bandwidth-bound, not CPU-primitive-bound (fix + PR #682)
The attribution above summed only
CUPTI_ACTIVITY_KIND_KERNELfor GPU time (8.66 ms across the whole gemma trace, 1.28 ms for the llama control) and concluded the GPU was idle and decode was CPU-bound on the per-step primitive count. That table only holds kernels launched OUTSIDE a CUDA graph; on GB10 MLX captures each forward into CUDA graphs, so the real GPU work lives inCUPTI_ACTIVITY_KIND_GRAPH_TRACE, which was not summed. Re-measured over a steady-state 20-decode-step window from the same sqlite exports: gemma-4-12b is 93.5% GPU-busy (1329 ms GPU graph execution in a 1422 ms wall, 71.1 ms/step); the llama-3.1-8b control's GRAPH_TRACE is 2333 ms at ~21.4 ms/step, also GPU-bound. Both models are GPU-execution-bound; the error was systematic.gemma-4-12b-it-4bit is a mixed-precision checkpoint: attention is 4-bit but every MLP gate/up/down is group_size=64 bits=8. Per-token active weight is ~10.1 GB, a ~37 ms/token floor at 273 GB/s (measured 71 ms is ~140 GB/s effective). The "44 ms context-independent overhead" is this weight streaming: bytes do not scale with k_len, so decode is flat from 512 to 32768.
Consequence: the biggest primitive-count lever (compiling the 8-bit MLP GeGLU) does cut ~530 primitives/step at decode (nsys-verified Multiply 394->59), but decode throughput moves under 1% because the GPU is already ~94% busy. The 1.8x bar is not reachable by reducing the primitive count; it needs fewer bytes/token (requantizing the MLP to 4-bit, which changes the checkpoint) or a faster MLX 8-bit qmm GEMV kernel (MLX overlay), both out of scope.
PR #682 lands the compiled-MLP fusion gated to decode (non-regressing on prefill/peak, verified 512/2048/8192 plus a llama control and a gemma-4-31b family point) as a correct latent fix, and documents the above. Bench data: benchmarks/cuda_gb10_gemma_decode_680_2026-07-07.csv.
test(mlxcel-core): cover compiled qgelu MLP decode gate and paths (#680)
perf(gemma4): cut decode-path primitive count on CUDA (#680) (#682)
fix(gemma4): make MTP verify byte-identical to decode on CUDA
perf(cuda): extend sdpa_vector decode kernels to head_dim 256/288 via…
perf(gemma4): cut decode-path primitive count on CUDA (lablup#680) (l…