cuda : CUDA GGML_OP_LIGHTNING_INDEXER implementation (generic vector kernel + wmma kernel)

[fairydreaming]

Overview

This PR adds CUDA implementation of GGML_OP_LIGHTNING_INDEXER. It contains two kernels: generic portable vector kernel and second faster optimized kernel using WMMA.

Requires #24231 (Edit: already merged) to work.

Additional information

I originally wrote this months ago for #21149, now only updated it to match the CPU implementation (scales removed, added mask).

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES, AI was used to discuss possible optimizations.

[JohannesGaessler]

From a cursory look at the code I get the impression that it would benefit from the use of the primitives in mma.cuh. The interface in nvcuda::wmma does not have a defined memory layout in registers so any non-scalar operations need to go through shared memory. To be clear: from my side this is advice regarding the implementation, not a requirement for a review/merge.

[edenfunf]

Tested PR #25545 at e624d0f on an RTX 5070 under Windows 11.

Environment:

• NVIDIA GeForce RTX 5070, sm_120, 12 GB, VMM enabled
• CUDA 13.3 (nvcc V13.3.73), Visual Studio 2026, CMake VS generator
• CMAKE_CUDA_ARCHITECTURES=120a

Correctness:

test-backend-ops test -o LIGHTNING_INDEXER -b CUDA0
96/96 tests passed
Backend CUDA0: OK

Both kernel paths were covered on sm_120: the nh=64 cases with F16/quantized K types used the WMMA kernel, while the remaining cases used the vector kernel. I also observed approximately 14.5–15.0 TFLOPS for the larger nh=64 q4_0 test cases:

LIGHTNING_INDEXER(hsk=128,nh=64,kv=256,nb=512,ns=1,nm=1,type_K=q4_0):   6721 runs -  149.06 us/run -  14.52 TFLOPS
LIGHTNING_INDEXER(hsk=128,nh=64,kv=256,nb=2048,ns=4,nm=4,type_K=q4_0):   435 runs - 2309.64 us/run -  14.99 TFLOPS
LIGHTNING_INDEXER(hsk=128,nh=32,kv=256,nb=512,ns=1,nm=1,type_K=q4_0):  11904 runs -   84.32 us/run -  12.83 TFLOPS

[fairydreaming]

From a cursory look at the code I get the impression that it would benefit from the use of the primitives in mma.cuh. The interface in nvcuda::wmma does not have a defined memory layout in registers so any non-scalar operations need to go through shared memory. To be clear: from my side this is advice regarding the implementation, not a requirement for a review/merge.

@JohannesGaessler I'm aware of that, but I haven't gotten that far down the CUDA matrix multiplication rabbit hole yet. Leaving that as something for the future.

[am17an]

I agree with @JohannesGaessler that mma would be much faster but this that can addressed in a later PR.

[JohannesGaessler]

To clarify my position: the only thing that I consider a hard requirement for my approval is a check of the compilation time since that can potentially lead to a stall. Other than that I would consider my comments to only be suggestions.

[fairydreaming]

To clarify my position: the only thing that I consider a hard requirement for my approval is a check of the compilation time since that can potentially lead to a stall. Other than that I would consider my comments to only be suggestions.

@JohannesGaessler On my machine (Epyc 9374F):

$ time /usr/local/cuda-12.8/bin/nvcc -forward-unknown-to-host-compiler -DGGML_BACKEND_BUILD -DGGML_BACKEND_SHARED -DGGML_CUDA_PEER_MAX_BATCH_SIZE=128 -DGGML_CUDA_USE_GRAPHS -DGGML_SCHED_MAX_COPIES=4 -DGGML_SHARED -D_GNU_SOURCE -D_XOPEN_SOURCE=600 -Dggml_cuda_EXPORTS --options-file CMakeFiles/ggml-cuda.dir/includes_CUDA.rsp -O3 -DNDEBUG -std=c++17 "--generate-code=arch=compute_120a,code=[sm_120a]" -Xcompiler=-fPIC -use_fast_math -extended-lambda -compress-mode=size -Xcompiler "-Wmissing-declarations -Wmissing-noreturn -Wall -Wextra -Wpedantic -Wcast-qual -Wno-unused-function -Wno-array-bounds -Wextra-semi -Wno-pedantic" -MD -MT ggml/src/ggml-cuda/CMakeFiles/ggml-cuda.dir/lightning-indexer.cu.o -MF CMakeFiles/ggml-cuda.dir/lightning-indexer.cu.o.d -x cu -c /home/phm/projects/llama.cpp-deepseek-lid-cuda/ggml/src/ggml-cuda/lightning-indexer.cu -o CMakeFiles/ggml-cuda.dir/lightning-indexer.cu.o

real	0m4.357s
user	0m3.737s
sys	0m0.620s

[am17an]

Good job!

[fairydreaming]

Good job!

In case anyone wants to use it in the meantime it's currently broken. Good job indeed, me.

Edit: should be OK now

[FullstackSensei]

Just wanted to share that that this PR works on Mi50 using ROCm and P40 CUDA. Model is fully in VRAM. Pulled and compiled just before the merge

[jaholmesuk]

Datapoint from Ampere multi-GPU, since this landed with a big effect for us.

DeepSeek-V4-Flash UD-Q4_K_XL (145GB) fully resident on 4x A100 80GB PCIe, -fa on -ngl 999 --parallel 1, single stream:

• b9986 (pre-merge): 18.8 t/s decode / ~400 t/s prefill on a 1.9k-token prompt. The log showed resolve_fused_ops: Lightning Indexer not supported, set to disabled.
• master 86a9c79 (post-merge): 41.0 t/s decode / ~492 t/s prefill. 2.2x decode, and the warning is gone.

Same box, model, and args. The delta spans about two days of commits so not all of it is necessarily this PR, but the indexer was the only fused op reported as disabled on the earlier build. Worth noting we also measured 2-GPU vs 4-GPU as flat pre-merge (18.9 vs 18.8), which ruled out hop overhead and is consistent with the unfused indexer being the per-token bottleneck rather than interconnect.

Happy to run further tests on this hardware if it is useful.