perf(cuda): Q6_K decode MMQ via SoA reader-port (free AoS) — +17% N=8 (#204)

[pekkah]

Summary

Resolves #204 (umbrella #203). Routes the Q6_K trunk weights (lm-head + ffn_down-half) through the int8 tensor-core decode-MMQ tile — the last big batched-decode chunk (~8.3 ms of the ~25 ms N=8 step was on the bit-exact WS-sw matvec). Measured on Qwen3-8B Q4_K_M @ RTX 4070 Ti (aggregate decode t/s vs the AoS/WS-sw baseline):

N	baseline (AoS)	SoA reader-port	Δ
6	269.0	304.0	+13.0%
7	305.1	346.9	+13.7%
8	327.2	382.1	+16.8% (reproduced 324.7 → 379.6)

Peak VRAM 11394 / 12282 MiB — no spill.

Why a reader-port (the two contained approaches failed on 12 GB)

The decode MMQ needs coalesced staging → a SoA weight layout. Keeping that SoA as a companion beside the interleaved AoS weight ~2.3×'d the Q6_K weight VRAM and spilled the 12 GB card at N=8 (325 → 169 t/s). A zero-VRAM AoS-direct staging avoided the spill but regressed −7…−11% (uncoalesced 210-B block stride).

The fix: repack each 2-D Q6_K matmul weight to the int8 SoA layout once at upload and free the AoS copy (exactly like Q4_K's RepackQ4KSoa) — so SoA is the only copy (~+0.4 GB net, fits N=8). To free the AoS, every Q6_K reader was ported to SoA:

• New _soa variants of llm_matvec_q6k / _n2 / _gemm_n / _ws / _ws_sw and llm_dequant_q6k_to_f16. Each is BIT-IDENTICAL to its AoS counterpart: the SoA stores (q6−32) int8 + verbatim scales + d, so d·scale·(q6−32) and the per-element reduction order are unchanged.
• Dispatch picks the _soa kernel iff the weight handle is SoA-repacked, else AoS.
• The token embedding table (read by embed_lookup, untied from output.weight) is excluded from the repack and stays AoS.
• SHARPI_Q6K_SOA=0 reverts to AoS everywhere (kill-switch); SHARPI_Q6K_DECODE_MMQ=0 keeps SoA readers but disables just the decode-MMQ tile.

The decode-MMQ tile itself (per-16 scales via two m16n8k32 mmas split at the a0/a1|a2/a3 k-boundary, symmetric no-min fixup) is argmax-stable — the same contract as the Q4_K tile.

Validation

• Bit-identical (not just argmax-stable) for every value-carrying reader: CudaGemmQ6KTests.Q6KSoaReaders_AreBitIdenticalToAos (MatMul / N2 / Batched / WS / MatMulBatchedGemm SoA == AoS, byte-for-byte) + MatMulBatchedGemm_Q6K_TracksCpuReference.
• Decode MMQ vs SimdKernels.DotQ6K: CudaDecodeMmqTests.DecodeMmq_Q6K_Soa_TracksCpuReference.
• e2e: all 8 CudaBatchForwardMultiTests + all 5 CudaSpecBatchVerifyTests, incl. the 48-token single-user SpecDecode_GreedyParity_E2E (the ported single-user matvec/WS stays byte-stable) and the prefill-parity oracle.
• Release build clean (0 warnings, TreatWarningsAsErrors). All 18 relevant CUDA tests green on the 4070 Ti.

Closes #204. Refs #203 (umbrella). Caps the batched-decode arc: #262 (+11-28%) → #263 (+6.6%) → this (+17% @ N=8).

🤖 Generated with Claude Code

[pekkah]

/gemini review

[gemini-code-assist]

Code Review

This pull request implements support for repacking 2-D Q6_K weights into a scale-pre-unpacked Structure of Arrays (SoA) layout (Issue #204). This change frees the interleaved Array of Structures (AoS) copy to save VRAM and routes all Q6_K readers (single-token, N=2, GEMM-N, weight-stationary, prefill dequant, and decode-MMQ) to their respective SoA-aware kernels. Feedback on the changes suggests optimizing the 16-bit float scale loading in several CUDA kernels by replacing two 8-bit loads and bitwise operations with a single 16-bit load to reduce instruction count in hot paths.

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 dReg is aligned to 16 bytes and g * 4L is a multiple of 4, the address dReg + g * 4L is 4-byte aligned. We can optimize this by performing a single 16-bit load instead of two 8-bit loads followed by a shift and bitwise OR. This reduces instruction count in the hot path.

        unsigned int dbits = *(const unsigned short*)(dReg + g * 4L);

[gemini-code-assist]

Since dReg is aligned to 16 bytes and g * 4L is a multiple of 4, the address dReg + g * 4L is 4-byte aligned. We can optimize this by performing a single 16-bit load instead of two 8-bit loads followed by a shift and bitwise OR. This reduces instruction count in the hot path.

        unsigned int dbits = *(const unsigned short*)(dReg + g * 4L);

[gemini-code-assist]

Since dReg is aligned to 16 bytes and g * 4L is a multiple of 4, the address dReg + g * 4L is 4-byte aligned. We can optimize this by performing a single 16-bit load instead of two 8-bit loads followed by a shift and bitwise OR. This reduces instruction count in the hot path.

        unsigned int dbits = *(const unsigned short*)(dReg + g * 4L);

[gemini-code-assist]

Since dReg is aligned to 16 bytes and g * 4L is a multiple of 4, the address dReg + g * 4L is 4-byte aligned. We can optimize this by performing a single 16-bit load instead of two 8-bit loads followed by a shift and bitwise OR. This reduces instruction count in the hot path.

        unsigned int dbits = *(const unsigned short*)(dReg + g * 4L);

[gemini-code-assist]

Since dReg is aligned to 16 bytes and g * 4L is a multiple of 4, the address dReg + g * 4L is 4-byte aligned. We can optimize this by performing a single 16-bit load instead of two 8-bit loads followed by a shift and bitwise OR. This reduces instruction count in the hot path.

        unsigned int dbits = *(const unsigned short*)(dReg + g * 4L);

[gemini-code-assist]

Code Review

This pull request introduces support for repacking 2-D Q6_K weights into a scale-pre-unpacked Structure of Arrays (SoA) layout, freeing the original interleaved Array of Structures (AoS) weights to optimize VRAM usage. It implements several SoA-aware CUDA kernels—including single-token, batched, weight-stationary, and dequantization variants—along with a new int8 decode-MMQ tile. Comprehensive tests are added to ensure the new SoA readers are bit-identical to their AoS counterparts and track the CPU reference. The review feedback highlights a platform-dependent issue in the CUDA kernels where long is used for 64-bit offset calculations; since long is 32-bit on Windows (LLP64), it should be replaced with long long or int64_t to prevent potential overflow.

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]

In CUDA C/C++ (and standard C++ on Windows), long is a 32-bit integer (LLP64 model), whereas on Linux it is 64-bit (LP64 model). To ensure consistent 64-bit precision for offset and size calculations across all platforms (preventing potential overflow/truncation issues on Windows if tensors ever grow extremely large), it is highly recommended to use long long (or int64_t) instead of long, and the LL suffix instead of L for literals.

This pattern also applies to the other newly added _soa kernels and templates in this PR.

    long long sb = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    int nb = cols >> 8;
    long long total = (long long)rows * nb;
    if (sb >= total) return;

    long long srcOff = sb * 210LL;
    signed char*   qDst = (signed char*)(dst + sb * 256LL);
    unsigned char* sDst = dst + total * 256LL + sb * 16LL;
    unsigned char* dDst = dst + total * (256LL + 16LL) + sb * 4LL;

[pekkah]

Applied the 16-bit aligned d load across all 5 Q6_K SoA readers (commit pushed). It's a per-super-block load (amortized 1/256), so e2e-neutral on these bandwidth-bound kernels — consistent with the #145 finding — but it's bit-identical and a clean instruction-count reduction, and the bit-identity + 48-token single-user greedy-parity oracles stay green.