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