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

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

* chore : remove indentation of #pragma unroll

* cuda : remove unnecessary kernel template declarations

* cuda : add WARPS_PER_BLOCK and K_VECS_PER_BLOCK template parameters in lightning indexer kernels to avoid duplication of constants.

* cuda : relax MMA architecture requirements to Turing in lightning indexer implementation

* chore : renamed variables

* chore : rename ggml_cuda_op_lightning_indexer() to ggml_cuda_lightning_indexer()

* chore : TODO for AMD rocWMMA

* chore : whitespace formatting

* chore : another variable rename to fix problems caused by shadowing

* chore : yet another rename, this time uppercased all constants

* cuda : added alignment checks for Q and K tensors in lightning indexer implementation

---------

Co-authored-by: Stanisław Szymczyk <sszymczy@gmail.com>

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -65,6 +65,7 @@
 #include "ggml-cuda/tri.cuh"
 #include "ggml-cuda/cumsum.cuh"
 #include "ggml-cuda/fill.cuh"
+#include "ggml-cuda/lightning-indexer.cuh"
 #include "ggml.h"
 
 #include <algorithm>
@@ -2257,6 +2258,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
         case GGML_OP_FILL:
             ggml_cuda_op_fill(ctx, dst);
             break;
+        case GGML_OP_LIGHTNING_INDEXER:
+            ggml_cuda_lightning_indexer(ctx, dst);
+            break;
         default:
             return false;
     }
@@ -4987,6 +4991,8 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
         case GGML_OP_DIAG:
         case GGML_OP_SOLVE_TRI:
             return true;
+        case GGML_OP_LIGHTNING_INDEXER:
+            return ggml_cuda_lightning_indexer_supported(dev_ctx->device, op);
 
         default:
             return false;