ggml : add GGML_OP_LIGHTNING_INDEXER that implements DeepSeek V3.2/V4 lightning indexer (llama/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 Georgi Gerganov
parent 4bc66b579f
commit 19c96ad196
No known key found for this signature in database
GPG Key ID: 449E073F9DC10735
6 changed files with 157 additions and 4 deletions

View File

@ -8,10 +8,10 @@ extern "C" {
#define RPC_PROTO_MAJOR_VERSION 4
#define RPC_PROTO_MINOR_VERSION 0
#define RPC_PROTO_PATCH_VERSION 1
#define RPC_PROTO_PATCH_VERSION 2
#ifdef __cplusplus
static_assert(GGML_OP_COUNT == 97, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
static_assert(GGML_OP_COUNT == 98, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
#endif
#define GGML_RPC_MAX_SERVERS 16

View File

@ -570,6 +570,7 @@ extern "C" {
GGML_OP_RWKV_WKV7,
GGML_OP_SOLVE_TRI,
GGML_OP_GATED_DELTA_NET,
GGML_OP_LIGHTNING_INDEXER,
GGML_OP_UNARY,
@ -2575,6 +2576,24 @@ extern "C" {
struct ggml_tensor * state,
int64_t K);
// DSA lightning indexer
//
// q: [n_embd_idx, n_head_idx, n_batch, ne3 ]
// k: [n_embd_idx, 1, n_kv, ne3 ]
// weights: [n_head_idx, n_batch, 1, ne3 ] !! prescaled !!
// mask: [n_kv, n_batch, 1, ne33] !! f16 !!
// res: [n_kv, n_batch, 1, ne3 ]
//
// broadcast:
// ne3 % ne33 == 0
//
GGML_API struct ggml_tensor * ggml_lightning_indexer(
struct ggml_context * ctx,
struct ggml_tensor * q,
struct ggml_tensor * k,
struct ggml_tensor * weights,
struct ggml_tensor * mask);
// custom operators
typedef void (*ggml_custom1_op_t)(struct ggml_tensor * dst , const struct ggml_tensor * a, int ith, int nth, void * userdata);

View File

@ -2060,6 +2060,10 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm
{
ggml_compute_forward_gated_delta_net(params, tensor);
} break;
case GGML_OP_LIGHTNING_INDEXER:
{
ggml_compute_forward_lightning_indexer(params, tensor);
} break;
case GGML_OP_MAP_CUSTOM1:
{
ggml_compute_forward_map_custom1(params, tensor);
@ -2380,6 +2384,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
case GGML_OP_FLASH_ATTN_BACK:
case GGML_OP_SSM_CONV:
case GGML_OP_SSM_SCAN:
case GGML_OP_LIGHTNING_INDEXER:
{
n_tasks = n_threads;
} break;
@ -2965,6 +2970,12 @@ struct ggml_cplan ggml_graph_plan(
{
GGML_ABORT("fatal error");
}
case GGML_OP_LIGHTNING_INDEXER:
{
// temp buffer for dequantizing lightning indexer keys
const int64_t ne10 = node->src[1]->ne[0];
cur += sizeof(float)*ne10*n_tasks;
} break;
default:
break;
}

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]);
}
}
}
}

View File

@ -105,6 +105,7 @@ void ggml_compute_forward_rwkv_wkv7(const struct ggml_compute_params * params, s
void ggml_compute_forward_solve_tri(const struct ggml_compute_params * params, struct ggml_tensor * dst);
void ggml_compute_forward_gla(const struct ggml_compute_params * params, struct ggml_tensor * dst);
void ggml_compute_forward_gated_delta_net(const struct ggml_compute_params * params, struct ggml_tensor * dst);
void ggml_compute_forward_lightning_indexer(const struct ggml_compute_params * params, struct ggml_tensor * dst);
void ggml_compute_forward_map_custom1(const struct ggml_compute_params * params, struct ggml_tensor * dst);
void ggml_compute_forward_map_custom2(const struct ggml_compute_params * params, struct ggml_tensor * dst);
void ggml_compute_forward_map_custom3(const struct ggml_compute_params * params, struct ggml_tensor * dst);

View File

@ -1079,6 +1079,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = {
"RWKV_WKV7",
"SOLVE_TRI",
"GATED_DELTA_NET",
"LIGHTNING_INDEXER",
"UNARY",
@ -1096,7 +1097,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = {
"GLU",
};
static_assert(GGML_OP_COUNT == 97, "GGML_OP_COUNT != 97");
static_assert(GGML_OP_COUNT == 98, "GGML_OP_COUNT != 98");
static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
"none",
@ -1190,6 +1191,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
"rwkv_wkv7(r, w, k, v, a, b, s)",
"A X = B, A triangular, solve X",
"gated_delta_net(q, k, v, g, beta, s)",
"lightning_indexer(q, k, weights, mask)",
"unary(x)",
@ -1207,7 +1209,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
"glu(x)",
};
static_assert(GGML_OP_COUNT == 97, "GGML_OP_COUNT != 97");
static_assert(GGML_OP_COUNT == 98, "GGML_OP_COUNT != 98");
static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2");
@ -6287,6 +6289,42 @@ struct ggml_tensor * ggml_gated_delta_net(
return result;
}
// ggml_lightning_indexer
struct ggml_tensor * ggml_lightning_indexer(
struct ggml_context * ctx,
struct ggml_tensor * q,
struct ggml_tensor * k,
struct ggml_tensor * weights,
struct ggml_tensor * mask) {
GGML_ASSERT( q->type == GGML_TYPE_F32);
GGML_ASSERT( weights->type == GGML_TYPE_F32);
GGML_ASSERT( mask->type == GGML_TYPE_F16);
GGML_ASSERT( q->ne[0] == k->ne[0]);
GGML_ASSERT( mask->ne[0] == k->ne[2]);
GGML_ASSERT( q->ne[1] == weights->ne[0]);
GGML_ASSERT( k->ne[1] == 1);
GGML_ASSERT( mask->ne[1] == q->ne[2]);
GGML_ASSERT( q->ne[2] == weights->ne[1]);
GGML_ASSERT(weights->ne[2] == 1);
GGML_ASSERT( mask->ne[2] == 1);
GGML_ASSERT( q->ne[3] == k->ne[3]);
GGML_ASSERT( k->ne[3] == weights->ne[3]);
GGML_ASSERT(weights->ne[3] % mask->ne[3] == 0);
int64_t ne[4] = { k->ne[2], q->ne[2], 1, q->ne[3] };
struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne);
result->op = GGML_OP_LIGHTNING_INDEXER;
result->src[0] = q;
result->src[1] = k;
result->src[2] = weights;
result->src[3] = mask;
return result;
}
////////////////////////////////////////////////////////////////////////////////
struct ggml_hash_set ggml_hash_set_new(size_t size) {