cuda: PQ2_0 mat-vec kernel for 3-8 columns on Ada

[sb32445]

Overview

Adds a mat-vec kernel for PQ2_0 weights with 3-8 activation columns on Ada (cc 8.9). The generic mmvq kernel drops to 36-66 % of DRAM bandwidth from 4 columns on. The new kernel feeds the raw 2-bit codes straight into dp4a with an integer correction term (new activation layout GGML_CUDA_Q8_1_PQ2), 4 rows per warp, __launch_bounds__(128, 3).

RTX 4070, m=5120 k=17408 (cold, ncu): 4 columns 105.6 -> 54.0 us, 8 columns 136.8 -> 78.1 us.
Ternary-Bonsai-2-27B-PQ2_0, parallel decode sequences: 4: +17 %, 6: +25 %, 8: +29 %.

Additional information

• Only Ada, PQ2_0, 3-8 columns, no fusion / ids / batch dims. Everything else takes the old path.
• The branch has two commits: the first contains the environment switch GGML_CUDA_PQ2_MULTICOL (=0 switches the kernel off) that was used for the measurements below, the last one removes it. To reproduce a measurement, build the first commit.
• test-backend-ops MUL_MAT 219/219 (CUDA vs CPU, incl. odd row counts). Perplexity with and without the path 6.378 (difference below 0.07 %, error +-0.53).
• Not measured on other GPUs. The condition is cc == GGML_CUDA_CC_ADA_LOVELACE exactly (8.9), the registers (168, 3 blocks per SM) are tuned for that GPU.
• Not bit-identical to the generic kernel (the integer part is exact, but the float accumulation order differs); details below. Within the new kernel the result of a column does not depend on how many columns share the call (3 to 8), which the generic kernel does not guarantee.
• ggml_cuda_q8_1_layout_host got a plain_2d argument; the layout choice for the quantizer and the kernel choice come from the same y_layout, and the PQ2 branch asserts that no ids, batch dims or fusion are involved (fusion with more than one column does not exist for PQ2_0).
• PTQ1_0 is not affected (own kernel).

Test results

• Hardware / software: RTX 4070 12 GB (AD104, cc 8.9, 504 GB/s, 48 MB L2), Linux 6.18, NVIDIA driver 615.71, CUDA 13.4, GCC 16.2; Release build, -DGGML_CUDA=ON -DCMAKE_CUDA_ARCHITECTURES=89.
• Base: speed numbers were measured on prism at 88c4bc60b; the four commits since (SYCL, WebGPU and cuda: fused FWHT quantizer for 64-wide warps (#303)) touch quantize.cu only in the 64-wide-warp FWHT path, not the code of this PR. The branch is rebased on 2459f68b5, builds, and test-backend-ops was repeated on it.
• Model: Ternary-Bonsai-2-27B in PQ2_0 (m=5120 x k=17408 is one FFN shape of it).
• test-backend-ops test -b CUDA0 -o MUL_MAT: 1586/1586 passed on the rebased branch (CUDA0 against CPU; includes the PQ2_0 cases with odd row counts). On the earlier base the PQ2_0 subset alone was 219/219.
• Kernel, cold L2 (ncu), m=5120 k=17408: 4 columns 105.6 -> 54.0 us, 8 columns 136.8 -> 78.1 us; the generic kernel reaches 36-66 % of the DRAM bandwidth from 4 columns on (46 % at 4, 36 % at 8 columns).
• Decode of Ternary-Bonsai-2-27B-PQ2_0 with parallel sequences (llama-batched-bench): 4 sequences +17 %, 6 sequences +25 %, 8 sequences +29 %.
• Perplexity (256 context, 6 chunks, -ub 8 / 5 / 4) before -> after: 6.3784 / 6.3784 / 6.3776 -> 6.3774 / 6.3774 / 6.3820, a difference of at most 0.07 %, far inside the reported error of +-0.53.
• Numerics, greedy decode with 4 parallel sequences (-np 4, 4 different prompts sent at once, 192 tokens, q4_0 K/V, two rounds per setting): with the new kernel both rounds are identical for all 4 prompts. With the generic kernel (GGML_CUDA_PQ2_MULTICOL=0) two identical runs already differ for 3 of 4 prompts (the arithmetic of the generic kernel depends on the number of columns, and the batch composition varies slightly between runs). Old against new: 2 of 4 prompts identical, the others first differ at character 161 of 848 and 407 of 838.
• Not tested: other GPUs, HIP, MoE (ids) models, column counts outside 3 to 8 (they take the old path), a task-level quality metric.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: The patches were developed with Claude Code (Anthropic's coding agent): it wrote the code, the measurement scripts and the first drafts of the commit messages and PR texts. I decided what to work on (which kernels and host paths to optimise, based on profiles of my own decode setup). The measurements and checks listed in the PR texts were run in the Claude Code sessions; I did not re-run them independently. I will maintain the changes. Commits where Claude Code was used carry a Co-Authored-By trailer.

[bri-prism]

@sb32445 @synchronic1, #243 contains two PQ2_0 multi-column MMVQ commits (2c9103d, 356b782) that cover the same ground as this PR. Could you compare notes and settle on one kernel? @synchronic1, #243 also bundles a chart, IQ4_XS/Q5_0 changes and a fused q8_0 FA path and conflicts in five files, so if there is a part you want reviewed, a separate focused PR per change would be much easier to land.

[synchronic1]

@bri-prism Thanks for flagging the overlap. I’ll split PQ2_0 out and coordinate with sb32445 on shared quantization and dispatch. Ampere here, Ada there; arch-specific paths may make sense, subject to tests.

@sb32445 Happy to coordinate. I’m preparing Ampere correctness and fallback checks; your Ada kernel won’t execute on my 3060 unchanged. Could we agree on shared quantization and dispatch, with each of us validating on our hardware?

[synchronic1]

@bri-prism @sb32445 Thanks for flagging the overlap, and thanks for the Ada work. We completed RTX 3060 (Ampere) validation of the isolated PQ2_0 changes: all 350 test calls passed, and PR306 correctly selected its fallback on Ampere. This supports compatibility of the tested path on Ampere; it does not yet validate a combined patch or establish that the changes merge cleanly.

@sb32445 Would you be willing to validate the isolated Ampere changes and any proposed shared quantization/dispatch changes on your Ada hardware? We don’t currently have Ada hardware available and can cover the Ampere side.

@bri-prism Would you prefer one focused, coordinated PQ2_0 PR retaining the appropriate Ampere and Ada paths, or separate PRs built on shared groundwork? I’ll separate the unrelated changes from PR243 as requested.

[sb32445]

@bri-prism @synchronic1 thanks for the quick Ampere check, and for splitting #243.

I read both kernels against each other, and I think they do not need to be merged into one, because they never select the same hardware:

• cuda: PQ2_0 mat-vec kernel for 3-8 columns on Ada #306 is gated on cc == GGML_CUDA_CC_ADA_LOVELACE (8.9 exactly), PQ2_0, plain 2D, 3-8 columns, and uses its own activation layout (GGML_CUDA_Q8_1_PQ2).
• The PQ2_0 weight-reuse path in Prism b10683 Ganglion Helios port multi slot concurrency speedup from kernel reuse #243 (2c9103d, 356b782) is gated on get_device_table_id(cc) == MMVQ_PARAMETERS_AMPERE. In Prism b10683 Ganglion Helios port multi slot concurrency speedup from kernel reuse #243 that table is cc >= 8.0 && cc < 8.9 (mmvq.cu, get_device_table_id), so Ada never takes it, and it covers 2-5 columns.

So the dispatch is disjoint by architecture: Ampere takes the #243 path, Ada takes #306, everything else (and 2 columns on Ada) stays on the generic kernel. If I read #243 wrongly, please correct me.

What follows from that:

1. Shared groundwork: cuda: PQ2_0 mat-vec kernel for 3-8 columns on Ada #306 does not depend on the Ampere table, and prism does not have it yet, it comes in with Prism b10683 Ganglion Helios port multi slot concurrency speedup from kernel reuse #243. I suggest separate PRs: the PQ2_0 part of Prism b10683 Ganglion Helios port multi slot concurrency speedup from kernel reuse #243 brings its table with it, cuda: PQ2_0 mat-vec kernel for 3-8 columns on Ada #306 stays as it is. Whichever lands second needs a small rebase of mmvq.cu (cuda: PQ2_0 mat-vec kernel for 3-8 columns on Ada #306 touches it with one include and one switch block; the rest is in the new mmvq-pq2_0.cuh, quantize.cu and common.cuh).
2. Validation split: I cannot run the Prism b10683 Ganglion Helios port multi slot concurrency speedup from kernel reuse #243 kernel on Ada (it is not selected there), so my RTX 4070 cannot validate its correctness; that has to stay with the Ampere runs. What I can do on Ada, once the split PQ2_0 PR is up: run test-backend-ops -o MUL_MAT and a PQ2_0 decode and check that it is unchanged against prism, i.e. that the new table and the calc_rows_per_block signature change do not alter Ada.
3. Numerics: Prism b10683 Ganglion Helios port multi slot concurrency speedup from kernel reuse #243 states bit-exactness against the generic kernel; cuda: PQ2_0 mat-vec kernel for 3-8 columns on Ada #306 is not bit-identical to the generic kernel (integer part exact, float accumulation order differs, details in the PR text). As the two paths run on different hardware, a user never mixes them in one run.

If you would rather have one combined PQ2_0 PR with both paths, I am fine with that too, but I think two focused PRs are easier to review and to validate per GPU.

[bri-prism]

Independent validation of the current PR head on an RTX 4090 (native sm_89), CUDA 12.8.93, driver 580.159.03.

Compared the exact parent against the unmodified PR using identical Release builds and the public PQ2_0 model referenced in the PR. The model was downloaded at a pinned revision and its SHA256 verified. Library-loader checks confirmed the intended parent/head libraries. Batch-invariant mode was unset.

Whole-model decode: three alternating AB / BA / AB comparisons with llama-batched-bench, pp128/tg128, F16 K/V, flash attention on, GPU layers 99, batch 2048/microbatch 512, threads 4. Requested context 8192; the executable adjusted capacity to 9216 consistently. These are aggregate synthetic-token decode rates, excluding prompt processing and model loading.

Parallel sequences	Parent median tok/s	PR median tok/s	Gain
4	253.18	255.67	+0.98%
6	325.85	346.91	+6.46%
8	370.83	414.22	+11.70%

Unchanged-dispatch controls at sequences 1/2/9 had median changes of -0.08%/+0.20%/-0.01%. However, individual runs varied substantially, including in those controls: raw paired gains were +0.85/-6.40/+9.51% at sequences 4, +6.24/-4.25/+17.92% at 6, and +11.39/+0.15/+20.68% at 8. Batch 8 is the strongest result; batch 4's small gain is unresolved, and batch 6 was not faster in every pair. This does not reproduce a blanket 17–29% gain on this different Ada GPU.

Operation timings: with a 192MiB cache sweep outside the CUDA event interval, K17408/M5120 improved 52.13→44.03µs at columns 4 and 59.39→45.79µs at 8. K5120/M17408 improved 53.25→45.06µs at 8, but regressed 43.26→45.06µs at 4. These results were repeatable across rounds. Each round used 101 samples after 20 warmups. The timer covers backend graph execution, activation quantization and scheduling gaps; it is not an isolated kernel timer or a reproduction of the reported NCU measurements.

Correctness: parent and head each passed 302/302 PQ2 MUL_MAT,MUL_MAT_ID CUDA-versus-CPU tests. Each also passed 648 finite-output cases covering row tails, columns 1–9 and batches 1/2. All 432 fallback cases were parent/head bit-identical, and 180 first-three-column comparisons across columns 3–8 were bit-identical within the PR. Maximum parent/head absolute difference was 1.52587890625e-5, consistent with the changed accumulation order.

Optional Nsight Compute profiling was denied GPU-counter access, so those attempts are excluded from the timings. No runtime kernel-name trace, generated-text quality, perplexity, long-context or speculative-acceptance result is claimed.

Overall: correctness checks pass and there is a useful, shape-dependent Ada performance benefit, strongest at eight sequences on this card. Please preserve the batch 4 operation regression and whole-model timing variability when describing expected gains.