diff --git a/ggml/include/ggml-rpc.h b/ggml/include/ggml-rpc.h index 5ad121ae57f1..efa5420a6982 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 ac133665d978..b2859ebe7df2 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 a82842fcffc0..2745a7dbb29e 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 fde939b4ad2a..b3288ce8ce40 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -11556,3 +11556,77 @@ 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 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)); + + 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) { + 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) { + k_to_float(src1_row, src1_row_f32, n_embd); + } else { + src1_row_f32 = (float *) src1_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 + 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); + // ReLU and weights (prescaled) + score += MAX(qk, 0.0f) * src2_row[i_head]; + } + // apply mask + dst_row[i_kv] = score + GGML_CPU_FP16_TO_FP32(src3_row[i_kv]); + } + } + } +} diff --git a/ggml/src/ggml-cpu/ops.h b/ggml/src/ggml-cpu/ops.h index a8e18c716db7..e956c25d3eda 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 de0321d9ffd9..8334440fc871 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]); + + 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/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 0c0aac2bba76..e2401ea42210 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; }