Repository navigation
perf(gemma4/cuda): gemma-4 decode carries ~44 ms/token model-specific fixed overhead on GB10, independent of context length #680
Description
Activity
- added a commit that references this issue
on 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:Metric (per step) gemma-4-12b (48L) llama-8b (32L) ratio MLX primitives evaluated ( gpu::evalranges)~2,383 (49.6/layer) ~943 (29.5/layer) 2.5x CUDA graph kernel nodes ~1,545 ~465 3.3x CUDA graphs launched ( CommandEncoder::commit)~102 ~37 2.8x Wall per step ~75 ms ~22 ms 3.4x 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
- Attention/SDPA: perf(cuda): decode for head_dim > 128 models runs the unfused SDPA fallback; extend the sdpa_vector kernels to cover them #675 A/B was 0.99-1.01x at every context (512-32768).
- Per-step mask construction: the
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. - CUDA graph size caps: MLX pins DGX Spark (cc 1210) to
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. - pipeline_hint:
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.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.
- added a commit that references this issue
on Jul 7, 2026 - added a commit that references this issue
on Jul 7, 2026 - added a commit that references this issue
on Oct 6, 2026
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:
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)
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
Related