cuda: read the WHT signs as packed bits in a warp-local kernel

[TranserZ]

Overview

Packing WHT rotation sign tables into bit masks & giving each 128-element group a single warp.
These changes remove the constant-cache serialization on the per-thread sign reads and output is bit-identical and prefill throughput rises about 6% on a Jetson Orin Nano.

Additional information

Problem

The WHT rotation kernel assigns one element per thread, so a 128-element group is handled by 128 threads that stage the data through shared memory for the butterfly. In __constant__ float SIGNS1[128], each thread reads a different address as SIGNS1[t]. The constant cache can broadcast only one address per warp, so a request for 32 distinct addresses is serialized into 32 replays. This read happens twice per kernel invocation. The last two butterfly stages (h=32, 64) also still use shared memory and barriers.

My idea to solve the problem

I implemented three changes.

First, I used float4 instead of float, so one thread reads 4 consecutive values instead of one. A group is now owned by a single warp (32 lanes) rather than by 128 threads (4 warps). With the group inside one warp, the first two butterfly stages complete in registers and the rest are handled by __shfl_xor_sync, so shared memory and the barriers disappear entirely.

Second, I store each sign as 1 bit instead of a float. The storage for all signs drops from 512 B (4 B * 128) to 16 B (4 B * 4). The signs for the 4 elements owned by lane t (4t..4t+3) land in exactly one contiguous nibble of word t>>3, so the addresses touched by the whole warp drop from 32 to 4, and 8 lanes share each word, which turns the access into a broadcast.

Third, I apply the sign with XOR instead of a floating-point multiply. In IEEE 754 the top bit is the sign, so x ^ (bit << 31) flips it without a -1.0f multiply.

Correctness

Stage order, pairing, and operand order are unchanged from the original kernel, and flipping the sign bit is the same operation as multiplying by -1.0f. The output is therefore expected to be bit-identical, and I confirmed this by measurement:

• Perplexity matches exactly. (wikitext-2, 8 chunks, turbo3/turbo3) Per-chunk values agree to the last digit: 5.8947 / 7.5995 / 8.9344 / 9.4007 / 10.1115 / 10.4895 / 10.9802 / 11.6840, final PPL = 11.6840 +/- 0.70844.

Result

Environment

Item	Value
Board	Jetson Orin Nano 8GB
GPU	Orin, sm_87, 7546 MiB
L4T	R39.2.0
Kernel	6.8.12-1021-tegra
CUDA	13.2.78
gcc	13.3.0
Power mode	15W
GPU clock	pinned at 1020 MHz
CPU	performance / 1497.6 MHz
EMC	pinned at 3199 MHz
Build	Release, GGML_CUDA_FA_ALL_QUANTS=ON
Model	Llama 3B Q4_K_M, 3.21B

Benchmark settings

Item	Value
Context	2048 / 8192
Outer rounds	3
Inner repetitions	-r 3
KV types	turbo3 and f16
Flash attention	on
GPU layers	99

• Stage order, pairing, and operand order are identical to the original kernel, and the sign flip is mathematically the same as a 1.0f multiply, so the output is bit-identical.
• The f16 difference converges to within measurement error.
• The fast path applies only for group size 128 with 16-byte aligned pointers, which is the shape the KV cache uses. Group sizes 64 and 32, and unaligned inputs, keep using the original kernel.
• it gives a meaningful gain : Prefill throughput increased by 6.55% and 5.75% for context length 2048 and 8192, respectively.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES
  I did the problem diagnosis and the design of the optimization (warp-level reorganization, packing the signs into bits, applying the sign with XOR). I used AI for: writing the kernel implementation, writing and running the benchmark harness, aggregating the measurement results, and drafting/translating this pull request description.

[TheTom]

Thanks, this is a clean piece of work. Reviewed and tested; merging once CI re-runs on the one-line push below.

What I checked:

• Arithmetic. The stage order (h = 1, 2, 4, ... 64), operand order (lower + upper / upper - lower) and the h >= 4 lane mapping are the same as k_turbo_wht_f32<dir,128>, and the only substitution is x * (+/-1.0f) becoming a sign-bit XOR, which is exact for every finite value and both zeros. A byte-level float32 simulation of both kernels over 400 random cases (both directions, magnitudes 1e-3 to 1e3, with and without a scale_inv) gave zero mismatches. The one difference is NaN inputs (multiply may return a NaN of either sign, XOR flips it deterministically); NaN in a KV tensor is already a bug, but a one-line comment would save the next reader the question.
• The packed masks. Recomputed from TURBO_WHT_SIGNS1/2 on base (bit e & 31 of word e >> 5, set for -1.0f) they reproduce 0xE46A0359, 0x42F85949, 0xA4C63BBB, 0x49F250A9 and 0x16ACEE90, 0x3F628FDC, 0xB7A5357A, 0xBEA56562 exactly. They are hand-typed though; a comment naming the derivation, or a CPU-side check in test-turbo-quant, would keep them from drifting from the float tables.
• GB10 (sm_121), head e9dba56: -o TURBO_WHT 54/54 (they ran, not 0/0), -o SET_ROWS_TURBO3 and SET_ROWS_TURBO4 pass, -o FLASH_ATTN_EXT 8774/8774 including all turbo K/V cases.
• The InnerQ scale_inv path is not reachable from any test-backend-ops case (ggml_turbo_wht(ctx, a, dir, 0, nullptr) everywhere), so I ran llama-perplexity with Gemma-4-31B Q4_K_M, -ctk turbo3 -ctv turbo3 -fa on, 3 chunks at 2048, base vs this branch: identical to the last digit on every chunk (171.4941, 395.3773, 660.5673). That is the gate for this path.
• wave64. The kernel hardcodes 32 (>> 5, & 31, full mask) where the fork usually uses ggml_cuda_get_physical_warp_size(). It is still correct on HIP wave64 because vendors/hip.h maps the 3-arg __shfl_xor_sync to width warpSize and every mask is 1, 2, 4, 8 or 16, so no shuffle crosses a 32-lane half; and the early return cannot split a warp since all 32 lanes share g. It will run two logical groups per wave there, and the "4 warps per block" tuning is CUDA-shaped. Not measured on AMD; fine to leave.
• Dispatch: the fast path takes every realistic shape (group 128, head_dim % 4 == 0, 16-byte aligned), so the old 128 kernel is reachable only for odd head dims or unaligned pointers. Group sizes 64/32 untouched. Note this does not touch the separate inlined WHT in set-rows.cu's turbo write path, which still reads the float sign tables per thread; the same trick probably applies there if you want a follow-up.
• Hosted CI green across cuda/musa/hip/ubuntu/windows and the HIP VGPR check.

I pushed e5d6c09: (int) cast on the fast-path grid dim, matching the neighbouring launches so MSVC does not warn on the int64_t narrowing.

[jasstrong]

Ran this on AMD, since the wave size is the one thing a Jetson cannot tell you. Three cards: MI210 (gfx90a, CDNA2, wave64), Radeon RX 7900 XTX (gfx1100, RDNA3, wave32) and Radeon PRO V620 (gfx1030, RDNA2, wave32). Short version: correct everywhere, bit-identical on wave64, and the kernel is roughly twice as fast on all three.

Correctness. test-backend-ops -o TURBO_WHT is 27 cases compared against the CPU reference (head_dim 128/256/512, n_heads 1/4/8, both directions, plus the round-trip), and every case covers the fast path since all those head dims are multiples of 4 with group size 128.

card	baseline	with this PR
MI210, wave64	27/27	27/27
7900 XTX, wave32	27/27	27/27
V620, wave32	27/27	27/27

Bit-identity on wave64. Perplexity over 8 chunks at n_ctx 2048, turbo3 K and V, Qwen3.8-27B TQ3_1S, both sides built fresh from the same commit:

baseline  [1]1.9587 [2]1.9579 [3]1.7966 [4]1.7457 [5]1.6889 [6]1.7022 [7]1.6643 [8]1.6028
patched   [1]1.9587 [2]1.9579 [3]1.7966 [4]1.7457 [5]1.6889 [6]1.7022 [7]1.6643 [8]1.6028
Final estimate: PPL = 1.6028 +/- 0.02521   (both)

Kernel throughput. test-backend-ops perf does not register TURBO_WHT, so I added cases locally to both sides to time the kernel in isolation. Largest case, head_dim 256 with 8192 heads, 16 MB moved per run:

card	baseline	with this PR	speedup
MI210	29.61 us, 528 GB/s	15.09 us, 1036 GB/s	1.96x
7900 XTX	26.86 us, 582 GB/s	12.24 us, 1277 GB/s	2.19x
V620	38.65 us, 404 GB/s	14.21 us, 1100 GB/s	2.72x

Across the four sizes I measured the speedup runs 1.8x to 2.1x on the MI210, 1.4x to 2.2x on the XTX and 1.6x to 2.7x on the V620, smallest at the smallest size in each case. Note the working set is resident in Infinity Cache on the two RDNA cards, so their absolute figures sit above DRAM bandwidth; the ratio is the meaningful part.

End to end it is much smaller, at least here. Prefill on the MI210 with the same model and cache type, native TQ kernels and MMQ prefill enabled:

	baseline	with this PR
pp2048	678.34 t/s	679.35 t/s
pp8192	663.13 t/s	664.22 t/s

About 0.15%, consistently positive but nowhere near your 6%. I read that as the rotation being a small share of prefill on a card with this much bandwidth and cache, rather than as a contradiction: your Jetson at 15 W is exactly where constant-cache replays and barriers cost the most. Worth knowing the kernel win is real and large regardless.

One review note on the wave64 path. This comment is not accurate on CDNA:

    // A group is a whole warp, so this return never splits one; the full-mask
    // shuffles below depend on that.

A wavefront there is 64 lanes, so one wavefront carries two independent groups and the early return does split it. The code is still correct, but for a different reason than the one given: every shuffle mask you use is at most 16, so lane ^ mask never leaves the aligned 32-lane block, and no lane ever reads from one that took the other branch. That is worth stating explicitly, because it is the invariant a future stage or a wider mask would break. Two smaller things in the same area: the new code uses the three-argument __shfl_xor_sync, which HIP maps to the runtime wave size, while k_turbo_wht_f32 right above it passes WARP_SIZE explicitly; and 0xffffffff describes 32 lanes, which HIP ignores but which reinforces the wrong mental model. None of this changes behaviour today.

Happy to re-run any of the above if you change the shuffle structure.