ggml-org/llama.cpp b11405

on GitHub

cuda: tile the lightning indexer over keys and tokens for 4 heads (#29901)

• cuda: tile the lightning indexer over keys and tokens for 4 heads

With too few heads for a wmma tile, a block scores 64 keys against 8
tokens: the keys are staged once in half precision, the queries one
head at a time, and each thread owns one key for two tokens, so no dot
product needs a cross thread reduction. Batches smaller than a token
tile keep the vector kernel. test-backend-ops measures 4 heads.

• cuda: multiply the lightning indexer tile in float

Address review from am17an: the half2 products overflow once a single
q * k exceeds the f16 range. The queries stay in float in shared memory
and each half2 of keys is widened once for both tokens, so every
product and sum is computed in float.

• cuda: widen each lightning indexer key once for all heads

The tile kernel stages the queries and weights of every head at once,
so each key element is widened from half once and feeds all heads,
with a single barrier. F16 keys are copied into the tile without a
float round trip. Keeping the keys in float in shared memory measures
slower, the occupancy drops.

• cuda: stop the lightning indexer tile from spilling registers on ROCm

Each thread of the tile kernel now scores two keys for a single token,
so a warp shares its token and the query reads are broadcasts: six
shared reads per element pair instead of nine for the same products.
The inner loop is unrolled by 8, which keeps gfx908 at 63 VGPRs with no
spill where the fully unrolled loop needed over a thousand, and makes
the kernel 36x faster on an R9700 and slightly faster on CUDA.