From f2bc4dfa873e417e2ffa7fc6c3670aa6c4fe2bf3 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Tue, 1 Sep 2026 12:09:25 +0300 Subject: [PATCH] metal : fuse gated_delta_net with cache cpy Add GGML_METAL_FUSE_GDN_CACHE to the fusion table: when the gated_delta_net kernel is followed by a cpy that scatters its recurrent state snapshots into the KV cache, the kernel writes the snapshots straight into the cache buffer and the trailing cpy is elided. The gdn output has other consumers (the attn scores view), so unlike the elision-chain patterns this is not a simple chain: a 'raw' flag on the fusion pattern skips the generic chain/shape and ggml_can_fuse_subgraph_ext checks, making the pattern-specific check callback the sole validator. Packing (ggml_metal_fuse_max) now matches on the same view-transparent node sequence that the compute phase uses, so the gdn + cache cpy group is packed along with any intermediate views and stays adjacent through the reorder. The fused cpy is a view consumer of the gdn (it writes the cache directly), so its mem-range is skipped in the encoder; the skip is restricted to CPY nodes consuming the previous fused node through a view so other fusions are unaffected. Add test_gated_delta_net_cache_fusion and register 5 cases. Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-0731 --- ggml/src/ggml-metal/ggml-metal-fuse.cpp | 173 +++++++++++++----- ggml/src/ggml-metal/ggml-metal-fuse.h | 5 + ggml/src/ggml-metal/ggml-metal-impl.h | 1 + ggml/src/ggml-metal/ggml-metal-ops.cpp | 51 +++++- .../ggml-metal/kernels/gated_delta_net.metal | 20 +- tests/test-backend-ops.cpp | 123 +++++++++++++ 6 files changed, 321 insertions(+), 52 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-fuse.cpp b/ggml/src/ggml-metal/ggml-metal-fuse.cpp index 7ded58534a..1b8fe11463 100644 --- a/ggml/src/ggml-metal/ggml-metal-fuse.cpp +++ b/ggml/src/ggml-metal/ggml-metal-fuse.cpp @@ -89,6 +89,63 @@ static bool ggml_metal_fuse_check_add_chain(const struct ggml_tensor * const * n return true; } +// GATED_DELTA_NET + CPY: the trailing cpy scatters the gdn state snapshots into the recurrent +// cache, so the gdn kernel writes them straight to the cache and the cpy is elided. +// mirrors ggml_metal_op_can_fuse_gdn_cache (PR #25788). the gdn output has other consumers (the +// attn scores view), so unlike the other patterns this is not an elision chain: the structural +// checks live entirely in this callback (raw = true). +static bool ggml_metal_fuse_check_gdn_cache(const struct ggml_tensor * const * nodes, + const struct ggml_metal_fuse * fuse, + enum ggml_metal_fuse_mode mode) { + const struct ggml_tensor * gdn = nodes[0]; + const struct ggml_tensor * cpy = nodes[1]; + + // the kernel skips the snapshot tail, so the gdn output must not be a graph output + if (gdn->type != GGML_TYPE_F32 || (gdn->flags & GGML_TENSOR_FLAG_OUTPUT)) { + return false; + } + + if (cpy->op != GGML_OP_CPY || (cpy->flags & GGML_TENSOR_FLAG_OUTPUT)) { + return false; + } + + const int64_t S_v = gdn->src[2]->ne[0]; + const int64_t H = gdn->src[2]->ne[1]; + const int64_t n_tokens = gdn->src[2]->ne[2]; + const int64_t n_seqs = gdn->src[2]->ne[3]; + const int64_t K = ggml_get_op_params_i32(gdn, 0); + const size_t tail_off = ggml_row_size(GGML_TYPE_F32, S_v * H * n_tokens * n_seqs); + + const int64_t D = S_v * S_v * H; + const int64_t n_written = std::min(n_tokens, K); + + const struct ggml_tensor * src = cpy->src[0]; // gdn snapshot tail view + const struct ggml_tensor * dst = cpy->src[1]; // cache view + + // src must be this gdn's snapshot tail (contiguous, at the tail offset) + if (src->op != GGML_OP_VIEW || src->view_src != gdn || + src->view_offs != tail_off || !ggml_is_contiguous(src)) { + return false; + } + + const int64_t expected_ne[GGML_MAX_DIMS] = { D, n_seqs, n_written, 1 }; + if (dst->type != GGML_TYPE_F32 || + !std::equal(expected_ne, expected_ne + GGML_MAX_DIMS, dst->ne) || + dst->nb[0] != ggml_type_size(GGML_TYPE_F32) || + dst->nb[1] != (size_t) ggml_row_size(GGML_TYPE_F32, D)) { + return false; + } + + if (mode == GGML_METAL_FUSE_FULL) { + // the cache must be allocated so the kernel can write straight to its buffer + if (dst->data == nullptr) { + return false; + } + } + + return true; +} + // MUL + SIN + SQR + MUL + ADD (snake activation) static bool ggml_metal_fuse_check_snake(const struct ggml_tensor * const * nodes, const struct ggml_metal_fuse * fuse, @@ -144,20 +201,22 @@ static const enum ggml_op ops_add_4[] = { GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, static const enum ggml_op ops_add_5[] = { GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD }; static const enum ggml_op ops_add_6[] = { GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD }; static const enum ggml_op ops_add_7[] = { GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD }; -static const enum ggml_op ops_snake[] = { GGML_OP_MUL, GGML_OP_SIN, GGML_OP_SQR, GGML_OP_MUL, GGML_OP_ADD }; +static const enum ggml_op ops_snake[] = { GGML_OP_MUL, GGML_OP_SIN, GGML_OP_SQR, GGML_OP_MUL, GGML_OP_ADD }; +static const enum ggml_op ops_gdn_cache[] = { GGML_OP_GATED_DELTA_NET, GGML_OP_CPY }; static const struct ggml_metal_fuse ggml_metal_fuses[] = { - { GGML_METAL_FUSE_NORM_MUL, ops_norm_mul, 2, nullptr, 0, ggml_metal_fuse_check_norm }, - { GGML_METAL_FUSE_NORM_MUL_ADD, ops_norm_mul_add, 3, nullptr, 0, ggml_metal_fuse_check_norm }, - { GGML_METAL_FUSE_NORM_MUL, ops_rms_norm_mul, 2, nullptr, 0, ggml_metal_fuse_check_norm }, - { GGML_METAL_FUSE_NORM_MUL_ADD, ops_rms_norm_mul_add, 3, nullptr, 0, ggml_metal_fuse_check_norm }, - { GGML_METAL_FUSE_ADD_CHAIN, ops_add_2, 2, nullptr, 0, ggml_metal_fuse_check_add_chain }, - { GGML_METAL_FUSE_ADD_CHAIN, ops_add_3, 3, nullptr, 0, ggml_metal_fuse_check_add_chain }, - { GGML_METAL_FUSE_ADD_CHAIN, ops_add_4, 4, nullptr, 0, ggml_metal_fuse_check_add_chain }, - { GGML_METAL_FUSE_ADD_CHAIN, ops_add_5, 5, nullptr, 0, ggml_metal_fuse_check_add_chain }, - { GGML_METAL_FUSE_ADD_CHAIN, ops_add_6, 6, nullptr, 0, ggml_metal_fuse_check_add_chain }, - { GGML_METAL_FUSE_ADD_CHAIN, ops_add_7, 7, nullptr, 0, ggml_metal_fuse_check_add_chain }, - { GGML_METAL_FUSE_SNAKE, ops_snake, 5, nullptr, 0, ggml_metal_fuse_check_snake }, + { GGML_METAL_FUSE_NORM_MUL, ops_norm_mul, 2, nullptr, 0, false, ggml_metal_fuse_check_norm }, + { GGML_METAL_FUSE_NORM_MUL_ADD, ops_norm_mul_add, 3, nullptr, 0, false, ggml_metal_fuse_check_norm }, + { GGML_METAL_FUSE_NORM_MUL, ops_rms_norm_mul, 2, nullptr, 0, false, ggml_metal_fuse_check_norm }, + { GGML_METAL_FUSE_NORM_MUL_ADD, ops_rms_norm_mul_add, 3, nullptr, 0, false, ggml_metal_fuse_check_norm }, + { GGML_METAL_FUSE_ADD_CHAIN, ops_add_2, 2, nullptr, 0, false, ggml_metal_fuse_check_add_chain }, + { GGML_METAL_FUSE_ADD_CHAIN, ops_add_3, 3, nullptr, 0, false, ggml_metal_fuse_check_add_chain }, + { GGML_METAL_FUSE_ADD_CHAIN, ops_add_4, 4, nullptr, 0, false, ggml_metal_fuse_check_add_chain }, + { GGML_METAL_FUSE_ADD_CHAIN, ops_add_5, 5, nullptr, 0, false, ggml_metal_fuse_check_add_chain }, + { GGML_METAL_FUSE_ADD_CHAIN, ops_add_6, 6, nullptr, 0, false, ggml_metal_fuse_check_add_chain }, + { GGML_METAL_FUSE_ADD_CHAIN, ops_add_7, 7, nullptr, 0, false, ggml_metal_fuse_check_add_chain }, + { GGML_METAL_FUSE_SNAKE, ops_snake, 5, nullptr, 0, false, ggml_metal_fuse_check_snake }, + { GGML_METAL_FUSE_GDN_CACHE, ops_gdn_cache, 2, nullptr, 0, true, ggml_metal_fuse_check_gdn_cache }, }; const struct ggml_metal_fuse * ggml_metal_fuse_all(int * n) { @@ -194,20 +253,25 @@ const struct ggml_metal_fuse * ggml_metal_fuse_next( continue; } - // the first node must match the pattern start - if (gf->nodes[node_idxs[idx]]->op != fuse->ops[0]) { - continue; - } - const struct ggml_tensor * nodes[GGML_METAL_FUSE_MAX]; - // common element-wise chain constraints: each node reads the previous one, - // and all nodes have the same shape + // the op sequence must match exactly bool ok = true; for (int j = 0; j < fuse->n_ops; j++) { nodes[j] = gf->nodes[node_idxs[idx + j]]; + if (nodes[j]->op != fuse->ops[j]) { + ok = false; + break; + } + } + if (!ok) { + continue; + } - if (j > 0) { + if (!fuse->raw) { + // common element-wise chain constraints: each node reads the previous one, + // and all nodes have the same shape + for (int j = 1; j < fuse->n_ops && ok; j++) { if (nodes[j]->src[0] != nodes[j - 1] && nodes[j]->src[1] != nodes[j - 1]) { ok = false; break; @@ -217,23 +281,23 @@ const struct ggml_metal_fuse * ggml_metal_fuse_next( break; } } - } - if (!ok) { - continue; + if (!ok) { + continue; + } + + // ggml_can_fuse_subgraph_ext expects outputs as absolute graph node indices + int outputs_buf[GGML_MAX_SRC]; + outputs_buf[0] = node_idxs[idx + fuse->n_ops - 1]; + const int * outputs = ggml_metal_fuse_outputs(fuse, outputs_buf); + const int n_outputs = fuse->n_outputs ? fuse->n_outputs : 1; + + // structural subgraph checks (op sequence, elidable uses, view containment) + if (!ggml_can_fuse_subgraph_ext(gf, node_idxs + idx, fuse->n_ops, fuse->ops, outputs, n_outputs)) { + continue; + } } - // ggml_can_fuse_subgraph_ext expects outputs as absolute graph node indices - int outputs_buf[GGML_MAX_SRC]; - outputs_buf[0] = node_idxs[idx + fuse->n_ops - 1]; - const int * outputs = ggml_metal_fuse_outputs(fuse, outputs_buf); - const int n_outputs = fuse->n_outputs ? fuse->n_outputs : 1; - - // structural subgraph checks (op sequence, elidable uses, view containment) - if (!ggml_can_fuse_subgraph_ext(gf, node_idxs + idx, fuse->n_ops, fuse->ops, outputs, n_outputs)) { - continue; - } - - // pattern-specific checks + // pattern-specific checks (the sole validator for raw patterns) if (fuse->check && !fuse->check(nodes, fuse, mode)) { continue; } @@ -248,28 +312,45 @@ const struct ggml_metal_fuse * ggml_metal_fuse_next( } // optimize phase: maximum number of nodes starting at idx (a raw sequential graph index) that -// could be fused, chaining patterns back-to-back +// could be fused, chaining patterns back-to-back. matching runs on the same filtered (view +// transparent) node sequence that the compute phase uses, so the returned count is the raw index +// span from idx to the last matched node (intermediate views are packed along). int ggml_metal_fuse_max(const struct ggml_cgraph * gf, int idx) { + // an empty/view node cannot start a pattern - pack it alone + if (ggml_op_is_empty(gf->nodes[idx]->op) || ggml_is_empty(gf->nodes[idx])) { + return 1; + } + + // collect the non-empty node indices starting at idx int idxs[GGML_METAL_FUSE_MAX]; + int n_idxs = 0; + for (int i = idx; i < gf->n_nodes && n_idxs < GGML_METAL_FUSE_MAX; i++) { + if (!ggml_op_is_empty(gf->nodes[i]->op) && !ggml_is_empty(gf->nodes[i])) { + idxs[n_idxs++] = i; + } + } + if (n_idxs == 0) { + return 1; + } int total = 0; - int i = idx; - - while (i < gf->n_nodes && total < GGML_METAL_FUSE_MAX) { - const int n_idxs = std::min(GGML_METAL_FUSE_MAX, (int) gf->n_nodes - i); - for (int j = 0; j < n_idxs; j++) { - idxs[j] = i + j; - } + int i_f = 0; + while (i_f < n_idxs && total < GGML_METAL_FUSE_MAX) { int len = 1; - const struct ggml_metal_fuse * fuse = ggml_metal_fuse_next(gf, idxs, n_idxs, 0, GGML_METAL_FUSE_STRUCTURAL, &len); + const struct ggml_metal_fuse * fuse = ggml_metal_fuse_next(gf, idxs, n_idxs, i_f, GGML_METAL_FUSE_STRUCTURAL, &len); if (!fuse || total + len > GGML_METAL_FUSE_MAX) { break; } total += len; - i += len; + i_f += len; } - return std::max(total, 1); + if (i_f == 0) { + return 1; + } + + // map the matched non-empty nodes back to the raw index span (views are included) + return std::min(GGML_METAL_FUSE_MAX, idxs[i_f - 1] - idx + 1); } diff --git a/ggml/src/ggml-metal/ggml-metal-fuse.h b/ggml/src/ggml-metal/ggml-metal-fuse.h index 9186926932..6a3b4a19af 100644 --- a/ggml/src/ggml-metal/ggml-metal-fuse.h +++ b/ggml/src/ggml-metal/ggml-metal-fuse.h @@ -34,6 +34,7 @@ enum ggml_metal_fuse_id { GGML_METAL_FUSE_NORM_MUL_ADD, // NORM/RMS_NORM + MUL + ADD GGML_METAL_FUSE_ADD_CHAIN, // ADD x N (N in [2, 7]) GGML_METAL_FUSE_SNAKE, // MUL + SIN + SQR + MUL + ADD + GGML_METAL_FUSE_GDN_CACHE, // GATED_DELTA_NET + CPY (write snapshots into the recurrent cache) }; struct ggml_metal_fuse { @@ -42,6 +43,10 @@ struct ggml_metal_fuse { int n_ops; // number of ops const int * outputs; // output node indices (absolute graph indices; nullptr => the last node) int n_outputs;// number of outputs (0 => default last node) + // if raw: the generic chain/shape + ggml_can_fuse_subgraph checks are skipped and the + // check callback below is the sole validator (used for patterns that are not elision chains, + // e.g. the gdn + cache-cpy write-through fusion) + bool raw; // extra backend constraints on top of ggml_can_fuse_subgraph // nodes[j] is the j-th node of the pattern bool (*check)(const struct ggml_tensor * const * nodes, diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index 30e40f527f..8de80624dd 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -978,6 +978,7 @@ typedef struct { uint64_t nb1; uint64_t nb2; uint64_t nb3; + int32_t state_out_stride; // 0 => snapshots are appended after the attn scores (unfused) } ggml_metal_kargs_gated_delta_net; typedef struct { diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index e5cc36bc10..d5d276e839 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -88,6 +88,19 @@ struct ggml_metal_op { return ggml_metal_fuse_next(gf, idxs.data(), (int) idxs.size(), i0, mode, n_out); } + // true if node i is a cpy fused into the previous node through a view: the fusing op already + // wrote its dst directly, so its mem-range must not be tracked (e.g. the gdn + cache cpy) + bool is_view_consumer(int i) const { + if (i <= 0) { + return false; + } + + const ggml_tensor * node = this->node(i); + const ggml_tensor * prev = this->node(i - 1); + + return node->op == GGML_OP_CPY && node->src[0] && node->src[0]->view_src == prev; + } + ggml_metal_device_t dev; ggml_metal_library_t lib; ggml_metal_encoder_t enc; @@ -510,6 +523,11 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) { // update the mem ranges in the encoding context for (int i = 0; i < n_fuse; ++i) { + // view consumers are handled by the fusing op (they wrote the dst directly) + if (ctx->is_view_consumer(idx + i)) { + continue; + } + if (!ggml_metal_op_concurrency_add(ctx, ctx->node(idx + i))) { ggml_metal_op_concurrency_reset(ctx); } @@ -1867,6 +1885,8 @@ int ggml_metal_op_gated_delta_net(ggml_metal_op_t ctx, int idx) { ggml_metal_library_t lib = ctx->lib; ggml_metal_encoder_t enc = ctx->enc; + const bool use_fusion = ctx->use_fusion; + const int debug_fusion = ctx->debug_fusion; GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); @@ -1879,6 +1899,29 @@ int ggml_metal_op_gated_delta_net(ggml_metal_op_t ctx, int idx) { auto pipeline = ggml_metal_library_get_pipeline_gated_delta_net(lib, op); + // when fused with the trailing cache cpy, the snapshots are written straight into the + // recurrent cache and the cpy is skipped (see GGML_METAL_FUSE_GDN_CACHE) + ggml_metal_buffer_id bid_state_out = ggml_metal_get_buffer_id(op); + int32_t state_out_stride = 0; + int n_fuse = 1; + + if (use_fusion) { + int n = 1; + const ggml_metal_fuse * fuse = ctx->can_fuse(idx, GGML_METAL_FUSE_FULL, &n); + + if (fuse && fuse->id == GGML_METAL_FUSE_GDN_CACHE) { + const ggml_tensor * dst_cache = ctx->node(idx + 1)->src[1]; // cache view + + bid_state_out = ggml_metal_get_buffer_id(dst_cache); + state_out_stride = (int32_t) (dst_cache->nb[2]/sizeof(float)); + n_fuse = 2; + + if (debug_fusion > 1) { + GGML_LOG_DEBUG("%s: fuse: GATED_DELTA_NET + CPY\n", __func__); + } + } + } + int ida = 0; ggml_metal_kargs_gated_delta_net args = { @@ -1917,23 +1960,25 @@ int ggml_metal_op_gated_delta_net(ggml_metal_op_t ctx, int idx) { /*.nb1 =*/ nb1, /*.nb2 =*/ nb2, /*.nb3 =*/ nb3, + /*.state_out_stride =*/ state_out_stride, }; ggml_metal_encoder_set_pipeline(enc, pipeline); - ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), ida++); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), ida++); // args ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), ida++); // q ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), ida++); // k ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[2]), ida++); // v ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[3]), ida++); // gate ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[4]), ida++); // beta ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[5]), ida++); // state - ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), ida++); // dst + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), ida++); // dst (attn) + ggml_metal_encoder_set_buffer (enc, bid_state_out, ida++); // state_out const int nsg = pipeline.nsg; ggml_metal_encoder_dispatch_threadgroups(enc, op->src[2]->ne[0]/nsg, op->src[2]->ne[1], op->src[2]->ne[3], 32, nsg, 1); - return 1; + return n_fuse; } int ggml_metal_op_solve_tri(ggml_metal_op_t ctx, int idx) { diff --git a/ggml/src/ggml-metal/kernels/gated_delta_net.metal b/ggml/src/ggml-metal/kernels/gated_delta_net.metal index 8422d8e29f..cbfec5c795 100644 --- a/ggml/src/ggml-metal/kernels/gated_delta_net.metal +++ b/ggml/src/ggml-metal/kernels/gated_delta_net.metal @@ -15,6 +15,7 @@ kernel void kernel_gated_delta_net_impl( device const char * b, device const char * s, device char * dst, + device char * state_out, uint3 tgpig[[threadgroup_position_in_grid]], uint3 tpitg[[thread_position_in_threadgroup]], uint3 ntg[[threads_per_threadgroup]]) { @@ -65,6 +66,12 @@ kernel void kernel_gated_delta_net_impl( // per-(seq,head) offset within a slot const uint state_out_base = (i23*args.ne21 + i21)*S_v*S_v + i20*S_v; + // when fused with the cache cpy, write the snapshots straight into the cache buffer using + // the slot stride; otherwise append them after the attn scores (state_out_stride == 0) + const bool fused = args.state_out_stride > 0; + const device float * state_out_ = fused ? (device float *)state_out : (device float *)dst + attn_size; + const uint slot_stride = fused ? (uint)args.state_out_stride : state_size_per_snap; + for (short t = 0; t < args.ne22; t++) { float s_k = 0.0f; @@ -116,7 +123,7 @@ kernel void kernel_gated_delta_net_impl( if (K > 1) { const int target_slot = (int)args.ne22 - 1 - (int)t; if (target_slot >= 0 && target_slot < (int)K) { - device float * dst_state = (device float *) (dst) + attn_size + (uint)target_slot * state_size_per_snap + state_out_base; + device float * dst_state = (device float *)state_out_ + (uint)target_slot * slot_stride + state_out_base; FOR_UNROLL (short j = 0; j < NSG; j++) { const short is = tx*NSG + j; dst_state[is] = ls[j]; @@ -126,7 +133,7 @@ kernel void kernel_gated_delta_net_impl( } if (K == 1) { - device float * dst_state = (device float *) (dst) + attn_size + state_out_base; + device float * dst_state = (device float *)state_out_ + state_out_base; FOR_UNROLL (short j = 0; j < NSG; j++) { const short is = tx*NSG + j; dst_state[is] = ls[j]; @@ -158,6 +165,7 @@ kernel void kernel_gated_delta_net_impl( device const char * b, device const char * s, device char * dst, + device char * state_out, uint3 tgpig[[threadgroup_position_in_grid]], uint3 tpitg[[thread_position_in_threadgroup]], uint3 ntg[[threads_per_threadgroup]]) { @@ -230,7 +238,13 @@ kernel void kernel_gated_delta_net_impl( dst_attn += args.ne21*S_v; } - device float * dst_state = (device float *) (dst) + args.ne23*args.ne22*args.ne21*S_v + (i23*args.ne21 + i21)*S_v*S_v + i20; + // when fused with the cache cpy, write the snapshots straight into the cache buffer using + // the slot stride; otherwise append them after the attn scores (state_out_stride == 0) + const bool fused = args.state_out_stride > 0; + const device float * state_out_ = fused ? (device float *)state_out : (device float *)dst + args.ne23*args.ne22*args.ne21*S_v; + const uint slot_stride = fused ? (uint)args.state_out_stride : S_v*S_v; + + device float * dst_state = (device float *)state_out_ + (i23*args.ne21 + i21)*slot_stride + i20; device T * dstt_state = (device T *) (dst_state); FOR_UNROLL (short j = 0; j < NSG; j++) { diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index eff63575f2..2272778d0d 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -4495,6 +4495,122 @@ struct test_gated_delta_net : public test_case { } }; +// GGML_OP_GATED_DELTA_NET + GGML_OP_CPY (recurrent cache fusion) +struct test_gated_delta_net_cache_fusion : public test_case { + const ggml_type type; + + const int64_t head_count; + const int64_t head_size; + const int64_t n_seq_tokens; + const int64_t n_seqs; + const int64_t K; // snapshot slot count (>1) + + ggml_tensor * cpy_node = nullptr; + + std::string vars() override { + return VARS_TO_STR6(type, head_count, head_size, n_seq_tokens, n_seqs, K); + } + + test_gated_delta_net_cache_fusion(ggml_type type = GGML_TYPE_F32, + int64_t head_count = 4, int64_t head_size = 32, int64_t n_seq_tokens = 2, int64_t n_seqs = 1, + int64_t K = 2) + : type(type), head_count(head_count), head_size(head_size), n_seq_tokens(n_seq_tokens), n_seqs(n_seqs), K(K) {} + + ggml_tensor * build_graph(ggml_context * ctx) override { + const int64_t S_v = head_size; + const int64_t H_v = head_count; + const int64_t H_k = head_count; + const int64_t D = S_v * S_v * H_v; + const int64_t n_written = std::min(n_seq_tokens, K); + + ggml_tensor * q = ggml_new_tensor_4d(ctx, type, head_size, H_k, n_seq_tokens, n_seqs); + ggml_tensor * k = ggml_new_tensor_4d(ctx, type, head_size, H_k, n_seq_tokens, n_seqs); + ggml_tensor * v = ggml_new_tensor_4d(ctx, type, head_size, H_v, n_seq_tokens, n_seqs); + ggml_set_name(q, "q"); + ggml_set_name(k, "k"); + ggml_set_name(v, "v"); + ggml_tensor * g = ggml_new_tensor_4d(ctx, type, 1, H_v, n_seq_tokens, n_seqs); + ggml_tensor * beta = ggml_new_tensor_4d(ctx, type, 1, H_v, n_seq_tokens, n_seqs); + ggml_tensor * state = ggml_new_tensor_4d(ctx, type, head_size, head_size, H_v, n_seqs); + ggml_set_name(g, "g"); + ggml_set_name(beta, "beta"); + ggml_set_name(state, "state"); + + q = ggml_l2_norm(ctx, q, 1e-6f); + k = ggml_l2_norm(ctx, k, 1e-6f); + + ggml_tensor * gdn_out = ggml_gated_delta_net(ctx, q, k, v, g, beta, state, K); + ggml_set_name(gdn_out, "gdn_out"); + + // attn scores view (first part of the gdn output) + ggml_tensor * attn = ggml_view_4d(ctx, gdn_out, + S_v, H_v, n_seq_tokens, n_seqs, + ggml_row_size(gdn_out->type, S_v), + ggml_row_size(gdn_out->type, S_v * H_v), + ggml_row_size(gdn_out->type, S_v * H_v * n_seq_tokens), 0); + ggml_set_name(attn, "attn"); + + // snapshot tail view [D, n_seqs, n_written] + const int64_t attn_score_elems = S_v * H_v * n_seq_tokens * n_seqs; + ggml_tensor * src = ggml_view_3d(ctx, gdn_out, + D, n_seqs, n_written, + ggml_row_size(gdn_out->type, D), + ggml_row_size(gdn_out->type, D * n_seqs), + ggml_row_size(gdn_out->type, attn_score_elems)); + + // recurrent cache view [D, n_seqs, n_written] + ggml_tensor * cache = ggml_new_tensor_3d(ctx, type, D, n_seqs, n_written); + ggml_set_name(cache, "cache"); + ggml_tensor * dst = ggml_view_3d(ctx, cache, + D, n_seqs, n_written, + ggml_row_size(cache->type, D), + ggml_row_size(cache->type, D * n_seqs), 0); + + ggml_tensor * cpy = ggml_cpy(ctx, src, dst); + ggml_set_name(cpy, "gdn_cache_cpy"); + cpy_node = cpy; + + // read the cpy output (not the plain dst view, which would not pull the cpy into the graph) + // so that neither the gdn nor the cpy is the graph output + ggml_tensor * out = ggml_sum(ctx, cpy); + return out; + } + + std::string op_desc(ggml_tensor * t) override { + GGML_UNUSED(t); + return "GATED_DELTA_NET_CACHE_FUSION"; + } + + bool run_whole_graph() override { return true; } + std::vector fusion_test_nodes() override { return { cpy_node }; } + + uint64_t op_flops(ggml_tensor * t) override { + GGML_UNUSED(t); + const uint64_t S_v = head_size; + const uint64_t H_v = head_count; + const uint64_t T = n_seq_tokens; + const uint64_t B = n_seqs; + return (4ull*S_v + 2ull*S_v*S_v) * H_v * T * B; + } + + void initialize_tensors(ggml_context * ctx) override { + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { + if (ggml_is_view_op(t->op)) { continue; } + if (strcmp(t->name, "g") == 0) { + init_tensor_uniform(t, -20.0f, -1e-4f); + } else if (strcmp(t->name, "beta") == 0) { + init_tensor_uniform(t, 0.0f, 1.0f); + } else if (strcmp(t->name, "v") == 0) { + init_tensor_uniform(t, -0.3f, 5.0f); + } else if (strcmp(t->name, "cache") == 0) { + init_tensor_uniform(t, 0.0f, 0.0f); + } else { + init_tensor_uniform(t); + } + } + } +}; + // GGML_OP_GATED_LINEAR_ATTN struct test_gla : public test_case { const ggml_type type; @@ -10487,6 +10603,13 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 32, 8, 1, 1, false, false, /*K=*/3)); test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 16, 2, 1, false, false, /*K=*/4)); + // gdn + cache cpy fusion (K > 1) + test_cases.emplace_back(new test_gated_delta_net_cache_fusion(GGML_TYPE_F32, 4, 32, 2, 1, 2)); + test_cases.emplace_back(new test_gated_delta_net_cache_fusion(GGML_TYPE_F32, 4, 64, 4, 1, 2)); + test_cases.emplace_back(new test_gated_delta_net_cache_fusion(GGML_TYPE_F32, 4, 32, 4, 1, 4)); + test_cases.emplace_back(new test_gated_delta_net_cache_fusion(GGML_TYPE_F32, 8, 32, 4, 2, 4)); + test_cases.emplace_back(new test_gated_delta_net_cache_fusion(GGML_TYPE_F32, 4, 32, 8, 1, 4)); + #if 0 // these tests are disabled to save execution time, sbut they can be handy for debugging test_cases.emplace_back(new test_llama(2, true));