From 00f5442cc4e805293280c8f85d21d8f9d4aad206 Mon Sep 17 00:00:00 2001 From: fairydreaming <166155368+fairydreaming@users.noreply.github.com> Date: Sat, 11 Jul 2026 11:39:07 +0200 Subject: [PATCH] ggml : add GGML_OP_LIGHTNING_INDEXER that implements DeepSeek V3.2/V4 lightning indexer (#24231) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 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 --- ggml/include/ggml-rpc.h | 4 +- ggml/include/ggml.h | 19 ++++++++ ggml/src/ggml-cpu/ggml-cpu.c | 11 +++++ ggml/src/ggml-cpu/ops.cpp | 84 ++++++++++++++++++++++++++++++++++ ggml/src/ggml-cpu/ops.h | 1 + ggml/src/ggml.c | 42 ++++++++++++++++- src/llama-context.cpp | 15 +++++++ src/llama-cparams.h | 2 + src/llama-graph.cpp | 6 +-- src/llama-graph.h | 1 + src/models/deepseek32.cpp | 67 ++++++++++++++------------- src/models/deepseek4.cpp | 37 ++++++++------- tests/test-backend-ops.cpp | 87 ++++++++++++++++++++++++++++++++++++ 13 files changed, 324 insertions(+), 52 deletions(-) diff --git a/ggml/include/ggml-rpc.h b/ggml/include/ggml-rpc.h index 5ad121ae57..efa5420a69 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 diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index ac133665d9..b2859ebe7d 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -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); diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index a82842fcff..2745a7dbb2 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -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; } diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index df0028cf15..85eaebb573 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -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]); + } + } + } +} diff --git a/ggml/src/ggml-cpu/ops.h b/ggml/src/ggml-cpu/ops.h index a8e18c716d..e956c25d3e 100644 --- a/ggml/src/ggml-cpu/ops.h +++ b/ggml/src/ggml-cpu/ops.h @@ -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); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index de0321d9ff..5d88049916 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -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) { diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 5edfc85abf..3a469bc90b 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 546ae1e2c1..58520caa36 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 a8fd11ebc3..c8ecb0a285 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 97141ef93b..c84cb6a452 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 9a20e2ce90..32262e6840 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); + 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); - // sum by q n_indexer_head dimension - indexer_score = ggml_sum_rows(ctx0, indexer_score); - cb(indexer_score, "indexer_score", il); + ggml_tensor * indexer_kq = ggml_mul_mat(ctx0, indexer_k, indexer_q); + cb(indexer_kq, "indexer_kq", 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); + // ReLU requires contiguous tensors + indexer_kq = ggml_cont(ctx0, ggml_permute(ctx0, indexer_kq, 2, 1, 0, 3)); + cb(indexer_kq, "indexer_kq", 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); + // 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 07aa477e1e..e64e43451f 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)); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 1fae3f5176..084344fb25 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -7097,6 +7097,67 @@ 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 int64_t nm; // ne[3] of mask + + const ggml_type type_K; + + std::string vars() override { + return VARS_TO_STR7(hsk, nh, kv, nb, ns, nm, type_K); + } + + double max_nmse_err() override { + 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, 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); + 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, nm); + 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; @@ -9393,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 (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, nm, type_K)); + } + } + } + } + } + return test_cases; } #ifdef _MSC_VER @@ -9722,6 +9796,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, ns, type_K)); + } + } + } + } + } + return test_cases; }