perf(gemma4-cuda): MMQ + flash-attention prefill + dp4a decode (#141/#142)

[pekkah]

Summary

Closes the Gemma 4 E4B CUDA performance gap to llama.cpp on prefill (issues #141 / #142), on an RTX 4070 Ti with gemma-4-E4B-it-Q8_0.

	before this branch	after	llama.cpp (same box)
Prefill (~1K ctx)	109 t/s	2853 (26×)	~8000
Prefill (~1.8K ctx)	—	~2920	—
Decode @d0	~49	59	~78

The two foundational commits (#141 cuBLAS-GEMM prefill + #142 dp4a/graphs decode) landed first (109→1564 prefill, 49→59 decode). The rest of this branch then closed most of the remaining prefill gap after profiling revealed the real bottleneck.

The key finding

The issue framed prefill as matmul-bound, but profiling at a realistic prompt length (N=1848) showed attention was 67% of prefill — it's O(n²) and the old "~30%" note was a short-prompt artifact — and memory-bound on redundant K/V reads (the scalar per-query kernel re-streamed each query's K/V window up to ~512× for SWA layers; ~79× off compute peak). So the lever was attention, not the matmul.

What's in the branch

Prefill (#141):

• int8 tensor-core MMQ matmul llm_mmq_q8_0 — inline-PTX mma.sync.m16n8k32.s8, Q8_0 weights × Q8_1 activations, register-prefetch double-buffered. Reads each weight once as int8 with no fp16 HBM dequant temp; beats cuBLAS GEMM (matmul 332→316ms). Default-on (SHARPI_PREFILL_MMQ=0 reverts).
• flash-attention prefill llm_flash_attn_prefill_f32 — shared K/V tiles reused across a 16-query tile + online softmax (FA2 style), half2 fp16x2 QK dot via inline PTX (K staged fp16, V fp32 for exact PV). Attention 929→245ms (3.8×). Default-on (SHARPI_PREFILL_FLASH=0 reverts).
• batched Q8_0 embedding llm_embed_lookup_q8_0_batched — collapses 2·N host launches into one (embed 62→17ms).
• cuBLAS-GEMM prefill foundation + F32 PLE via TF32 Sgemm.

Decode (#142): dp4a/Q8_1 int8 matvec + CUDA graphs default-on.

Net: Gemma 4 prefill ~1389 → 2920 t/s (2.05×) this session; 109 → 2853 across the whole branch. Decode untouched this session.

Correctness

MMQ and flash are argmax-stable, not bit-exact (int8 / fp16 rounding) — by design, consistent with the existing cuBLAS-GEMM prefill posture. The bit-exact parity oracles deliberately pin them off; new tolerance oracles cover the fast paths:

• CudaMmqQ8_0Tests (int8 dot vs fp32 DotQ8_0)
• CudaFlashAttnTests (vs both scalar attention kernels: GQA, head_dim 256/512, windowing, partial tiles)
• Gemma4_E4B_BatchedPrefill_{Mmq,FlashAttn}MatchesSequential (e2e argmax-stability)

22 focused CUDA tests green; the full CUDA suite (120) was green before the final additive (batched-embed) commit. Reviewed with pr-review-toolkit (code-reviewer + silent-failure-hunter): no critical/important findings.

Remaining gap

Prefill is now well-balanced (matmul 45% / attention 36% / embed+ple 19%); closing the rest needs a full tensor-core flash at d=512 (occupancy-limited even in llama.cpp) and decode-side work (dp4a qs-alignment).

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request introduces significant performance optimizations for Gemma 4 prefill and decode paths, including a memory-efficient flash-attention prefill kernel, an int8 tensor-core MMQ matmul, a dp4a-based Q8_0 decode matvec, and batched embedding lookups. The code review feedback suggests several high-impact optimizations: replacing runtime integer division and modulo operations with bitwise shifts and masks in the flash-attention kernel, consolidating dual 8-bit loads into single 16-bit loads for scale factors across multiple CUDA kernels, fixing a late overflow check in MatMulBatchedMmq, and utilizing CollectionsMarshal.AsSpan to avoid allocations when processing token lists.

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]

Since head_dim (and thus hd2) is always a power of two in LLMs (e.g., 256, 512), we can completely avoid the expensive runtime integer division and modulo operations by using bitwise shifts and masks. This significantly improves the performance of the key loading phase in the flash attention prefill kernel.

        int shift_k = __ffs(hd2) - 1;
        int mask_k = hd2 - 1;
        for (int idx = tid; idx < kt_tile * hd2; idx += (int)blockDim.x) {
            int kk = idx >> shift_k;
            int pr = idx & mask_k;
            unsigned int kh = 0u;
            if (kk < tile_keys) {
                long off = (long)(kt0 + kk) * kv_dim + (long)kv_head * head_dim + 2 * pr;
                kh = sharpi_f32x2_to_f16x2(k_cache[off], k_cache[off + 1]);
            }
            sKh[idx] = kh;
        }

[gemini-code-assist]

Similarly to the key loading phase, head_dim is always a power of two, so we can replace the runtime integer division and modulo with bitwise operations to optimize the value loading phase.

        int shift_v = __ffs(head_dim) - 1;
        int mask_v = head_dim - 1;
        for (int idx = tid; idx < kt_tile * head_dim; idx += (int)blockDim.x) {
            int kk = idx >> shift_v;
            int d = idx & mask_v;
            sV[idx] = (kk < tile_keys)
                ? v_cache[(long)(kt0 + kk) * kv_dim + (long)kv_head * head_dim + d]
                : 0.f;
        }

[gemini-code-assist]

Loading the 16-bit fp16 scale d using two separate sharpi_byte_at calls is inefficient. Since b0 is always 2-byte aligned, we can perform a single 16-bit load by casting the pointer to const unsigned short*, which compiles to a single ld.global.u16 instruction.

            unsigned int d_bits = *reinterpret_cast<const unsigned short*>(reinterpret_cast<const char*>(weights) + b0);
            float dw = sharpi_fp16_to_fp32(d_bits);

[gemini-code-assist]

We can optimize the scale loading inside the MMQ_LOAD_TILE macro by performing a single 16-bit load instead of two sharpi_byte_at calls.

        if (tid < MMQ_BM) { long wb = ((long)(row_block + tid) * nb + (KB)) * 34L; \
            rWd = (row_block + tid < rows) ? sharpi_fp16_to_fp32(*reinterpret_cast<const unsigned short*>(reinterpret_cast<const char*>(weights) + wb)) : 0.f; } \

[gemini-code-assist]

We can optimize the scale loading inside llm_dequant_q8_0_to_f16 by performing a single 16-bit load instead of two sharpi_byte_at calls.

        unsigned int d_bits = *reinterpret_cast<const unsigned short*>(reinterpret_cast<const char*>(weights) + b0);
        float d = sharpi_fp16_to_fp32(d_bits);

[gemini-code-assist]

The overflow check (long)cols * nTok > int.MaxValue is performed after casting the value to int on line 1824. If the value exceeds int.MaxValue, qN will already contain an overflowed/truncated value before the check runs. Performing the check on the long value first is safer and cleaner.

            long totalElements = (long)cols * nTok;
            if (totalElements > int.MaxValue)
                throw new InvalidOperationException(
                    $"MatMulBatchedMmq: cols*nTok ({totalElements}) exceeds int range.");
            int qN = (int)totalElements;

[gemini-code-assist]

To avoid array allocations when tokens is a List<int>, we can use CollectionsMarshal.AsSpan to get a direct span of the list's underlying array. This completely avoids allocations for both int[] and List<int> collections.

            ReadOnlySpan<int> idsSpan = tokens is int[] arr ? arr :
                                        tokens is System.Collections.Generic.List<int> list ? System.Runtime.InteropServices.CollectionsMarshal.AsSpan(list) :
                                        System.Linq.Enumerable.ToArray(tokens);
            var idTensor = _gpu.UploadRaw(
                System.Runtime.InteropServices.MemoryMarshal.AsBytes<int>(idsSpan),
                TensorShape.D1(N), DType.Float32);

[pekkah]

Reviewed the Gemini code-assist suggestions — all 7 are legitimate medium-priority micro-opts. Deferred to #145 rather than amending a CI-green, tested, review-toolkit-clean PR at merge time (the two shift/mask suggestions also bake a power-of-two head_dim assumption into an otherwise-general kernel — needs a guard first). Remaining perf gap tracked in #141/#142; full tensor-core flash at d=512 in #146.