From 57a667e99e605593bd2992bbd303ebe3abf3fd53 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stanis=C5=82aw=20Szymczyk?= Date: Sat, 6 Jun 2026 11:13:17 +0200 Subject: [PATCH 01/13] ggml : add GGML_OP_LIGHTNING_INDEXER that implements DeepSeek V3.2/V4 lightning indexer --- ggml/include/ggml-rpc.h | 4 +- ggml/include/ggml.h | 9 +++++ ggml/src/ggml-cpu/ggml-cpu.c | 11 ++++++ ggml/src/ggml-cpu/ops.cpp | 73 ++++++++++++++++++++++++++++++++++++ ggml/src/ggml-cpu/ops.h | 1 + ggml/src/ggml.c | 40 +++++++++++++++++++- 6 files changed, 134 insertions(+), 4 deletions(-) diff --git a/ggml/include/ggml-rpc.h b/ggml/include/ggml-rpc.h index 6fcf5a43393..5ad121ae57f 100644 --- a/ggml/include/ggml-rpc.h +++ b/ggml/include/ggml-rpc.h @@ -8,10 +8,10 @@ extern "C" { #define RPC_PROTO_MAJOR_VERSION 4 #define RPC_PROTO_MINOR_VERSION 0 -#define RPC_PROTO_PATCH_VERSION 0 +#define RPC_PROTO_PATCH_VERSION 1 #ifdef __cplusplus -static_assert(GGML_OP_COUNT == 96, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION"); +static_assert(GGML_OP_COUNT == 97, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION"); #endif #define GGML_RPC_MAX_SERVERS 16 diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index f6725265504..dc737362829 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -567,6 +567,7 @@ extern "C" { GGML_OP_RWKV_WKV7, GGML_OP_SOLVE_TRI, GGML_OP_GATED_DELTA_NET, + GGML_OP_LIGHTNING_INDEXER, GGML_OP_UNARY, @@ -2555,6 +2556,14 @@ extern "C" { struct ggml_tensor * beta, struct ggml_tensor * state); + GGML_API struct ggml_tensor * ggml_lightning_indexer( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * weights, + float scale_embd, + float scale_heads); + // custom operators typedef void (*ggml_custom1_op_t)(struct ggml_tensor * dst , const struct ggml_tensor * a, int ith, int nth, void * userdata); diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index cd5c61a8187..023dfbaeb87 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -2047,6 +2047,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); @@ -2366,6 +2370,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; @@ -2951,6 +2956,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; } diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 3a1912ae91b..c35fae10133 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -11414,3 +11414,76 @@ 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 * src0 = dst->src[0]; // q + const ggml_tensor * src1 = dst->src[1]; // k + const ggml_tensor * src2 = dst->src[2]; // weights + + const float scale_embd = ggml_get_op_params_f32(dst, 0); + const float scale_heads = ggml_get_op_params_f32(dst, 1); + + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT(src2->type == GGML_TYPE_F32); + + GGML_TENSOR_TERNARY_OP_LOCALS + + GGML_ASSERT( nb0 == sizeof(float)); + GGML_ASSERT(nb00 == sizeof(float)); + + int n_embd = src0->ne[0]; + int n_head = src0->ne[1]; + int n_batch = src0->ne[2]; + int n_stream = src0->ne[3]; + int n_kv = src1->ne[2]; + + ggml_to_float_t const k_to_float = ggml_get_type_traits(src1->type)->to_float; + GGML_ASSERT((src1->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 * src1_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 i_stream = 0; i_stream < n_stream; ++i_stream) { + for (int i_batch = 0; i_batch < n_batch; ++i_batch) { + for (int i_kv = ir0; i_kv < ir1; ++i_kv) { + char * src1_row = (char *) src1->data + i_kv*nb12 + i_stream*nb13; + if (k_to_float) { + k_to_float(src1_row, src1_row_f32, n_embd); + } else { + src1_row_f32 = (float *) src1_row; + } + float * src2_row = (float *) ((char *) src2->data + i_batch*nb21 + i_stream*nb23); + float * dst_row = (float *) ((char *) dst->data + i_batch*nb1 + i_stream*nb3); + float score = 0.0f; + for (int i_head = 0; i_head < n_head; ++i_head) { + // dot product of q and k for head i_head + float qk = 0.0f; + float * src0_row = (float *) ((char *) src0->data + i_head*nb01 + i_batch*nb02 + i_stream*nb03); + ggml_vec_dot_f32(n_embd, &qk, 0, src0_row, 0, src1_row_f32, 0, 1); + qk *= scale_embd; + // ReLU and weights + score += MAX(qk, 0.0f) * src2_row[i_head]; + } + score *= scale_heads; + dst_row[i_kv] = score; + } + } + } +} diff --git a/ggml/src/ggml-cpu/ops.h b/ggml/src/ggml-cpu/ops.h index 7398e561894..3648f574b42 100644 --- a/ggml/src/ggml-cpu/ops.h +++ b/ggml/src/ggml-cpu/ops.h @@ -104,6 +104,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); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 8815c67d8bc..14df3a55476 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -1063,6 +1063,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "RWKV_WKV7", "SOLVE_TRI", "GATED_DELTA_NET", + "LIGHTNING_INDEXER", "UNARY", @@ -1080,7 +1081,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "GLU", }; -static_assert(GGML_OP_COUNT == 96, "GGML_OP_COUNT != 96"); +static_assert(GGML_OP_COUNT == 97, "GGML_OP_COUNT != 97"); static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "none", @@ -1173,6 +1174,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, scale_embd, scale_heads)", "unary(x)", @@ -1190,7 +1192,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "glu(x)", }; -static_assert(GGML_OP_COUNT == 96, "GGML_OP_COUNT != 96"); +static_assert(GGML_OP_COUNT == 97, "GGML_OP_COUNT != 97"); static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); @@ -6230,6 +6232,40 @@ 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, + float scale_embd, + float scale_heads) { + + GGML_ASSERT(q->type == GGML_TYPE_F32); + GGML_ASSERT(weights->type == GGML_TYPE_F32); + GGML_ASSERT(q->ne[0] == k->ne[0]); + GGML_ASSERT(q->ne[1] == weights->ne[0]); + GGML_ASSERT(k->ne[1] == 1); + GGML_ASSERT(q->ne[2] == weights->ne[1]); + GGML_ASSERT(weights->ne[2] == 1); + GGML_ASSERT(q->ne[3] == k->ne[3]); + GGML_ASSERT(k->ne[3] == weights->ne[3]); + + 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); + + ggml_set_op_params_f32(result, 0, scale_embd); + ggml_set_op_params_f32(result, 1, scale_heads); + + result->op = GGML_OP_LIGHTNING_INDEXER; + result->src[0] = q; + result->src[1] = k; + result->src[2] = weights; + + return result; +} + //////////////////////////////////////////////////////////////////////////////// struct ggml_hash_set ggml_hash_set_new(size_t size) { From cf0167fbf38948c81650d72b7fd6ec78a3cb40f5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stanis=C5=82aw=20Szymczyk?= Date: Wed, 8 Jul 2026 13:22:33 +0200 Subject: [PATCH 02/13] ggml : remove scale parameters from lightning indexer OP, add f16 mask parameter --- ggml/include/ggml.h | 14 ++++++++++++-- ggml/src/ggml-cpu/ops.cpp | 19 ++++++++++--------- ggml/src/ggml.c | 14 ++++++++------ 3 files changed, 30 insertions(+), 17 deletions(-) diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 403bdfae34d..b2859ebe7df 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -2576,13 +2576,23 @@ 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, - float scale_embd, - float scale_heads); + struct ggml_tensor * mask); // custom operators diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 148f6c6d78c..b3288ce8ce4 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -11566,15 +11566,16 @@ void ggml_compute_forward_lightning_indexer( const ggml_tensor * src0 = dst->src[0]; // q const ggml_tensor * src1 = dst->src[1]; // k const ggml_tensor * src2 = dst->src[2]; // weights - - const float scale_embd = ggml_get_op_params_f32(dst, 0); - const float scale_heads = ggml_get_op_params_f32(dst, 1); + const ggml_tensor * src3 = dst->src[3]; // mask GGML_ASSERT(dst->type == GGML_TYPE_F32); GGML_ASSERT(src0->type == GGML_TYPE_F32); GGML_ASSERT(src2->type == GGML_TYPE_F32); + GGML_ASSERT(src3->type == GGML_TYPE_F16); GGML_TENSOR_TERNARY_OP_LOCALS + GGML_TENSOR_LOCALS(int64_t, ne3, src3, ne) + GGML_TENSOR_LOCALS(size_t, nb3, src3, nb) GGML_ASSERT( nb0 == sizeof(float)); GGML_ASSERT(nb00 == sizeof(float)); @@ -11604,6 +11605,9 @@ void ggml_compute_forward_lightning_indexer( for (int i_stream = 0; i_stream < n_stream; ++i_stream) { for (int i_batch = 0; i_batch < n_batch; ++i_batch) { + const float * src2_row = (float *) ((char *) src2->data + i_batch*nb21 + i_stream*nb23); + const ggml_fp16_t * src3_row = (ggml_fp16_t *) ((char *) src3->data + i_batch*nb31 + (i_stream%ne33)*nb33); + float * dst_row = (float *) ((char *) dst->data + i_batch*nb1 + i_stream*nb3); for (int i_kv = ir0; i_kv < ir1; ++i_kv) { char * src1_row = (char *) src1->data + i_kv*nb12 + i_stream*nb13; if (k_to_float) { @@ -11611,20 +11615,17 @@ void ggml_compute_forward_lightning_indexer( } else { src1_row_f32 = (float *) src1_row; } - float * src2_row = (float *) ((char *) src2->data + i_batch*nb21 + i_stream*nb23); - float * dst_row = (float *) ((char *) dst->data + i_batch*nb1 + i_stream*nb3); float score = 0.0f; for (int i_head = 0; i_head < n_head; ++i_head) { // dot product of q and k for head i_head float qk = 0.0f; float * src0_row = (float *) ((char *) src0->data + i_head*nb01 + i_batch*nb02 + i_stream*nb03); ggml_vec_dot_f32(n_embd, &qk, 0, src0_row, 0, src1_row_f32, 0, 1); - qk *= scale_embd; - // ReLU and weights + // ReLU and weights (prescaled) score += MAX(qk, 0.0f) * src2_row[i_head]; } - score *= scale_heads; - dst_row[i_kv] = score; + // apply mask + dst_row[i_kv] = score + GGML_CPU_FP16_TO_FP32(src3_row[i_kv]); } } } diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index e1ce9362c4d..8334440fc87 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -1191,7 +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, scale_embd, scale_heads)", + "lightning_indexer(q, k, weights, mask)", "unary(x)", @@ -6296,29 +6296,31 @@ struct ggml_tensor * ggml_lightning_indexer( struct ggml_tensor * q, struct ggml_tensor * k, struct ggml_tensor * weights, - float scale_embd, - float scale_heads) { + 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]); 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); - ggml_set_op_params_f32(result, 0, scale_embd); - ggml_set_op_params_f32(result, 1, scale_heads); - result->op = GGML_OP_LIGHTNING_INDEXER; result->src[0] = q; result->src[1] = k; result->src[2] = weights; + result->src[3] = mask; return result; } From bb2e155050061ccb5b9c78bb4c51ee25920cfec2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stanis=C5=82aw=20Szymczyk?= Date: Wed, 8 Jul 2026 15:50:11 +0200 Subject: [PATCH 03/13] tests : add GGML_OP_LIGHTNING_INDEXER tests --- tests/test-backend-ops.cpp | 81 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 81 insertions(+) diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 0c0aac2bba7..e2401ea4221 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -7108,6 +7108,61 @@ struct test_diag : public test_case { } }; +// GGML_OP_LIGHTNING_INDEXER +struct test_lightning_indexer : public test_case { + const int64_t hsk; // indexer K head size + const int64_t nh; // num indexer heads + const int64_t kv; // kv size + const int64_t nb; // batch size + const int64_t ns; // num streams + + const ggml_type type_K; + + std::string vars() override { + return VARS_TO_STR6(hsk, nh, kv, nb, ns, type_K); + } + + double max_nmse_err() override { + return 1e-6; + } + + test_lightning_indexer(int64_t hsk = 128, int64_t nh = 64, int64_t kv = 256, int64_t nb = 128, int64_t ns = 1, ggml_type type_K = GGML_TYPE_F16) + : hsk(hsk), nh(nh), kv(kv), nb(nb), ns(ns), type_K(type_K) {} + + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, hsk, nh, nb, ns); + ggml_set_param(q); + ggml_set_name(q, "q"); + + ggml_tensor * k = ggml_new_tensor_4d(ctx, type_K, hsk, 1, kv, ns); + ggml_set_param(k); + ggml_set_name(k, "k"); + + ggml_tensor * w = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, nh, nb, 1, ns); + ggml_set_param(w); + ggml_set_name(w, "w"); + + ggml_tensor * m = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, kv, nb, 1, ns); + ggml_set_param(m); + ggml_set_name(m, "m"); + + ggml_tensor * out = ggml_lightning_indexer(ctx, q, k, w, m); + ggml_set_name(out, "out"); + + return out; + } + + void initialize_tensors(ggml_context * ctx) override { + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { + if (strcmp(t->name, "m") == 0) { + init_tensor_kq_mask(t); + } else { + init_tensor_uniform(t); + } + } + } +}; + // Deserializable generic test case struct input_tensor { ggml_type type; @@ -9399,6 +9454,19 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_falcon(2)); #endif + // lightning_indexer + for (int kv : { 256 }) { + for (int bs : { 1, 512 }) { + for (int nh : { 32, 64 }) { + for (int ns : { 1, 4 }) { + for (ggml_type type_K : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_Q8_0, GGML_TYPE_Q5_1, GGML_TYPE_Q5_0, GGML_TYPE_Q4_1, GGML_TYPE_Q4_0, GGML_TYPE_IQ4_NL}) { + test_cases.emplace_back(new test_lightning_indexer(128, nh, kv, bs, ns, type_K)); + } + } + } + } + } + return test_cases; } #ifdef _MSC_VER @@ -9724,6 +9792,19 @@ static std::vector> make_test_cases_perf() { test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 128, 1024, 1)); // 4h PP-1024 test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 32, 128, 64, 1, 1, false, true)); // KDA PP-64 + // lightning_indexer + for (int kv : { 256, 4096, 65536 }) { + for (int bs : { 1, 512, 2048 }) { + for (int nh : { 32, 64 }) { + for (int ns : { 1, 4 }) { + for (ggml_type type_K : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_Q8_0, GGML_TYPE_Q5_1, GGML_TYPE_Q5_0, GGML_TYPE_Q4_1, GGML_TYPE_Q4_0, GGML_TYPE_IQ4_NL}) { + test_cases.emplace_back(new test_lightning_indexer(128, nh, kv, bs, ns, type_K)); + } + } + } + } + } + return test_cases; } From c7a29777592f6f47c8f992e8c73df3f4b425cb93 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stanis=C5=82aw=20Szymczyk?= Date: Wed, 8 Jul 2026 16:31:45 +0200 Subject: [PATCH 04/13] ggml : bump RPC version --- ggml/include/ggml-rpc.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ggml/include/ggml-rpc.h b/ggml/include/ggml-rpc.h index 5ad121ae57f..efa5420a698 100644 --- a/ggml/include/ggml-rpc.h +++ b/ggml/include/ggml-rpc.h @@ -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 From 115b777d27c0884a206eb0b38557ae7e8540ceb8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stanis=C5=82aw=20Szymczyk?= Date: Wed, 8 Jul 2026 17:28:13 +0200 Subject: [PATCH 05/13] chore : check if lightning indexer input tensors are not transposed --- ggml/src/ggml-cpu/ops.cpp | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index b3288ce8ce4..95154c96158 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -11577,8 +11577,11 @@ void ggml_compute_forward_lightning_indexer( GGML_TENSOR_LOCALS(int64_t, ne3, src3, ne) GGML_TENSOR_LOCALS(size_t, nb3, src3, nb) - GGML_ASSERT( nb0 == sizeof(float)); - GGML_ASSERT(nb00 == sizeof(float)); + GGML_ASSERT( nb0 == ggml_type_size( dst->type)); + GGML_ASSERT(nb00 == ggml_type_size(src0->type)); + GGML_ASSERT(nb10 == ggml_type_size(src1->type)); + GGML_ASSERT(nb20 == ggml_type_size(src2->type)); + GGML_ASSERT(nb30 == ggml_type_size(src3->type)); int n_embd = src0->ne[0]; int n_head = src0->ne[1]; From 4da159d05e6fc09072a1f83decefa0fc3989a97e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stanis=C5=82aw=20Szymczyk?= Date: Wed, 8 Jul 2026 17:46:42 +0200 Subject: [PATCH 06/13] tests : count flops instead of bandwidth in lightning indexer test --- tests/test-backend-ops.cpp | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index e2401ea4221..200d2f1e0b4 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -7126,6 +7126,11 @@ struct test_lightning_indexer : public test_case { return 1e-6; } + uint64_t op_flops(ggml_tensor * t) override { + GGML_UNUSED(t); + return ((2 * hsk + 2) * nh + 1) * kv * nb * ns; + } + test_lightning_indexer(int64_t hsk = 128, int64_t nh = 64, int64_t kv = 256, int64_t nb = 128, int64_t ns = 1, ggml_type type_K = GGML_TYPE_F16) : hsk(hsk), nh(nh), kv(kv), nb(nb), ns(ns), type_K(type_K) {} From 755d9cadae9f0738f4d14091853acd7636ad5471 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stanis=C5=82aw=20Szymczyk?= Date: Wed, 8 Jul 2026 18:29:35 +0200 Subject: [PATCH 07/13] chore : add missing const --- ggml/src/ggml-cpu/ops.cpp | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 95154c96158..a7e6621751b 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -11583,11 +11583,11 @@ void ggml_compute_forward_lightning_indexer( GGML_ASSERT(nb20 == ggml_type_size(src2->type)); GGML_ASSERT(nb30 == ggml_type_size(src3->type)); - int n_embd = src0->ne[0]; - int n_head = src0->ne[1]; - int n_batch = src0->ne[2]; - int n_stream = src0->ne[3]; - int n_kv = src1->ne[2]; + const int n_embd = src0->ne[0]; + const int n_head = src0->ne[1]; + const int n_batch = src0->ne[2]; + const int n_stream = src0->ne[3]; + const int n_kv = src1->ne[2]; ggml_to_float_t const k_to_float = ggml_get_type_traits(src1->type)->to_float; GGML_ASSERT((src1->type == GGML_TYPE_F32 || k_to_float) && "lightning indexer: unsupported K-type"); @@ -11622,7 +11622,7 @@ void ggml_compute_forward_lightning_indexer( for (int i_head = 0; i_head < n_head; ++i_head) { // dot product of q and k for head i_head float qk = 0.0f; - float * src0_row = (float *) ((char *) src0->data + i_head*nb01 + i_batch*nb02 + i_stream*nb03); + const float * src0_row = (float *) ((char *) src0->data + i_head*nb01 + i_batch*nb02 + i_stream*nb03); ggml_vec_dot_f32(n_embd, &qk, 0, src0_row, 0, src1_row_f32, 0, 1); // ReLU and weights (prescaled) score += MAX(qk, 0.0f) * src2_row[i_head]; From 7f59826cc6e732dff4e2845e7f7d074b32e7be1d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stanis=C5=82aw=20Szymczyk?= Date: Wed, 8 Jul 2026 21:41:43 +0200 Subject: [PATCH 08/13] chore : whitespace --- ggml/src/ggml-cpu/ops.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index a7e6621751b..b44eb063b4f 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -11610,7 +11610,7 @@ void ggml_compute_forward_lightning_indexer( for (int i_batch = 0; i_batch < n_batch; ++i_batch) { const float * src2_row = (float *) ((char *) src2->data + i_batch*nb21 + i_stream*nb23); const ggml_fp16_t * src3_row = (ggml_fp16_t *) ((char *) src3->data + i_batch*nb31 + (i_stream%ne33)*nb33); - float * dst_row = (float *) ((char *) dst->data + i_batch*nb1 + i_stream*nb3); + float * dst_row = (float *) ((char *) dst->data + i_batch*nb1 + i_stream*nb3 ); for (int i_kv = ir0; i_kv < ir1; ++i_kv) { char * src1_row = (char *) src1->data + i_kv*nb12 + i_stream*nb13; if (k_to_float) { From 86ab142967db518076e8cedecb168c5e713df157 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stanis=C5=82aw=20Szymczyk?= Date: Thu, 9 Jul 2026 14:35:13 +0200 Subject: [PATCH 09/13] ggml : renamed variables in CPU lightning indexer implementation --- ggml/src/ggml-cpu/ops.cpp | 91 +++++++++++++++++++++------------------ 1 file changed, 49 insertions(+), 42 deletions(-) diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index b44eb063b4f..026038c6306 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -11563,41 +11563,48 @@ void ggml_compute_forward_lightning_indexer( const ggml_compute_params * params, ggml_tensor * dst) { - const ggml_tensor * src0 = dst->src[0]; // q - const ggml_tensor * src1 = dst->src[1]; // k - const ggml_tensor * src2 = dst->src[2]; // weights - const ggml_tensor * src3 = dst->src[3]; // mask + 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(src0->type == GGML_TYPE_F32); - GGML_ASSERT(src2->type == GGML_TYPE_F32); - GGML_ASSERT(src3->type == GGML_TYPE_F16); - - GGML_TENSOR_TERNARY_OP_LOCALS - GGML_TENSOR_LOCALS(int64_t, ne3, src3, ne) - GGML_TENSOR_LOCALS(size_t, nb3, src3, nb) - - GGML_ASSERT( nb0 == ggml_type_size( dst->type)); - GGML_ASSERT(nb00 == ggml_type_size(src0->type)); - GGML_ASSERT(nb10 == ggml_type_size(src1->type)); - GGML_ASSERT(nb20 == ggml_type_size(src2->type)); - GGML_ASSERT(nb30 == ggml_type_size(src3->type)); - - const int n_embd = src0->ne[0]; - const int n_head = src0->ne[1]; - const int n_batch = src0->ne[2]; - const int n_stream = src0->ne[3]; - const int n_kv = src1->ne[2]; - - ggml_to_float_t const k_to_float = ggml_get_type_traits(src1->type)->to_float; - GGML_ASSERT((src1->type == GGML_TYPE_F32 || k_to_float) && "lightning indexer: unsupported K-type"); + 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 * src1_row_f32 = (float *) params->wdata + ith*(1*n_embd + CACHE_LINE_SIZE_F32); + 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; @@ -11606,29 +11613,29 @@ void ggml_compute_forward_lightning_indexer( const int ir0 = dr*ith; const int ir1 = MIN(ir0 + dr, nr); - for (int i_stream = 0; i_stream < n_stream; ++i_stream) { - for (int i_batch = 0; i_batch < n_batch; ++i_batch) { - const float * src2_row = (float *) ((char *) src2->data + i_batch*nb21 + i_stream*nb23); - const ggml_fp16_t * src3_row = (ggml_fp16_t *) ((char *) src3->data + i_batch*nb31 + (i_stream%ne33)*nb33); - float * dst_row = (float *) ((char *) dst->data + i_batch*nb1 + i_stream*nb3 ); - for (int i_kv = ir0; i_kv < ir1; ++i_kv) { - char * src1_row = (char *) src1->data + i_kv*nb12 + i_stream*nb13; + 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(src1_row, src1_row_f32, n_embd); + k_to_float(k_row, k_row_f32, n_embd); } else { - src1_row_f32 = (float *) src1_row; + k_row_f32 = (float *) k_row; } float score = 0.0f; - for (int i_head = 0; i_head < n_head; ++i_head) { - // dot product of q and k for head i_head + for (int h = 0; h < n_head; ++h) { + // dot product of q and k for head h float qk = 0.0f; - const float * src0_row = (float *) ((char *) src0->data + i_head*nb01 + i_batch*nb02 + i_stream*nb03); - ggml_vec_dot_f32(n_embd, &qk, 0, src0_row, 0, src1_row_f32, 0, 1); + 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) * src2_row[i_head]; + score += MAX(qk, 0.0f) * w_row[h]; } // apply mask - dst_row[i_kv] = score + GGML_CPU_FP16_TO_FP32(src3_row[i_kv]); + dst_row[ik] = score + GGML_CPU_FP16_TO_FP32(m_row[ik]); } } } From adfba4dacce72d59e1a766dd6436dec311969a11 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stanis=C5=82aw=20Szymczyk?= Date: Thu, 9 Jul 2026 16:01:13 +0200 Subject: [PATCH 10/13] ggml : fix lightning indexer mask broadcasting --- ggml/src/ggml.c | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 8334440fc87..9d63f579640 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -6311,7 +6311,7 @@ struct ggml_tensor * ggml_lightning_indexer( 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]); + 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); From 84dd68c13fc929f8cf7d626acbdd960093b66aaa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stanis=C5=82aw=20Szymczyk?= Date: Thu, 9 Jul 2026 16:01:42 +0200 Subject: [PATCH 11/13] tests : tests for lightning indexer mask broadcasting --- tests/test-backend-ops.cpp | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 200d2f1e0b4..4513f508786 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -7115,11 +7115,12 @@ struct test_lightning_indexer : public test_case { const int64_t kv; // kv size const int64_t nb; // batch size const int64_t ns; // num streams + const int64_t nm; // ne[3] of mask const ggml_type type_K; std::string vars() override { - return VARS_TO_STR6(hsk, nh, kv, nb, ns, type_K); + return VARS_TO_STR7(hsk, nh, kv, nb, ns, nm, type_K); } double max_nmse_err() override { @@ -7131,8 +7132,8 @@ struct test_lightning_indexer : public test_case { return ((2 * hsk + 2) * nh + 1) * kv * nb * ns; } - test_lightning_indexer(int64_t hsk = 128, int64_t nh = 64, int64_t kv = 256, int64_t nb = 128, int64_t ns = 1, ggml_type type_K = GGML_TYPE_F16) - : hsk(hsk), nh(nh), kv(kv), nb(nb), ns(ns), type_K(type_K) {} + test_lightning_indexer(int64_t hsk = 128, int64_t nh = 64, int64_t kv = 256, int64_t nb = 128, int64_t ns = 1, int64_t nm = 1, ggml_type type_K = GGML_TYPE_F16) + : hsk(hsk), nh(nh), kv(kv), nb(nb), ns(ns), nm(nm), type_K(type_K) {} ggml_tensor * build_graph(ggml_context * ctx) override { ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, hsk, nh, nb, ns); @@ -7147,7 +7148,7 @@ struct test_lightning_indexer : public test_case { ggml_set_param(w); ggml_set_name(w, "w"); - ggml_tensor * m = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, kv, nb, 1, ns); + ggml_tensor * m = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, kv, nb, 1, nm); ggml_set_param(m); ggml_set_name(m, "m"); @@ -9463,9 +9464,9 @@ static std::vector> make_test_cases_eval() { for (int kv : { 256 }) { for (int bs : { 1, 512 }) { for (int nh : { 32, 64 }) { - for (int ns : { 1, 4 }) { + for (auto [ns, nm] : { std::pair{1, 1}, std::pair{4, 4}, std::pair{4, 1} }) { for (ggml_type type_K : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_Q8_0, GGML_TYPE_Q5_1, GGML_TYPE_Q5_0, GGML_TYPE_Q4_1, GGML_TYPE_Q4_0, GGML_TYPE_IQ4_NL}) { - test_cases.emplace_back(new test_lightning_indexer(128, nh, kv, bs, ns, type_K)); + test_cases.emplace_back(new test_lightning_indexer(128, nh, kv, bs, ns, nm, type_K)); } } } @@ -9803,7 +9804,7 @@ static std::vector> make_test_cases_perf() { for (int nh : { 32, 64 }) { for (int ns : { 1, 4 }) { for (ggml_type type_K : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_Q8_0, GGML_TYPE_Q5_1, GGML_TYPE_Q5_0, GGML_TYPE_Q4_1, GGML_TYPE_Q4_0, GGML_TYPE_IQ4_NL}) { - test_cases.emplace_back(new test_lightning_indexer(128, nh, kv, bs, ns, type_K)); + test_cases.emplace_back(new test_lightning_indexer(128, nh, kv, bs, ns, ns, type_K)); } } } From 123137778b4c002a7f1adde067c11581774eca63 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stanis=C5=82aw=20Szymczyk?= Date: Thu, 9 Jul 2026 16:03:07 +0200 Subject: [PATCH 12/13] chore : whitespace --- ggml/src/ggml.c | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 9d63f579640..5d880499169 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -6298,19 +6298,19 @@ struct ggml_tensor * ggml_lightning_indexer( 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( 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( 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] }; From 428831bb7c0d1e5633ef067aa43c33f1afd99a30 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stanis=C5=82aw=20Szymczyk?= Date: Fri, 10 Jul 2026 19:09:34 +0200 Subject: [PATCH 13/13] llama : use GGML_OP_LIGHTNING_INDEXER in DeepSeek V3.2 and DeepSeek V4 models --- src/llama-context.cpp | 15 ++++++++ src/llama-cparams.h | 2 ++ src/llama-graph.cpp | 6 ++-- src/llama-graph.h | 1 + src/models/deepseek32.cpp | 73 +++++++++++++++++++++------------------ src/models/deepseek4.cpp | 37 ++++++++++++-------- 6 files changed, 83 insertions(+), 51 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 5edfc85abfd..3a469bc90bd 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -55,6 +55,12 @@ static const llm_fused_op_probe llm_fused_op_gdn_ch_probe = { /*.n_tokens_per_seq =*/ 16, }; +static const llm_fused_op_probe llm_fused_op_lid_probe = { + /*.op =*/ LLM_FUSED_OP_LIGHTNING_INDEXER, + /*.name =*/ "Lightning Indexer", + /*.n_tokens_per_seq =*/ 1, +}; + llama_context::llama_context( const llama_model & model, llama_context_params params) : @@ -226,6 +232,9 @@ llama_context::llama_context( cparams.fused_gdn_ch = true; cparams.auto_fgdn = true; + cparams.fused_lid = true; + cparams.auto_flid = true; + // with causal attention, the batch size is limited by the context size cparams.n_batch = cparams.causal_attn ? std::min(cparams.n_ctx, params.n_batch) : params.n_batch; @@ -522,6 +531,12 @@ void llama_context::resolve_fused_ops(const llama_memory_context_i * mctx, uint3 resolve(llm_fused_op_gdn_ch_probe, cparams.fused_gdn_ch); cparams.auto_fgdn = false; } + + if (cparams.auto_flid) { + LLAMA_LOG_INFO("%s: resolving fused Lightning Indexer support:\n", func); + resolve(llm_fused_op_lid_probe, cparams.fused_lid); + cparams.auto_flid = false; + } } void llama_context::sched_reserve() { diff --git a/src/llama-cparams.h b/src/llama-cparams.h index 546ae1e2c12..58520caa365 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -41,6 +41,8 @@ struct llama_cparams { bool fused_gdn_ar; // use fused gated delta net (autoregressive) bool fused_gdn_ch; // use fused gated delta net (chunked) bool auto_fgdn; + bool fused_lid; // use fused lightning indexer + bool auto_flid; bool no_perf; bool warmup; // TODO: remove [TAG_LLAMA_GRAPH_NO_WARMUP] bool op_offload; diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index a8fd11ebc38..c8ecb0a2854 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -842,7 +842,7 @@ static void dsv4_build_comp_inputs( GGML_ASSERT(n_stream > 0); GGML_ASSERT(n_tokens%n_stream == 0); - inp.kq_mask = ggml_new_tensor_4d(ctx, cparams.flash_attn && strcmp(name, "lid") != 0 ? GGML_TYPE_F16 : GGML_TYPE_F32, plan.n_kv, n_tokens/n_stream, 1, n_stream); + inp.kq_mask = ggml_new_tensor_4d(ctx, (strcmp(name, "lid") != 0 && cparams.flash_attn) || (strcmp(name, "lid") == 0 && cparams.fused_lid) ? GGML_TYPE_F16 : GGML_TYPE_F32, plan.n_kv, n_tokens/n_stream, 1, n_stream); ggml_set_input(inp.kq_mask); ggml_set_name(inp.kq_mask, (std::string("dsv4_") + name + "_kq_mask").c_str()); } @@ -3025,9 +3025,9 @@ llm_graph_input_attn_k_dsa * llm_graph_context::build_attn_inp_k_dsa() const { { inp->self_k_idxs_lid = mctx_cur->get_lid()->build_input_k_idxs(ctx0, ubatch); - // ensure F32 mask + // ensure that mask type matches fused lightning indexer use (requires f16 mask) auto cparams_copy = cparams; - cparams_copy.flash_attn = false; + cparams_copy.flash_attn = cparams.fused_lid; inp->self_kq_mask_lid = build_attn_inp_kq_mask(ctx0, mctx_cur->get_lid(), ubatch, cparams_copy); inp->self_kq_mask_lid_cnv = inp->self_kq_mask_lid; diff --git a/src/llama-graph.h b/src/llama-graph.h index 97141ef93be..c84cb6a452b 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -42,6 +42,7 @@ enum llm_fused_op { LLM_FUSED_OP_FLASH_ATTN, LLM_FUSED_OP_GDN_AR, LLM_FUSED_OP_GDN_CH, + LLM_FUSED_OP_LIGHTNING_INDEXER, }; enum llm_ffn_op_type : int { diff --git a/src/models/deepseek32.cpp b/src/models/deepseek32.cpp index 9a20e2ce907..32262e6840b 100644 --- a/src/models/deepseek32.cpp +++ b/src/models/deepseek32.cpp @@ -301,43 +301,50 @@ llama_model_deepseek32::graph::graph(const llama_model & model, const llm_graph_ indexer_q = ggml_view_4d(ctx0, indexer_q, indexer_q->ne[0], indexer_q->ne[1], indexer_q->ne[2]/n_stream, n_stream, indexer_q->nb[1], indexer_q->nb[2], indexer_q->nb[3]/n_stream, 0); indexer_weights = ggml_view_4d(ctx0, indexer_weights, 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); - // calculate indexer kq - indexer_q = ggml_permute(ctx0, indexer_q, 0, 2, 1, 3); - cb(indexer_q, "indexer_q", il); - indexer_k = ggml_permute(ctx0, indexer_k, 0, 2, 1, 3); - cb(indexer_k, "indexer_k", il); - - ggml_tensor * indexer_kq = ggml_mul_mat(ctx0, indexer_k, indexer_q); - cb(indexer_kq, "indexer_kq", il); - - // ReLU requires contiguous tensors - indexer_kq = ggml_cont(ctx0, ggml_permute(ctx0, indexer_kq, 2, 1, 0, 3)); - cb(indexer_kq, "indexer_kq", il); - - // apply ReLU - ggml_tensor * indexer_score = ggml_relu(ctx0, indexer_kq); - cb(indexer_score, "indexer_score", il); - // pre-scale weights to avoid scaling operations on huge indexer_score tensor indexer_weights = ggml_scale(ctx0, indexer_weights, 1.0f / sqrtf(float(n_embd_indexer_head * n_indexer_head))); cb(indexer_weights, "indexer_weights", il); - // multiply scores by indexer weights - indexer_score = ggml_mul(ctx0, indexer_score, indexer_weights); - cb(indexer_score, "indexer_score", il); - - // sum by q n_indexer_head dimension - indexer_score = ggml_sum_rows(ctx0, indexer_score); - cb(indexer_score, "indexer_score", il); - - // permute result to match KQ mask - indexer_score = ggml_cont(ctx0, ggml_permute(ctx0, indexer_score, 2, 1, 0, 3)); - cb(indexer_score, "indexer_score", il); - - // mask indexer scores - ggml_tensor * indexer_kq_mask = inp_attn_dsa->get_kq_mask_lid(); - indexer_score = ggml_add(ctx0, indexer_score, indexer_kq_mask); - cb(indexer_score, "indexer_score", il); + ggml_tensor * indexer_score = nullptr; + if (cparams.fused_lid) { + indexer_score = ggml_lightning_indexer(ctx0, indexer_q, indexer_k, indexer_weights, inp_attn_dsa->get_kq_mask_lid()); + cb(indexer_score, "indexer_score", il); + res->add_fused_node({LLM_FUSED_OP_LIGHTNING_INDEXER, indexer_score, il}); + } else { + // calculate indexer kq + indexer_q = ggml_permute(ctx0, indexer_q, 0, 2, 1, 3); + cb(indexer_q, "indexer_q", il); + indexer_k = ggml_permute(ctx0, indexer_k, 0, 2, 1, 3); + cb(indexer_k, "indexer_k", il); + + ggml_tensor * indexer_kq = ggml_mul_mat(ctx0, indexer_k, indexer_q); + cb(indexer_kq, "indexer_kq", il); + + // ReLU requires contiguous tensors + indexer_kq = ggml_cont(ctx0, ggml_permute(ctx0, indexer_kq, 2, 1, 0, 3)); + cb(indexer_kq, "indexer_kq", il); + + // apply ReLU + indexer_score = ggml_relu(ctx0, indexer_kq); + cb(indexer_score, "indexer_score", il); + + // multiply scores by indexer weights + indexer_score = ggml_mul(ctx0, indexer_score, indexer_weights); + cb(indexer_score, "indexer_score", il); + + // sum by q n_indexer_head dimension + indexer_score = ggml_sum_rows(ctx0, indexer_score); + cb(indexer_score, "indexer_score", il); + + // permute result to match KQ mask + indexer_score = ggml_cont(ctx0, ggml_permute(ctx0, indexer_score, 2, 1, 0, 3)); + cb(indexer_score, "indexer_score", il); + + // mask indexer scores + ggml_tensor * indexer_kq_mask = inp_attn_dsa->get_kq_mask_lid(); + indexer_score = ggml_add(ctx0, indexer_score, indexer_kq_mask); + cb(indexer_score, "indexer_score", il); + } // get indices of top k indexer scores uint32_t n_top_k = indexer_score->ne[0] < n_indexer_top_k ? indexer_score->ne[0] : n_indexer_top_k; diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp index 07aa477e1e4..e64e43451f7 100644 --- a/src/models/deepseek4.cpp +++ b/src/models/deepseek4.cpp @@ -556,25 +556,32 @@ 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_score = nullptr; + if (cparams.fused_lid) { + indexer_score = ggml_lightning_indexer(ctx0, indexer_q, indexer_k, indexer_weights, inp_lid.kq_mask); + cb(indexer_score, "lid_score_masked", il); + res->add_fused_node({LLM_FUSED_OP_LIGHTNING_INDEXER, indexer_score, il}); + } else { + 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); + 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); + 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)); - cb(indexer_score, "lid_score", il); + 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)); + cb(indexer_score, "lid_score", il); - indexer_score = ggml_add(ctx0, indexer_score, inp_lid.kq_mask); - cb(indexer_score, "lid_score_masked", il); + indexer_score = ggml_add(ctx0, indexer_score, inp_lid.kq_mask); + cb(indexer_score, "lid_score_masked", il); + } const uint32_t n_top_k = indexer_score->ne[0] < hparams.indexer_top_k ? indexer_score->ne[0] : hparams.indexer_top_k; ggml_tensor * top_k = ggml_cont(ctx0, ggml_top_k(ctx0, indexer_score, n_top_k));