perf(cuda): generalize Gemma 4 batched prefill + flash attn + decode graphs to dense Q4_K models

[pekkah]

Goal

Bring the Gemma 4 CUDA performance work to other dense models that fully offload via CudaForwardPass — primarily Qwen3-8B Q4_K_M. The batched-prefill orchestration and CUDA-graph machinery are already largely generic; they are just structurally gated behind _isGemma4Like. This issue generalizes them.

Out of scope: MoE/GDN models (separate CudaHybridForwardPass / CudaHybridGdnForwardPass), and Vulkan.

Background

What the Gemma 4 work (#136/#141/#142/#146/#147/#149/#154) splits into:

• Already general — Q4_K decode already uses dp4a (the Gemma 4 GPU decode is ~1.8-2.9x slower than llama.cpp (kernel efficiency + low-ctx clock cliff) #142 work brought Q8_0 up to Q4_K); the perf(sampler): top-k-first fast path — Gemma 4 E4B decode 48.7→70.3 t/s (#142) #154 sampler fix is model-agnostic and already benefits everyone.
• Q8_0-specific — int8 MMQ prefill (Gemma 4 GPU prefill is ~70x slower than llama.cpp — adopt cuBLAS GEMM for prefill #141) + SoA repack (perf(cuda): split-K + better-tiled int8 MMQ GEMM — 23-34% → 60-80% of TC peak (#141 follow-up) #149) are tied to the Q8_0 block layout.
• Quant-agnostic but Gemma-gated — batched-trunk prefill (Gemma 4 CUDA perf: collapse per-token kernel launches + enable batched prefill #136), flash-attention prefill (Gemma 4 GPU prefill is ~70x slower than llama.cpp — adopt cuBLAS GEMM for prefill #141/perf(cuda): full tensor-core (mma.sync) flash-attention prefill for d=512 — beat the half2 kernel (#141 follow-up) #146/perf(cuda): multi-warp / d-split TC flash prefill — fix the single-warp occupancy limit (#146 follow-up) #147), decode CUDA graphs (Gemma 4 CUDA perf: collapse per-token kernel launches + enable batched prefill #136). This is the opportunity.

Two config knobs that looked like blockers are already first-class: attnScale (kernels do <=0 => 1/sqrt(headDim), Gemma hardcodes 1f) and hp.FfnActivation (Silu vs GeluApprox, batched body hardcodes GeluTanhMul).

Work Item A — Generalize batched prefill + flash attention (biggest win, low risk)

Three hard Gemma assumptions in GpuLayerBatchedTrunkGemma4 (CudaForwardPass.cs:2011):

• A1. int layerHd = _hp.LayerHeadDim![layer] (:2013) NPEs on non-Gemma → _hp.LayerHeadDim?[layer] ?? _headDim
• A2. _gpu.GeluTanhMul(...) (:2103) hardcoded → dispatch on _hp.FfnActivation (SiLuMul for Silu)
• A3. attnScale: 1f on all 4 attention calls (:2077-2090) → cached field (1f if _isGemma4Like else -1f)
• A4. Generalize the gate IsGemma4BatchedPrefillSupported (:1779) — keep real guards (!_isMoE, !_tqEnabled, !_hasAttnBias, IsNeoxRope, batchable dtypes), drop the implicit Gemma assumption
• A5. Open the Prefill gate (:1732) from _isGemma4Like to "dense + batchable"
• A6. Bit-parity oracle test: batched prefill vs per-token Forward on Qwen3-8B, argmax-stable (flash/GEMM not byte-exact, matches Gemma 4 GPU prefill is ~70x slower than llama.cpp — adopt cuBLAS GEMM for prefill #141)
• A7. Measure prefill t/s before/after on Qwen3-8B Q4_K_M; update README row

What Q4_K gains: flash-attention prefill (biggest lever, quant-agnostic, head_dim=128 ✓) + launch collapse. NOT the compute-bound trunk GEMM (Q4_K stays on memory-bound GEMM-N; that's Item C).

Work Item B — Decode CUDA graphs for non-Gemma (launch-overhead win)

ForwardGemma4 is split prefix(embed/PLE) → capturable device region → download. The non-Gemma Forward (:890) inlines everything with no capture.

• B1. Refactor Forward into token-varying prefix (embed) + RunDeviceRegion(position) (layer loop + final norm/output) + download, keeping TQ ring-advance (:1056) and SnapKV capture (:955) OUTSIDE the region
• B2. Wrap with existing generic capture helpers (TryBeginGraphCapture / TryEndGraphCaptureAndInstantiate / LaunchGraphForPosition); reuse the _tqEnabled || _kvEvictedCount > 0 bail from :1315
• B3. Verify bit-identical (graph replay == direct launch); measure decode t/s on Qwen3-8B; update README

Work Item C — Compute-bound Q4_K prefill (optional, larger; gate on A's profiler)

• C1. (cheaper) Extend MatMulBatchedGemm to dequant Q4_K->fp16->cuBLAS, OR
• C2. (maximal) Q4_K int8 MMQ kernel (llama.cpp mul_mat_q Q4_K reference)

Only pursue if A's prefill profiling shows trunk matmuls (not attention) dominating after flash lands.

Verification

• dotnet build clean (TreatWarningsAsErrors), dotnet test --filter ...ForwardPass
• Argmax-stable parity oracle for A; bit-identical check for B
• Before/after t/s on Qwen3-8B Q4_K_M (model at C:\p\sharpi\models\ or E:\models\)

Sequencing: A first (biggest win + clean oracle), B second, C gated on A's profiler.

[pekkah]

Progress — Work Item A landed (A1–A6 done, A7 in progress)

Branch: perf/cuda-dense-q4k-batched-prefill-156

• A1–A3: 3 Gemma assumptions fixed in the batched layer body (GpuLayerBatchedTrunk): LayerHeadDim null-coalesce → _headDim; FFN activation dispatch on hp.FfnActivation (SiLuMul vs GeluTanhMul); cached _attnScale field (1f Gemma / -1f else → kernel derives 1/sqrt(head_dim)).
• A4–A5: gate generalized — IsGemma4BatchedPrefillSupported → IsBatchedPrefillSupported; Prefill gate dropped _isGemma4Like. Methods renamed PrefillBatchedTrunkGemma4/GpuLayerBatchedTrunkGemma4 → PrefillBatchedTrunk/GpuLayerBatchedTrunk.
• A6: parity test Qwen3CudaBatchedPrefillTests (default argmax-stable + flash-off bit-exact) — both pass. Gemma suite still 5/5 (no regression).

Key finding (latent bug surfaced by the generalization): the batched body applied QK-norm → RoPE, but the dense per-token Forward applies RoPE → QK-norm for weighted (non-L2) QK-norm models like Qwen3. RoPE does not commute with per-channel-weighted RMSNorm, so this diverged by ~9 logits (argmax-stable but wrong). Fixed by ordering the two to match the matching per-token oracle: Gemma keeps norm→rope, dense uses rope→norm. Also now honors NoRopeLayerStep in the batched RoPE like the per-token path. After the fix, flash-off batched vs per-token is bit-exact on Qwen3-8B Q4_K, proving the batched matvec/norm/rope/SwiGLU primitives match and the ordering was the sole discrepancy.

• A7: measure Qwen3-8B Q4_K prefill before/after (SHARPI_BATCHED_PREFILL=0 vs default); update README — in progress.

Next session: Work Item B (decode CUDA graphs for the non-Gemma Forward).

[pekkah]

✅ Work Item A complete & committed

Branch perf/cuda-dense-q4k-batched-prefill-156 (2 commits: 1b4bb4c, 1aff3b6).

Result: Qwen3-8B Q4_K_M prefill 61.2 → 119.8 t/s (1.96×) at ~1K ctx, full GPU offload, RTX 4070 Ti. Decode unchanged (~58 t/s — that's Item B). The win is flash-attention prefill + launch collapse; Q4_K trunk matmuls stay on the memory-bound GEMM-N (compute-bound Q4_K GEMM is Item C).

• A1–A7 all done.
• Tests: Qwen3CudaBatchedPrefillTests (default argmax-stable + flash-off bit-exact) pass. Regression: Gemma4CudaBatchedPrefillTests 5/5, CudaForwardPassSnapKvTests + Gemma4CudaForwardPassTests + Gemma4CudaGraphParityTests 10/10.
• README Qwen3 8B CUDA row updated.

Latent bug fixed along the way: QK-norm/RoPE ordering — the batched body did norm→rope (Gemma order); dense per-token Forward does rope→norm for weighted QK-norm. Now ordered to match the per-token oracle per model family.

Remaining

• Work Item B — decode CUDA graphs for the non-Gemma Forward (launch-overhead win; decode is HBM-bound so the gain is the ~1k-launch/token host overhead). Independent of A. Next session.
• Work Item C (optional) — compute-bound Q4_K prefill (dequant→fp16→cuBLAS, or Q4_K MMQ). Gate on whether A's prefill profiling shows trunk matmuls dominating after flash.

[pekkah]

Follow-ups split into standalone issues

To keep tracking clean now that Item A is done, the remaining items have dedicated issues:

• perf(cuda): decode CUDA graphs for non-Gemma dense Forward (#156 Item B) #158 — Work Item B: decode CUDA graphs for the non-Gemma dense Forward.
• perf(cuda): compute-bound Q4_K prefill matmul — dequant->cuBLAS or Q4_K MMQ (#156 Item C) #159 — Work Item C (optional): compute-bound Q4_K prefill matmul.
• investigate: dense Forward applies RoPE before QK-norm (Qwen3 expects norm before RoPE) #157 — investigation: dense Forward applies RoPE before QK-norm (Qwen3 expects norm before RoPE) — a pre-existing correctness question surfaced while matching the batched path to the per-token oracle in Item A.

This issue (#156) remains the umbrella / Item A. Item A is complete on branch perf/cuda-dense-q4k-batched-prefill-156 (1.96× Qwen3-8B prefill) and can be closed once that branch is merged.

[pekkah]

Work Item C (C1) done — Q4_K compute-bound prefill GEMM

Committed on perf/cuda-dense-q4k-batched-prefill-156 (50bac9d).

Profiler gate (confirmed trunk dominance)

SHARPI_PREFILL_PROFILE=1, Qwen3-8B Q4_K, warm 1008-token prompt, -g -1:

layers=8401ms (attn=40ms  matmul=8329ms)   # before

Attention is 0.5% of layer time after flash; the Q4_K trunk matmuls are 99.5%. The matvec GEMM-N re-streams each weight once per token (llm_matvec_q4k_gemm_n, ~4.5 GB × N of HBM traffic) — memory-bound, exactly C's target. Gate passed → implemented C1.

C1: dequant→fp16→cuBLAS GEMM

• New llm_dequant_q4k_to_f16 kernel — decodes the Q4_K super-block (d*sc*nibble - dmin*mn, identical to llm_embed_lookup_q4k/llm_matvec_q4k) to fp16, one 256-thread block per weight row.
• Generalized MatMulBatchedGemm from Q8_0-only to Q8_0 | Q4_K (dtype-select the dequant kernel + 256-element cols alignment for Q4_K). Everything downstream (activation fp16 convert + cublasGemmEx fp16×fp16→fp32) is shared with the Gemma 4 GPU prefill is ~70x slower than llama.cpp — adopt cuBLAS GEMM for prefill #141 Q8_0 path.
• GpuMatMulBatched routes Q4_K → GEMM (weight read once per batch) with a defensive matvec fallback when cols % 256 != 0.

Parity + perf (same-session A/B, RTX 4070 Ti)

• New oracle Qwen3_8B_BatchedPrefill_Q4KGemm_ArgmaxStable: GEMM vs bit-exact matvec, argmax-equal + top-5 overlap (fp16, not byte-exact — same contract as the Q8_0 Gemma 4 GPU prefill is ~70x slower than llama.cpp — adopt cuBLAS GEMM for prefill #141 GEMM). FlashOff_MatchesSequential pinned to PrefillGemmEnabled=false to keep its bit-exact-matvec purpose. 27 Qwen3/Gemma4 CUDA + MMQ-SoA tests green.
• Prefill: 119.8 → 432 t/s (3.6×), SHARPI_PREFILL_GEMM=0 to bisect. --no-thinking 61.8 → 427. Decode + --tq rows (batched prefill disabled for TQ) unchanged. README updated.

C2 (Q4_K int8 MMQ) — remaining gap, not done

llama.cpp b8585 pp1008 = 5764 t/s on the same GPU; C1 closed the gap from ~48× to ~13×. The rest is llama.cpp's int8 Q4_K mul_mat_q kernel (reads the weight once as int8 with no fp16 HBM temp — the cost C1's fp16 dequant still pays). That's the C2 follow-up; it's a substantial kernel port (Q4_K get_scale_min decode + Q8_1 input + m16n8k32 s8 mma), mirroring how #141 shipped the Q8_0 GEMM and MMQ as separate steps. Deferring unless prefill becomes a priority again.

With A + B + C1 all landed, the umbrella's headline dense-Q4_K wins (Qwen3-8B): prefill 61.2 → 432 t/s (7.1×), decode +7% from CUDA graphs.

[pekkah]

Item C2 — int8 MMQ Q4_K prefill (llm_mmq_q4k) — done, default-on

The maximal Item C: each Q4_K weight is read once as int8 (4-bit nibbles expanded, get_scale_min_k4 decode) and fed straight to the m16n8k32 s8 mma — no fp16 dequant temp to HBM, the cost that capped C1's dequant→fp16→cuBLAS GEMM.

Implementation. Mirrors the llm_mmq_q8_0 tiling/fragment map byte-for-byte (validated by #141), so the mma geometry is the already-proven one. The only Q4_K-specific work:

• 4-bit → int8 weight expansion on load (same lo/hi nibble split as llm_matvec_q4k_gemm_n);
• per-(row, sub-block) (scale, min) unpack via get_scale_min_k4 (K-tile = one 32-element Q4_K sub-block);
• the asymmetric min-bias −super_dmin·mn·(d_a·Σq_a) added once per sub-block. The activation sum d_a·Σq_a is now packed as the fp16 s half of each q8_1 block by llm_quantize_q8_1 (every other reader masks the d-word with 0xffff, so that high half was inert — this mirrors ggml's block_q8_1.ds = {d, s}).

Wired as the default Q4_K prefill path under PrefillMmqEnabled (cols % 256); SHARPI_PREFILL_MMQ=0 reverts to C1, SHARPI_PREFILL_GEMM=0 to the bit-exact matvec. Argmax-stable, not bit-exact.

Perf (Qwen3-8B Q4_K, RTX 4070 Ti, same-session A/B vs C1):

Prompt	C1 GEMM	C2 MMQ	Δ
~100 tok	284 t/s	355 t/s	+25%
~866 tok	428 t/s	434 t/s	+1.3%
1008 tok	430 t/s	432 t/s	~tie

Profiled at N=110 the Q4_K matmul itself drops 353 → 285 ms (−19%). The win is biggest at short context, where C1 still pays its fp16-temp write; it converges to a tie by ~1K ctx as cuBLAS amortizes that temp across the large N. So the 1K README column is unchanged (~432) but short-context prefill and prefill VRAM both improve, and C2 is a strict win-or-neutral everywhere (no regression observed).

Remaining gap to llama.cpp (b8585 pp1008 5764 t/s): its MMQ is cp.async-pipelined, hiding the weight re-read across token tiles that cuBLAS amortizes via L2. A single non-pipelined MMQ kernel matches but does not beat tuned cuBLAS at large N — same kernel-family gap already flagged for Gemma 4 (#149). That pipelining is the follow-up if the 1K column is the target; C2 already captures the no-fp16-temp + short-prompt win.

Tests: CudaMmqQ4KTests (GPU MMQ vs CPU DotQ4K fp32 ref — 0 mismatches, isolates the nibble/scale decode + min-bias + fragment map) and Qwen3_8B_BatchedPrefill_Q4KMmq_ArgmaxStable (MMQ vs bit-exact matvec GEMM-N: argmax + top-5 overlap). 220/220 ForwardPass tests green — the shared quantize_q8_1 change is parity-safe across all q8_1 decode/MMQ paths.

Commit cb72245.

[pekkah]

All work items complete in PR #161 (branch perf/cuda-dense-q4k-batched-prefill-156):

• Item A (batched prefill + flash attn generalized to dense Q4_K) — 1b4bb4c. A1–A5 plumbed (LayerHeadDim?[layer] ?? _headDim, FfnActivation dispatch, cached attnScale, generalized gates); A6 oracle = Qwen3CudaBatchedPrefillTests (argmax-stable vs per-token Forward); A7 README updated.
• Item B (decode CUDA graphs) — shipped under perf(cuda): decode CUDA graphs for non-Gemma dense Forward (#156 Item B) #158, e1de163; Qwen3CudaGraphParityTests bit-identical.
• Item C — C1 dequant→fp16→cuBLAS GEMM (50bac9d) + C2 int8 tensor-core MMQ (cb72245), both argmax-stable, SHARPI_PREFILL_MMQ/_GEMM bisect knobs.
• Follow-up (perf(cuda): Q4_K decode SoA — scale-pre-unpacked repack lifts matvec 74%→89% HBM, +7% decode (#156 follow-up) #160): decode SoA repack + default-on, +7% decode bit-identical.

Net Qwen3-8B Q4_K_M: prefill 61.2 → 432 t/s, decode 65 → 74.7 t/s. Closing via PR #161.