Skip to content

cuda: opt-in split-K for the PQ2_0 one-column ffn_down mat-vec on Hopper (GGML_CUDA_PQ2_SPLITK) - #424

Merged
bri-prism merged 1 commit into
prismfrom
cuda/pq2-splitk-ffn-down
Oct 31, 2026
Merged

bri-prism merged 1 commit into
prismfrom
cuda/pq2-splitk-ffn-down

Conversation

@bri-prism

Copy link
Copy Markdown
Collaborator

What

  • Opt-in split-K for the PQ2_0 one-column ffn_down mat-vec on Hopper, behind GGML_CUDA_PQ2_SPLITK=1 (default off).
  • The PQ2_0 multi-column mat-vec is parameterized by rows per group, warps per CTA, split-K factor and unroll; the default instantiation keeps the existing loop and reduction order.

Why

  • At K x M = 12288 x 4096 and 17408 x 5120 the one-column launch leaves warps idle; two warps per row over halves of K keep more of them busy.

How

  • With the switch on, each row's K range is split between two warps whose partial sums are added in warp order through shared memory. This changes the float summation order, so output is not bit-identical (mean KLD about 2e-4 against the default, at the level of ordinary batching differences). Default off: launch path and results are unchanged.

Testing

  • Switch unset or 0: the SASS of all default instantiations is byte-identical to the base build, and greedy output is identical on 6 prompts on two PQ2_0 models.
  • test-backend-ops MUL_MAT passes with the switch unset and on, plus 4 extra cases on the split-K shapes.
  • H200 tg128 with the switch on, alternating runs: +2.8% and +2.6% on the two models.

…per (GGML_CUDA_PQ2_SPLITK)

Parameterize the PQ2_0 multi-column mat-vec by rows per group, warps per CTA,
split-K factor and unroll, and route all weight reads through one helper that
knows the block layout. The default instantiation (rows as before, 4 warps,
split-K 1, unroll 1) keeps the existing loop and reduction order.

GGML_CUDA_PQ2_SPLITK=1 selects 2 rows, 8 warps, split-K 2 for the one-column
plain (non-gate) mat-vec on Hopper when K x M is 12288 x 4096 or 17408 x 5120.
Each row's K range is split between two warps whose partial sums are added in
warp order through shared memory, so the float summation order differs from the
default path. Default off: with the variable unset or 0 the launch path and the
results are unchanged.
@bri-prism
bri-prism merged commit 5628c73 into prism Oct 31, 2026
6 of 14 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant