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