Gemma 4 CUDA perf: collapse per-token kernel launches + enable batched prefill

[pekkah]

✅ RESOLVED (2026-06-14)

The headline launch-collapse + batched-prefill work shipped: PLE projection fused (llm_ple_proj_norm), batched-trunk prefill enabled for Gemma 4 (#137), and CUDA-graph decode capture/replay made real and default-on (#139/#142 — SHARPI_CUDA_GRAPH=0 reverts), plus the attnScale and dual-HeadNorm fusions. Prefill/decode then moved to cuBLAS GEMM / dp4a (#141/#142). Decode is now HBM-bandwidth-bound (matvecs ~90% HBM), so the two remaining P1 micro-fusions (layer_output_scale, EmbeddingScale bake) are moot — there is no launch-count lever left. Remaining prefill compute lever lives in #152. Closing.

---

Performance review of the Gemma 4 CUDA execution path (CudaForwardPass.ForwardGemma4, CudaHybridForwardPass.GpuLayerGemma4/Prefill).

Key finding: Gemma 4 decode is launch-bound, not compute-bound. The inner loop issues ~28 kernel launches/layer (~950/token over ~34 layers) plus ~200 upfront for PLE projection prep — ~1,100–1,200 kernel launches per token. At ~4–6 µs host launch cost that's ~5–7 ms/token of pure host overhead (effective ceiling ~150–200 t/s before any compute). PLE alone is ~half the launches. Numerics (SWA bounding, dual-RoPE, KV-share, per-layer head_dim) are already correct — the wins are in launch count and prefill batching.

The recent GDN/MoE/MTP perf work is architecturally irrelevant to Gemma 4 (dense, no SSM, no experts), but two pieces transfer: the non-owning View slice type (#111) and the batched-trunk GEMM-N machinery (#123, currently gated off for Gemma 4).

Todo

P0 — Collapse PLE launches (biggest decode win)

• Replace PLE staging CopyDeviceRegion round-trips with CudaBackend.View(parent, elemOffset, elemCount) (added in perf(engine,cuda): batch the GDN-hybrid prefill trunk (GEMM-N attn/GDN projections) — trunk is now ~62% of prefill (#110 follow-up) #111). The "Tensor can't encode a device-pointer offset" comments at CudaForwardPass.cs:1122-1126 and 1156-1158 are stale. Zero numeric change.
• Replace the L-iteration per-slice loop in BuildPerLayerProjectionsGpu (CudaForwardPass.cs:1128-1140) with a single fused batched kernel (llm_ple_proj_norm: per-pleWidth-slice RmsNorm + add PLE-row slice + scale 1/√2 over all L rows). ~200 launches/token → 1.
• Hybrid path: replace the per-layer H2D UploadInto of the PLE slice (CudaHybridForwardPass.cs:1803-1804) with a View into a once-uploaded _projPerLayer.

P0 — Enable batched-trunk prefill for Gemma 4 (biggest TTFT win)

• Lift the !_isGemma4Like gate in IsBatchedPrefillSupported (CudaHybridForwardPass.cs:975-982). Blockers to resolve:
• Thread per-layer LayerHeadDim[i] through the GEMM-N trunk shapes (decode path already proves per-layer View slicing works).
• Add batched RoPE variants for dual-RoPE: RoPEWithFactors (global) and plain-θ (SWA). Current batched RoPE kernel is NEOX-only.
• Batched PLE over N tokens (reuses the P0 fused projection kernel with a batch dim).

P1 — Fold scalar fix-up kernels into neighbors

• Add explicit attnScale param to Attention/AttentionSwa (pass 1.0f for Gemma 4) and delete the Q prescale ScaleInPlace(qView, √layerHd) (CudaForwardPass.cs:1018, CudaHybridForwardPass.cs:1746). ~34 launches/token gone.
• Fuse layer_output_scale (CudaForwardPass.cs:1064-1065) into the preceding PLE/post-FFN AddInPlace as a fused scale-add.
• Bake EmbeddingScale (√embDim, :936-937) into embedding dequant or first RmsNorm weight at load.

P1 — Make the hybrid command-buffer real (CUDA Graphs)

• Implement BeginRecord/RecordBarrier/EndRecordAndSubmit (currently no-ops at CudaBackend.cs:685-688) as real CUDA Graph capture + replay with updatable kernel-node params for the changing position/token. Structural fix for the ~1,000-launch/token problem.

P2 — Minor fusions

• Fuse dual Q+K HeadNorm (CudaForwardPass.cs:974-978) into a single two-input kernel (MatVecDual-style). ~34 launches/token.

---

🤖 Generated with Claude Code

---

Update (batched prefill landed): all-GPU batched-trunk prefill for Gemma 4 shipped. Bonus primitives added along the way: llm_matvec_q8_0_gemm_n (the on-disk E4B is Q8_0, previously unbatchable) and llm_gelu_tanh_mul_strided (fully-batched PLE without a transpose). All new batched kernels have bit-exact oracle tests; end-to-end Gemma4_E4B_BatchedPrefill_MatchesSequential confirms parity vs the per-token loop. Follow-ups: (a) >4096 context uses the per-token fallback (global layers need AttentionBatchedWave; SWA is window-bounded so fine); (b) SnapKV-active prefill falls back (no Q-capture in the batched trunk); (c) Q8_0 GEMM-N re-reads weights per token — weight-reuse tiling is a perf follow-up.

[pekkah]

Progress update — branch perf/gemma4-cuda-launch-reduction-136 (8 commits)

Both P0 blocks plus the top P1 and the P2 launch-fusions are landed and verified. Every new CUDA kernel has a bit-exact oracle; the end-to-end Gemma4_E4B_BatchedPrefill_MatchesSequential confirms the batched prefill matches the per-token loop on the real Q8_0 E4B GGUF.

Done

Commit	Item	Effect
1b5f022	PLE staging copies → non-owning View slices	numerically inert; removes per-layer copies + L−1 hybrid uploads
0ba387e	PLE per-slice loop → RmsNormBatched+2 ops	~204 → 3 launches/token
de6ce2d	RoPEWithFactorsBatched kernel	batched-prefill primitive (oracle ✓)
7aa24f5	AttentionSwaBatched kernel	batched-prefill primitive (oracle ✓)
56f19d5	llm_matvec_q8_0_gemm_n + llm_gelu_tanh_mul_strided	unblocks Q8_0 model; gather-free batched PLE (oracles ✓)
c8f8a95	All-GPU batched-trunk prefill	collapses per-position prefill launches; end-to-end parity ✓
d70079e	Drop Q-prescale via attn_scale param	−1 launch/layer; Gemma now matches CPU scale=1.0 exactly; non-Gemma bit-identical via -1f sentinel
26d3718	Fuse Q+K HeadNorm (llm_head_norm_qk + batched)	−1 launch/layer (decode + prefill)

New kernels (all bit-exact-oracle tested)

RoPEWithFactorsBatched, AttentionSwaBatched, llm_matvec_q8_0_gemm_n, llm_gelu_tanh_mul_strided, attn_scale param on the 4 attention kernels, llm_head_norm_qk (+ _batched).

Still open (deliberately deferred)

• CUDA graphs (P1) — the structural fix for the ~1,000-launch/token decode ceiling; large, best as its own focused effort. The BeginRecord/RecordBarrier no-op scaffolding is already present.
• layer_output_scale fold / EmbeddingScale bake (P1) — ~1 launch/layer each; only fold into a conditional residual add, marginal vs. the residual-path risk. Left unchecked.

Validation

Per-kernel oracles + Gemma 4 CUDA suite green at every step; 0 warnings (TreatWarningsAsErrors). Full 110-test CUDA regression (non-Gemma attn_scale sentinel + shared kernel-array additions) running now — will confirm here when it completes.

🤖 Generated with Claude Code

[pekkah]

✅ Full CUDA regression green: 111/111 (0 failed, 0 skipped, 1h44m) — the entire CUDA test surface, incl. all non-Gemma model paths (Qwen / GDN / MoE / Coder). Confirms the attn_scale -1f sentinel is bit-identical for non-Gemma callers and the shared CudaBackend kernel-array additions disturbed nothing. Branch perf/gemma4-cuda-launch-reduction-136 (8 commits) is ready to PR.

[pekkah]

Measured impact (README row re-measured)

All-GPU Gemma 4 E4B-it Q8 on RTX 4070 Ti, -g -1 -c 2048, table's ~1K warm methodology (scripts/bench-gemma4-136.ps1, same prompt as bench-allrows-1k.ps1):

	Before #136	After #136	Δ
Prefill t/s (~1K warm)	47.7	109.4	~2.05×
Decode t/s (near-zero-ctx)	44.1	49.1	+11%

• Prefill — clean same-code A/B via SHARPI_BATCHED_PREFILL=0: 53.4 → 109.4 t/s (batched-trunk prefill collapses the per-position attention/FFN/PLE launches into batched GEMM-N + batched SWA/full attention).
• Decode — measured near-zero-ctx (the methodology the README decode column uses): the launch-bound model gains from the fused PLE projection + dropped Q-prescale + fused Q/K HeadNorm.
• Hybrid -g 22 row unchanged — its prefill is still per-token (batched prefill landed only in the all-GPU CudaForwardPass) and decode is CPU-dense-FFN bandwidth-bound, so the fusions don't move it (6.7/6.5 ≈ prior 6.6/6.8, within noise).

README line 54 updated accordingly. Branch perf/gemma4-cuda-launch-reduction-136 is now 9 commits; full CUDA regression green (111/111).

[pekkah]

Merged via #137 (commit 378e407) into master. Review toolkit + gemini findings addressed; the one deferred (pre-existing SWA KV cache OOB at ctx>window) is tracked separately as #138. Remaining open boxes here — CUDA graphs and the two marginal scalar folds — are the next follow-ups.

[pekkah]

CUDA Graph capture/replay for the Gemma 4 decode loop is implemented in PR #139.

Summary: capture the static-topology decode region once on the first decode token and
replay per token (one cuGraphLaunch + per-node position-param updates instead of ~1k
cuLaunchKernel), for both the all-GPU CudaForwardPass and hybrid CudaHybridForwardPass.
Bit-parity proven against direct launches.

Outcome — smaller than this issue anticipated. After #137's launch-fusions, Gemma 4 E4B
decode is largely memory-bandwidth-bound at realistic context, so graphs only reclaim the
launch time not already hidden behind GPU work:

• all-GPU decode +5–6% at ~1K ctx (42→45 t/s); hybrid +1.6%; prefill unchanged.
• −9.7% at short context — the ~168 cuGraphExecKernelNodeSetParams/token overhead exceeds
  the launch-collapse savings when per-token GPU work is small.

Because of the short-context regression the feature ships default-off behind
SHARPI_CUDA_GRAPH, with a direct-launch fallback. #140 tracks the device-side-position
follow-up that would remove that regression and could justify flipping the default to on.

Also in PR #139: CUDA 12.x runtime alignment (cudart64_12/cublas64_12), and a fix for a
pre-existing SnapKV post-eviction decode bug on the dense path (plus disabling SnapKV for
Gemma-4-style models, whose SWA ring caches + per-layer head_dim can't be SnapKV-compacted).