perf(cuda): split-K + better-tiled int8 MMQ GEMM — 23-34% → 60-80% of TC peak (#141 follow-up)

[pekkah]

✅ RESOLVED / superseded by #152 (2026-06-14)

Of the approaches here: split-K was ruled out (occupancy isn't the bottleneck at real prefill batch sizes — end-to-end wash, reverted), cp.async staging was slower on Ada (reverted), and the SoA Q8_0 repack shipped (+10–12% prefill, bit-identical). The int8 MMQ still sits at ~40–43% of TC peak; that remaining "different tiling regime" lever — the only one not ruled out — is consolidated in #152. Closing in favor of #152 to avoid duplicate tracking.

---

Follow-up to #141. After the TC flash work (#146/#147) moved the prefill bottleneck to the matmul/FFN GEMMs, a roofline probe (CudaMmqRooflineProbe) shows the int8 MMQ kernel runs at only 23-34% of the 4070 Ti's ~160 TOPS dense int8 TC peak — ~3× headroom, matching the 2.2× end-to-end prefill gap to llama.cpp (~8475 vs our 3818 t/s).

Worst shape: ffn-down [out=2048 × in=8192], 23% peak. Few output rows → few CTA tiles → the SMs are occupancy-starved over the long (8192) contraction, exactly the class of problem the single-warp flash kernel had before #147's d-split.

Scope (pieces, not estimates):

• Split-K the contraction across CTAs for tall/thin-output GEMMs (ffn-down) so all SMs get work; reduce the partials (atomic or a second pass). This alone should lift ffn-down's 23%.
• Better register tiling / larger per-warp output tiles to raise arithmetic intensity (more mma per weight load).
• cp.async double-buffered weight/activation staging — complicated by the Q8_0 qs 2-byte misalignment (the same funnelshift issue noted for decode); may need a re-aligned Q8_0 staging or a 16-byte-aligned shared layout.
• Cross-reference llama.cpp ggml/src/ggml-cuda/mmq.cuh (on disk at C:\p\llama.cpp) for their tiling/pipelining constants.
• Parity: the existing CudaMmqQ8_0Tests oracle (argmax-stable vs CPU DotQ8_0); A/B with the roofline probe (target: push the three FFN/qkv shapes toward 60-80% peak).

This is the main remaining lever for the Gemma 4 prefill half of the #141/#142 "match llama.cpp" goal.

[pekkah]

Tried split-K. Negative result — occupancy is not the lever at real prefill batch sizes. Redirecting this issue.

Implemented split-K (grid.z over the K-loop, atomicAdd partials) and benched it:

• Isolated GEMM (CudaMmqRooflineProbe, nTok=1024): ffn-down 36.4 → 46.1 TOPS (23→29%, +27%) — split-K does fix the isolated tall/thin shape.
• End-to-end prefill A/B (Gemma 4, all-GPU, warm): N=965 +2.3% (3391→3469), N=2054 −2% (4049→3966). Both inside the ~6% run-to-run variance — a prompt-length-conditional wash, not a win. Reverted (it added branches + a Clear + atomics to the hottest kernel for a noise-floor effect).

Why the isolated win didn't translate: the probe used nTok=1024, but real prefill batches all prompt tokens at once. ffn-down's tok-block count is ceil(N/128): at N=965 that's 8 (→256 CTAs, starved, matches the probe), but at N=2054 it's 17 (→544 CTAs, not starved). So at realistic prefill sizes the MMQ GEMMs already have enough CTAs — occupancy isn't the bottleneck.

The real lever (revised scope): the ~30% TC-peak ceiling is inner-loop-efficiency-bound, not occupancy-bound:

• the Q8_0 qs 2-byte misalignment taxes every weight word with a __funnelshift (no aligned vector load, no cp.async);
• register-prefetch double-buffering instead of true cp.async shared-staging pipelining;
• per-block fp32 scale accumulation overhead in the inner loop.

So matching llama.cpp's MMQ needs a re-aligned Q8_0 staging + cp.async-pipelined inner loop (cross-ref ggml/src/ggml-cuda/mmq.cuh), not split-K. That's the actual remaining prefill lever — a deeper kernel rewrite. The CudaMmqRooflineProbe stays as the measurement harness (but future probes must use the real prefill nTok, not 1024, to avoid the occupancy-mismeasurement that misled this attempt).

[pekkah]

Update: validated the SoA fix in isolation — it works, +10-19% bit-identical. (Commit 4d65210, branch perf/gemma4-cuda-flash-tc-146.)

Built llm_mmq_q8_0_soa + MatMulBatchedMmqSoa reading Q8_0 weights from a struct-of-arrays repack (quants 16 B-aligned, fp16 scales separate), so every weight word is a plain aligned load instead of the AoS __funnelshift. Validated (CudaMmqSoaTests):

• Bit-identical to the interleaved MMQ (maxAbs 0, 0 diffs) across 5 shapes incl. ffn-down & qkv.
• Speed at real prefill nTok=2048 (not the 1024 that mismeasured split-K):

  GEMM	AoS	SoA	Δ
  ffn-gate/up	54.6 TOPS	61.4	+11.1%
  ffn-down	39.6 TOPS	49.1	+19.4%
  qkv	54.6 TOPS	60.5	+9.7%

A per-instruction win (not occupancy), so it should translate end-to-end — and it's bit-identical, so the migration is verifiable.

Remaining: the model-wide migration (the high-risk part). Each Q8_0 weight is read by both prefill MMQ and decode matvec, so a weight is either interleaved or SoA — all its readers must switch together (no per-weight VRAM doubling on a 12 GB card). Coordinated change:

• Upload: repack Q8_0 → SoA at load (GPU repack kernel or CPU pre-upload). Layout: per tensor [quants rows*cols B][scales rows*nb fp16]; host passes (quantsPtr, scalesPtr=base+rows*cols) to each kernel (host knows rows/cols).
• Convert readers (each verifiable bit-identical): llm_matvec_q8_0, llm_matvec_q8_0_dp4a (decode Gemma 4 GPU decode is ~1.8-2.9x slower than llama.cpp (kernel efficiency + low-ctx clock cliff) #142 — same funnelshift class, so this helps decode too), llm_matvec_q8_0_gemm_n, llm_embed_lookup_q8_0 (+batched), llm_dequant_q8_0_to_f16, and wire llm_mmq_q8_0_soa in.
• Verify with the full bit-exact Gemma4 suite (the format change is load-bearing; focused tests aren't enough).

This is the next session's main task — mechanical but coordinated, gated on the full suite. The lever is proven; execution is de-risked by bit-identicality.

[pekkah]

Wired SoA end-to-end into Gemma 4 — +10-12% prefill, bit-identical. (Commit 131fe06.)

The validated SoA MMQ kernel is now integrated into the model:

• llm_q8_0_repack_soa + RepackQ8_0Soa: one-time GPU repack of 2-D Q8_0 GEMM weights at upload (interleaved 34 B/block → [quants rows*cols B][scales rows*nb fp16]), source freed, handle marked in _soaHandles.
• llm_mmq_q8_0_soa (refactored to the AoS signature) + llm_matvec_q8_0_dp4a_soa: aligned-load variants. MatMulBatchedMmq and the decode matvec auto-route per repacked handle — no caller changes. output.weight auto-routes via the set.
• mmqSoa flag (SHARPI_MMQ_SOA, default off).

Verification:

• Bit-identical end-to-end (CudaMmqSoaE2ETests: Gemma 4 prefill + 4 greedy decode steps, maxAbs(SoA−AoS) = 0).
• Kernel-level bit-identical + isolated speed (CudaMmqSoaTests).
• 26 focused Gemma4/Cuda tests green (SoA off by default → existing paths untouched).

End-to-end A/B (all-GPU, warm, gemma-4-E4B Q8_0, 4070 Ti):

N	SoA off	SoA on	Δ
965	3297	3698	+12.2%
2054	3849	4240	+10.2%
Decode neutral (BW-bound, low arithmetic intensity).

Prefill now ~4240 t/s (was 2853 at the start of this push) — 50% of llama.cpp's ~8475, up from 34%.

Remaining for default-on (follow-up): the 3 non-default Q8_0 readers — llm_matvec_q8_0 (fp32, decode when dp4a off), llm_matvec_q8_0_gemm_n, llm_dequant_q8_0_to_f16 — still read interleaved. They're only exercised by the bit-exact/gemm-off oracles, but a half-migration would let those oracles falsely pass on garbage (they compare batched-vs-sequential, both reading SoA-as-interleaved → self-consistent but wrong). So default-on requires converting all three (each a mechanical bit-identical change, same pattern as the dp4a variant), then flip SHARPI_MMQ_SOA default and re-run the full bit-exact suite. Kept opt-in until then.

[pekkah]

SoA Q8_0 is now DEFAULT-ON — migration complete. (Commit 0955e99.)

Converted the three remaining interleaved readers (llm_matvec_q8_0 fp32, llm_matvec_q8_0_gemm_n, llm_dequant_q8_0_to_f16) to bit-identical SoA variants. All five Q8_0 readers (MMQ, dp4a, fp32 matvec, GEMM-N, dequant) now auto-route to the aligned-load SoA kernel per repacked handle; the embedding stays interleaved (uploaded separately).

mmqSoa defaults on (SHARPI_MMQ_SOA=0 reverts). Because every reader is now SoA-aware, the bit-exact Gemma 4 oracles exercise the SoA path and pass bit-exactly — GemmOff_MatchesSequentialBitExact included (this is the key check: a half-migration would have let it falsely pass on garbage, since it compares batched-vs-sequential and both sides must read the format correctly). 26 focused Gemma4/Cuda tests green; full ForwardPass suite running as the merge gate.

End-to-end, now by default (all-GPU, warm, gemma-4-E4B Q8_0, 4070 Ti):

N	off (=0)	default	Δ
965	3297	3698	+12.2%
2054	3849	4240	+10.2%

Prefill ~4240 t/s = ~50% of llama.cpp's ~8475 (was 34% before the TC-flash + SoA work). Closing #149 once the full suite confirms; the next lever toward parity is cp.async-pipelined staging on top of the now-aligned SoA layout (the alignment this unblocks) + the decode track (#142).

[pekkah]

Tried cp.async-pipelined weight staging on the SoA MMQ → negative result, reverted.

Built llm_mmq_q8_0_soa_cpa: 2-stage double-buffered cp.async of the (now 16-byte-aligned) weight quant tile global→shared, while register-prefetching scales+activations. Bit-identical to the SoA MMQ (correctness fine), but slower on the 4070 Ti at real nTok=2048:

GEMM	SoA (reg-prefetch)	cp.async .cg	cp.async .ca
ffn-gate/up	61.9 TOPS	−9.6%	−5.4%
ffn-down	48.7	−0.3%	+0.3%
qkv	61.2	−13.1%	−8.5%

.ca (L1-cache) beat .cg (L1-bypass) since the weight rows are re-read per tok-block (the gy grid), but both lose to the existing register-prefetch double-buffer. On Ada the register-prefetch already hides the global-load latency well, and cp.async's commit_group/wait_group barriers add overhead without a compensating win at this tile size. Reverted (no point carrying a slower kernel).

Conclusion for #149: the SoA repack (+10-12%, shipped, default-on) is the practical win for this MMQ structure; cp.async is not the next lever on Ada for this kernel. Closing the remaining prefill gap to llama.cpp would need a different tiling regime (larger output tiles / more N-tiles per warp to raise arithmetic intensity, or a fundamentally different MMQ schedule) — a deeper rewrite with uncertain payoff. The other open lever is decode (#142). Marking #149's SoA work done; cp.async explicitly ruled out.