ggml : add GGML_OP_LIGHTNING_INDEXER that implements DeepSeek V3.2/V4 lightning indexer (#24231)

* ggml : add GGML_OP_LIGHTNING_INDEXER that implements DeepSeek V3.2/V4 lightning indexer

* ggml : remove scale parameters from lightning indexer OP, add f16 mask parameter

* tests : add GGML_OP_LIGHTNING_INDEXER tests

* ggml : bump RPC version

* chore : check if lightning indexer input tensors are not transposed

* tests : count flops instead of bandwidth in lightning indexer test

* chore : add missing const

* chore : whitespace

* ggml : renamed variables in CPU lightning indexer implementation

* ggml : fix lightning indexer mask broadcasting

* tests : tests for lightning indexer mask broadcasting

* chore : whitespace

* llama : use GGML_OP_LIGHTNING_INDEXER in DeepSeek V3.2 and DeepSeek V4 models

---------

Co-authored-by: Stanisław Szymczyk <sszymczy@gmail.com>
This commit is contained in:
fairydreaming
2026-07-11 11:39:07 +02:00
committed by GitHub
co-authored by Stanisław Szymczyk
parent 76f2798059
commit 00f5442cc4
13 changed files with 324 additions and 52 deletions
+84
View File
@@ -11568,3 +11568,87 @@ void ggml_compute_forward_fwht(const ggml_compute_params * params, ggml_tensor *
}
}
}
// ggml_compute_forward_lightning_indexer
void ggml_compute_forward_lightning_indexer(
const ggml_compute_params * params,
ggml_tensor * dst) {
const ggml_tensor * q = dst->src[0];
const ggml_tensor * k = dst->src[1];
const ggml_tensor * w = dst->src[2]; // weights
const ggml_tensor * m = dst->src[3]; // mask
GGML_ASSERT(dst->type == GGML_TYPE_F32);
GGML_ASSERT( q->type == GGML_TYPE_F32);
GGML_ASSERT( w->type == GGML_TYPE_F32);
GGML_ASSERT( m->type == GGML_TYPE_F16);
GGML_TENSOR_LOCALS(int64_t, neq, q, ne)
GGML_TENSOR_LOCALS(size_t, nbq, q, nb)
GGML_TENSOR_LOCALS(int64_t, nek, k, ne)
GGML_TENSOR_LOCALS(size_t, nbk, k, nb)
GGML_TENSOR_LOCALS(int64_t, new, w, ne)
GGML_TENSOR_LOCALS(size_t, nbw, w, nb)
GGML_TENSOR_LOCALS(int64_t, nem, m, ne)
GGML_TENSOR_LOCALS(size_t, nbm, m, nb)
GGML_TENSOR_LOCALS(int64_t, ne, dst, ne)
GGML_TENSOR_LOCALS(size_t, nb, dst, nb)
GGML_ASSERT( nb0 == ggml_type_size(dst->type));
GGML_ASSERT(nbq0 == ggml_type_size( q->type));
GGML_ASSERT(nbk0 == ggml_type_size( k->type));
GGML_ASSERT(nbw0 == ggml_type_size( w->type));
GGML_ASSERT(nbm0 == ggml_type_size( m->type));
const int n_embd = q->ne[0];
const int n_head = q->ne[1];
const int n_tokens = q->ne[2];
const int n_stream = q->ne[3];
const int n_kv = k->ne[2];
ggml_to_float_t const k_to_float = ggml_get_type_traits(k->type)->to_float;
GGML_ASSERT((k->type == GGML_TYPE_F32 || k_to_float) && "lightning indexer: unsupported K-type");
const int nr = n_kv;
const int ith = params->ith;
const int nth = params->nth;
// (temporary) buffer for K converted to float
float * k_row_f32 = (float *) params->wdata + ith*(1*n_embd + CACHE_LINE_SIZE_F32);
// rows per thread
const int dr = (nr + nth - 1)/nth;
// row range for this thread
const int ir0 = dr*ith;
const int ir1 = MIN(ir0 + dr, nr);
for (int s = 0; s < n_stream; ++s) {
for (int t = 0; t < n_tokens; ++t) {
const float * w_row = (float *) ((char *) w->data + t*nbw1 + s*nbw3);
const ggml_fp16_t * m_row = (ggml_fp16_t *) ((char *) m->data + t*nbm1 + (s%nem3)*nbm3);
float * dst_row = (float *) ((char *) dst->data + t*nb1 + s*nb3 );
for (int ik = ir0; ik < ir1; ++ik) {
char * k_row = (char *) k->data + ik*nbk2 + s*nbk3;
if (k_to_float) {
k_to_float(k_row, k_row_f32, n_embd);
} else {
k_row_f32 = (float *) k_row;
}
float score = 0.0f;
for (int h = 0; h < n_head; ++h) {
// dot product of q and k for head h
float qk = 0.0f;
const float * q_row = (float *) ((char *) q->data + h*nbq1 + t*nbq2 + s*nbq3);
ggml_vec_dot_f32(n_embd, &qk, 0, q_row, 0, k_row_f32, 0, 1);
// ReLU and weights (prescaled)
score += MAX(qk, 0.0f) * w_row[h];
}
// apply mask
dst_row[ik] = score + GGML_CPU_FP16_TO_FP32(m_row[ik]);
}
}
}
}