From cdfbdcbfb56b91d82daff98a7478200a58afb0c7 Mon Sep 17 00:00:00 2001 From: Spencer Zaid Date: Thu, 2 Jul 2026 12:33:50 -0700 Subject: [PATCH 1/2] deepseek4 : wire ggml_lightning_indexer into build_lid_top_k --- src/models/deepseek4.cpp | 19 ++++--------------- 1 file changed, 4 insertions(+), 15 deletions(-) diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp index 759654228e36..d6c4eb27021a 100644 --- a/src/models/deepseek4.cpp +++ b/src/models/deepseek4.cpp @@ -582,21 +582,10 @@ ggml_tensor * llama_model_deepseek4::graph::build_lid_top_k( indexer_weights->ne[0], indexer_weights->ne[1]/n_stream, indexer_weights->ne[2], n_stream, indexer_weights->nb[1], indexer_weights->nb[2]/n_stream, indexer_weights->nb[3]/n_stream, 0); - indexer_q = ggml_permute(ctx0, indexer_q, 0, 2, 1, 3); - cb(indexer_q, "lid_q", il); - indexer_k = ggml_permute(ctx0, indexer_k, 0, 2, 1, 3); - cb(indexer_k, "lid_k", il); - - ggml_tensor * indexer_kq = ggml_mul_mat(ctx0, indexer_k, indexer_q); - cb(indexer_kq, "lid_kq", il); - - indexer_kq = ggml_cont(ctx0, ggml_permute(ctx0, indexer_kq, 2, 1, 0, 3)); - cb(indexer_kq, "lid_kq", il); - - ggml_tensor * indexer_score = ggml_relu(ctx0, indexer_kq); - indexer_score = ggml_mul(ctx0, indexer_score, indexer_weights); - indexer_score = ggml_sum_rows(ctx0, indexer_score); - indexer_score = ggml_cont(ctx0, ggml_permute(ctx0, indexer_score, 2, 1, 0, 3)); + // fused DSA lightning indexer: equivalent to the permute/mul_mat/relu/mul/ + // sum_rows path, but avoids the [n_kv x n_head x n_batch] score intermediate. + // weights are pre-scaled above, so scale_embd = scale_heads = 1.0. + ggml_tensor * indexer_score = ggml_lightning_indexer(ctx0, indexer_q, indexer_k, indexer_weights, 1.0f, 1.0f); cb(indexer_score, "lid_score", il); indexer_score = ggml_add(ctx0, indexer_score, inp_lid.kq_mask); From fbb92d285a80193fd3a6bf54a9130ed36ae523b0 Mon Sep 17 00:00:00 2001 From: Spencer Zaid Date: Thu, 2 Jul 2026 12:34:01 -0700 Subject: [PATCH 2/2] ggml-cuda : add GGML_OP_LIGHTNING_INDEXER kernel --- ggml/src/ggml-cuda/ggml-cuda.cu | 6 + ggml/src/ggml-cuda/lightning-indexer.cu | 148 +++++++++++++++++++++++ ggml/src/ggml-cuda/lightning-indexer.cuh | 3 + 3 files changed, 157 insertions(+) create mode 100644 ggml/src/ggml-cuda/lightning-indexer.cu create mode 100644 ggml/src/ggml-cuda/lightning-indexer.cuh diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index f3fb32452d2c..f76603208c98 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -28,6 +28,7 @@ #include "ggml-cuda/fwht.cuh" #include "ggml-cuda/getrows.cuh" #include "ggml-cuda/im2col.cuh" +#include "ggml-cuda/lightning-indexer.cuh" #include "ggml-cuda/mmf.cuh" #include "ggml-cuda/mmq.cuh" #include "ggml-cuda/mmvf.cuh" @@ -3082,6 +3083,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg case GGML_OP_ARGSORT: ggml_cuda_op_argsort(ctx, dst); break; + case GGML_OP_LIGHTNING_INDEXER: + ggml_cuda_op_lightning_indexer(ctx, dst); + break; case GGML_OP_FLASH_ATTN_EXT: ggml_cuda_flash_attn_ext(ctx, dst); break; @@ -5431,6 +5435,8 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g #else return true; #endif + case GGML_OP_LIGHTNING_INDEXER: + return op->src[1]->type == GGML_TYPE_F16 || op->src[1]->type == GGML_TYPE_F32; case GGML_OP_SUM_ROWS: case GGML_OP_MEAN: case GGML_OP_GROUP_NORM: diff --git a/ggml/src/ggml-cuda/lightning-indexer.cu b/ggml/src/ggml-cuda/lightning-indexer.cu new file mode 100644 index 000000000000..dca4dd739072 --- /dev/null +++ b/ggml/src/ggml-cuda/lightning-indexer.cu @@ -0,0 +1,148 @@ +#include "convert.cuh" +#include "lightning-indexer.cuh" + +// DeepSeek V3.2/V4 lightning indexer, fused (see ggml_compute_forward_lightning_indexer): +// dst[i_kv, i_batch, i_stream] = scale_heads * +// sum_h relu(scale_embd * dot_e(q[e,h,i_batch,i_stream], k[e,i_kv,i_stream])) * w[h,i_batch,i_stream] +// Fusing the reduction over heads avoids materializing the [n_kv, n_head, n_batch] +// score tensor that the unfused graph produces. +// +// One warp per i_kv; the dot product over n_embd is split across the warp lanes and +// finished with a shuffle reduction. q and w for the block's (i_batch, i_stream) are +// staged in shared memory once and reused by every i_kv in the block; the k row is +// held in registers. + +#define LID_MAX_EMB_PER_LANE 8 // supports n_embd up to 32*8 = 256 + +template +static __global__ void lightning_indexer_kernel( + const char * __restrict__ q, + const char * __restrict__ k, + const char * __restrict__ w, + char * __restrict__ dst, + const int n_embd, + const int n_head, + const int n_kv, + const int64_t nb_q_head, const int64_t nb_q_batch, const int64_t nb_q_stream, + const int64_t nb_k_kv, const int64_t nb_k_stream, + const int64_t nb_w_batch, const int64_t nb_w_stream, + const int64_t nb_dst_batch, const int64_t nb_dst_stream, + const float scale_embd, + const float scale_heads) { + extern __shared__ float smem[]; + float * sq = smem; // [n_head*n_embd] q rows for this (i_batch, i_stream) + float * sw = smem + n_head*n_embd; // [n_head] indexer weights + + const int i_batch = blockIdx.y; + const int i_stream = blockIdx.z; + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + const int n_warp = blockDim.x >> 5; + + const char * q_base = q + (int64_t) i_batch*nb_q_batch + (int64_t) i_stream*nb_q_stream; + for (int idx = threadIdx.x; idx < n_head*n_embd; idx += blockDim.x) { + const int h = idx / n_embd; + const int e = idx - h*n_embd; + sq[idx] = *(const float *)(q_base + (int64_t) h*nb_q_head + (int64_t) e*sizeof(float)); + } + const char * w_base = w + (int64_t) i_batch*nb_w_batch + (int64_t) i_stream*nb_w_stream; + for (int h = threadIdx.x; h < n_head; h += blockDim.x) { + sw[h] = *(const float *)(w_base + (int64_t) h*sizeof(float)); + } + __syncthreads(); + + const int i_kv = blockIdx.x*n_warp + warp; + if (i_kv >= n_kv) { + return; + } + + const kv_t * k_row = (const kv_t *)(k + (int64_t) i_kv*nb_k_kv + (int64_t) i_stream*nb_k_stream); + float kreg[LID_MAX_EMB_PER_LANE]; +#pragma unroll + for (int j = 0; j < LID_MAX_EMB_PER_LANE; ++j) { + const int e = lane + j*32; + kreg[j] = e < n_embd ? ggml_cuda_cast(k_row[e]) : 0.0f; + } + + float score = 0.0f; // meaningful on lane 0 + for (int h = 0; h < n_head; ++h) { + const float * sqh = sq + h*n_embd; + float part = 0.0f; +#pragma unroll + for (int j = 0; j < LID_MAX_EMB_PER_LANE; ++j) { + const int e = lane + j*32; + if (e < n_embd) { + part += sqh[e]*kreg[j]; + } + } + part += __shfl_down_sync(0xffffffff, part, 16); + part += __shfl_down_sync(0xffffffff, part, 8); + part += __shfl_down_sync(0xffffffff, part, 4); + part += __shfl_down_sync(0xffffffff, part, 2); + part += __shfl_down_sync(0xffffffff, part, 1); + if (lane == 0) { + score += fmaxf(part*scale_embd, 0.0f)*sw[h]; + } + } + + if (lane == 0) { + float * dst_row = (float *)(dst + (int64_t) i_batch*nb_dst_batch + (int64_t) i_stream*nb_dst_stream); + dst_row[i_kv] = score*scale_heads; + } +} + +void ggml_cuda_op_lightning_indexer(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * q = dst->src[0]; + const ggml_tensor * k = dst->src[1]; + const ggml_tensor * w = dst->src[2]; + + GGML_ASSERT(q->type == GGML_TYPE_F32); + GGML_ASSERT(w->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(q->nb[0] == sizeof(float)); + GGML_ASSERT(w->nb[0] == sizeof(float)); + GGML_ASSERT(dst->nb[0] == sizeof(float)); + GGML_ASSERT(k->nb[0] == ggml_type_size(k->type)); + + const int n_embd = q->ne[0]; + const int n_head = q->ne[1]; + const int n_batch = q->ne[2]; + const int n_stream = q->ne[3]; + const int n_kv = k->ne[2]; + + GGML_ASSERT(n_embd <= 32*LID_MAX_EMB_PER_LANE); + + const float scale_embd = ggml_get_op_params_f32(dst, 0); + const float scale_heads = ggml_get_op_params_f32(dst, 1); + + cudaStream_t stream = ctx.stream(); + + const int n_warp = 8; + const int block = n_warp*32; + const size_t smem = ((size_t) n_head*n_embd + n_head)*sizeof(float); + GGML_ASSERT(smem <= 48*1024); // default per-block shared limit; raise via cudaFuncSetAttribute if exceeded + + const dim3 grid((n_kv + n_warp - 1)/n_warp, n_batch, n_stream); + + const char * q_d = (const char *) q->data; + const char * k_d = (const char *) k->data; + const char * w_d = (const char *) w->data; + char * d_d = (char *) dst->data; + + switch (k->type) { + case GGML_TYPE_F16: + lightning_indexer_kernel<<>>( + q_d, k_d, w_d, d_d, n_embd, n_head, n_kv, + q->nb[1], q->nb[2], q->nb[3], k->nb[2], k->nb[3], + w->nb[1], w->nb[3], dst->nb[1], dst->nb[3], scale_embd, scale_heads); + break; + case GGML_TYPE_F32: + lightning_indexer_kernel<<>>( + q_d, k_d, w_d, d_d, n_embd, n_head, n_kv, + q->nb[1], q->nb[2], q->nb[3], k->nb[2], k->nb[3], + w->nb[1], w->nb[3], dst->nb[1], dst->nb[3], scale_embd, scale_heads); + break; + default: + GGML_ABORT("lightning_indexer: unsupported K type\n"); + } +} diff --git a/ggml/src/ggml-cuda/lightning-indexer.cuh b/ggml/src/ggml-cuda/lightning-indexer.cuh new file mode 100644 index 000000000000..31fcc7d5ae0a --- /dev/null +++ b/ggml/src/ggml-cuda/lightning-indexer.cuh @@ -0,0 +1,3 @@ +#include "common.cuh" + +void ggml_cuda_op_lightning_indexer(ggml_backend_cuda_context & ctx, ggml_tensor * dst);