perf(gemma4-cuda): collapse per-token launches + all-GPU batched-trunk prefill (#136)

[pekkah]

Gemma 4 CUDA performance: collapse per-token launches + batched-trunk prefill

Closes the work tracked in #136. Gemma 4 decode on CUDA is launch-bound (~1,000+ kernel launches/token), and prefill ran a per-token loop. This PR collapses launches on the decode path and adds an all-GPU batched-trunk prefill, plus two launch fusions.

Measured impact (E4B-it Q8, RTX 4070 Ti, -g -1 -c 2048, ~1K warm)

	Before	After	Δ
Prefill t/s	47.7	109.4	~2.05× (SHARPI_BATCHED_PREFILL=0 A/B: 53.4 → 109.4)
Decode t/s (near-zero-ctx)	44.1	49.1	+11%

What's in it (8 perf commits + 1 docs)

P0 — decode launch collapse

• PLE staging copies → non-owning View slices (the "Tensor can't encode a device offset" comment was stale; View from perf(engine,cuda): batch the GDN-hybrid prefill trunk (GEMM-N attn/GDN projections) — trunk is now ~62% of prefill (#110 follow-up) #111 does exactly that).
• PLE per-slice projection loop → RmsNormBatched + 2 full-buffer ops (~204 → 3 launches/token).

P0 — all-GPU batched-trunk prefill (CudaForwardPass)

• Embeds N tokens, builds PLE projections batched, runs every layer batched across N, final norm/output on the last token. Per-layer head_dim via max-sized scratch + exact-N views.
• Four new batched kernels, each with a bit-exact oracle: RoPEWithFactorsBatched, AttentionSwaBatched, llm_matvec_q8_0_gemm_n (unblocks the Q8_0 model), llm_gelu_tanh_mul_strided (gather-free batched PLE).
• Gated by IsGemma4BatchedPrefillSupported (dense, NEOX, non-L2 QK-norm, no attn-bias, TQ off, batchable dtypes, ctx ≤ 4096); SnapKV-active / unsupported configs fall back to the bit-exact per-token loop. SHARPI_BATCHED_PREFILL=0 forces sequential.

P1/P2 — launch fusions

• Drop the per-layer Q-prescale via an attn_scale kernel param (sentinel -1f keeps every non-Gemma caller bit-identical; Gemma passes 1.0 and now matches the CPU scale=1.0 reference exactly).
• Fuse Q+K HeadNorm into one launch (llm_head_norm_qk + batched).

Verification

• Every new kernel has a *_BitwiseMatchesSingle/Sequential/Separate oracle.
• End-to-end Gemma4_E4B_BatchedPrefill_MatchesSequential confirms batched ≡ per-token post-prompt logits on the real Q8_0 GGUF.
• Full CUDA regression green: 111/111 (all non-Gemma model paths — Qwen/GDN/MoE/Coder — confirm the attn_scale sentinel + shared kernel-array additions are inert). 0 warnings (TreatWarningsAsErrors).

Deferred (tracked as open boxes on #136)

• CUDA graphs (the structural fix for the launch-bound ceiling; large, its own effort — no-op BeginRecord/RecordBarrier scaffolding already present).
• layer_output_scale fold / EmbeddingScale bake (~1 launch/layer each; marginal vs. residual-path risk).

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request implements an all-GPU batched-trunk prefill optimization for Gemma 4 (Issue #136), which collapses per-position attention, FFN, and PLE launches into batched GEMM-N and batched-attention launches to significantly improve prefill performance. It introduces several new batched CUDA kernels (such as dual Q+K HeadNorm, batched SWA, and strided GELU-tanh multiplication) and integrates them into the forward pass logic, backed by comprehensive unit tests to ensure bit-identical outputs. The review feedback suggests improving the benchmark script by parameterizing the hardcoded model directory for better portability and wrapping the environment variable modifications in a try-finally block to ensure robust cleanup.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

The model path is hardcoded, which makes the script less portable. Consider making the model directory a script parameter with a default value. This would allow running the benchmark against models in different locations without modifying the script. It's a good practice to define parameters at the beginning of the script.

param(
    [string]$ModelDir = "E:\models"
)
$model = Join-Path $ModelDir "gemma-4-E4B-it-Q8_0.gguf"

[gemini-code-assist]

The SHARPI_BATCHED_PREFILL environment variable is modified but not robustly cleaned up. If an error occurs during the benchmark run, the variable might persist in the shell session. Using a try...finally block ensures that the environment variable is always cleaned up, regardless of whether the script runs successfully or encounters an error.

function Run($tag, $a, $envOff) {
    try {
        if ($envOff) { $env:SHARPI_BATCHED_PREFILL = "0" } else { Remove-Item env:SHARPI_BATCHED_PREFILL -ErrorAction SilentlyContinue }
        # Run twice, keep the second (warm) — matches the discard-first guidance.
        $null = .\scripts\bench-textgen.ps1 -Model $model -Tag "$tag-w" -NTokens 60 -Prompt $prompt -TimeoutSec 900 -ExtraArgs $a
        $r = .\scripts\bench-textgen.ps1 -Model $model -Tag $tag -NTokens 60 -Prompt $prompt -TimeoutSec 900 -ExtraArgs $a
        $script:rows += [PSCustomObject]@{ Tag=$tag; PrefTok=$r.PrefillTok; PrefillTps=$r.PrefillTps; DecodeTps=$r.DecodeTps; Wall=$r.WallSec; TO=$r.TimedOut }
        Write-Host ("  {0,-22} pref={1,7} t/s  dec={2,6} t/s  ({3} tok)" -f $tag,$r.PrefillTps,$r.DecodeTps,$r.PrefillTok) -ForegroundColor Green
    }
    finally {
        Remove-Item env:SHARPI_BATCHED_PREFILL -ErrorAction SilentlyContinue
    }
}

[pekkah]

Review addressed (toolkit + gemini)

Ran the pr-review-toolkit (code-reviewer, silent-failure-hunter, pr-test-analyzer, comment-analyzer) plus the gemini-code-assist comments. Fixes pushed in 5b1a859 (toolkit) and 9ae0329 (gemini). 28/28 Gemma 4 CUDA + batched-kernel oracles green after the changes.

Code / safety

• Fail-safe prefill gate — new BatchableWeight() treats an unregistered weight handle as not-batchable (was WDType → Q4_K default, which biased the gate toward "supported"); also gate on _gpuPerLayerModelProj's dtype. An untracked weight now forces the per-token fallback instead of being misdispatched into GEMM-N.

Tests (the top finding — the central change was under-covered)

• Added AttentionExplicitScale_BitwiseMatchesSingle_AndDiffersFromSentinel: exercises the attn_scale=1f path (Gemma's production path) bit-exact for both batched SWA and full attention, and asserts 1f ≠ the rsqrt sentinel so the param is provably plumbed. Previously only end-to-end coverage hit it.
• Tightened the end-to-end parity test to assert a >95% bit-exact fraction (the exact count was computed but unused), not just the loose maxAbs<1e-2 envelope.

Comments / docs

• <param name="attnScale"> added to all 4 attention dispatches (the ≤0 → 1/sqrt(headDim) sentinel contract).
• Fixed the PLE-collapse note (≈6 ops×L, not 3×L) and the batched-scratch note (Q/K/V/attn buffers are max-head_dim-sized + viewed per layer).

Bench script (gemini)

• $ModelDir param instead of a hardcoded path; try/finally around the SHARPI_BATCHED_PREFILL toggle.

Deferred to a separate issue (#138), not a blocker

• silent-failure-hunter flagged the Gemma SWA KV cache (window-sized buffer + absolute indexing → potential OOB when ctx > sliding_window). It's pre-existing and identical in the per-token and batched paths (so parity holds), and the benched/tested configs don't hit it. Filed as Gemma 4 SWA KV cache: window-sized buffer vs absolute indexing → OOB when ctx > sliding_window #138 for a proper ring-buffer-vs-full-length fix; out of scope for this parity-preserving perf PR.