diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index 8ac461bd3edf..67e88efc466a 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -2953,7 +2953,13 @@ struct ggml_cplan ggml_graph_plan( case GGML_OP_GATED_DELTA_NET: { const int64_t S_v = node->src[2]->ne[0]; - const int64_t K = node->src[5]->ne[1]; // state is (D, K, n_seqs) + // K = snapshot-slot count, from op_params -- shared by both + // op variants. src[5]->ne[1] is only K for the legacy + // (D,K,n_seqs) state; in rows mode src[5] is the 2D cache + // view whose ne[1] is the cache row count, so reading it + // there undersizes the scratch (overflow for a 1-row cache + // with K>1, i.e. batch-1 block decode). + const int64_t K = ggml_get_op_params_i32(node, 0); const int64_t per_thread = S_v + (K > 1 ? S_v * S_v : 0); cur = per_thread * sizeof(float) * n_tasks; } break; diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 3c1eb188db00..32f41f15df6c 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -1624,11 +1624,31 @@ static int ggml_metal_gdn_write_rows( *fused_set_rows = nullptr; const ggml_tensor * gdn = ctx->node(idx); - if (gdn->op != GGML_OP_GATED_DELTA_NET || gdn->src[6] == nullptr || + // honor the backend-wide fusion switch, like every other Metal fusion + if (!ctx->use_fusion || + gdn->op != GGML_OP_GATED_DELTA_NET || gdn->src[6] == nullptr || getenv("GGML_GDN_WRITE_FOLD_DISABLE") != nullptr) { return 1; } + // expected geometry of the recurrent-ring snapshot the kernel will scatter: + // the GDN output is [attn scores | K state snapshots]; the fold only applies + // to a SET_ROWS of the snapshot tail, whose per-row width is the full state + // D = S_v*S_v*H_v and whose row count is min(T, K)*n_seqs. + const int64_t S_v = gdn->src[2]->ne[0]; // value head dim + const int64_t H_v = gdn->src[2]->ne[2]; // value heads + const int64_t n_seqs = gdn->src[2]->ne[3]; + const int64_t T = gdn->src[0]->ne[2]; // tokens this step + const int64_t K = (int64_t) ggml_get_op_params_i32(gdn, 0); + const int64_t D = S_v * S_v * H_v; + const int64_t n_slots = (T < K ? T : K); // snapshot slots written + const int64_t n_write = n_slots * n_seqs; + // byte offset of the snapshot tail within the GDN output, matching the + // kernel: dst = base + attn_size + (K - n_slots)*state_size_per_snap. + const int64_t attn_size = T * H_v * S_v * n_seqs; + const int64_t state_size_per_snap = D * n_seqs; + const int64_t snap_off_elems = attn_size + (K - n_slots) * state_size_per_snap; + for (int j = idx + 1; j < ctx->n_nodes(); ++j) { ggml_tensor * set_rows = ctx->node(j); if (set_rows->op != GGML_OP_SET_ROWS || set_rows->src[0] == nullptr) { @@ -1648,6 +1668,32 @@ static int ggml_metal_gdn_write_rows( continue; } + // Descent from the GDN output is necessary but NOT sufficient: a caller + // could scatter a differently-shaped view, or a same-sized view at a + // different offset (e.g. an attention-output slice). Verify the fold + // target is exactly the snapshot tail -- per-row state width, index + // count, destination row width, AND that the view begins at the + // snapshot-tail byte offset within the GDN output (tensors are + // allocated at encode time, so the data pointers are valid here). + // The fused epilogue scatters the CONTIGUOUS snapshot tail, but + // ggml_set_rows only requires contiguous rows (nb[0]); it permits an + // arbitrary row stride nb[1] that its own kernel would honor. Require + // the compact [D, n_write] layout (row width D, unit element stride, + // row stride == D) so a strided view is left to the real SET_ROWS. + const ggml_tensor * view = set_rows->src[0]; + const size_t ts = ggml_type_size(view->type); + if (ggml_nelements(view) != D * n_write || + view->ne[0] != D || view->nb[0] != ts || view->nb[1] != (size_t) D * ts || + set_rows->src[1]->ne[0] != n_write || + set_rows->src[2]->ne[0] != D) { + continue; + } + if (view->data == nullptr || gdn->data == nullptr || + (size_t) ((const char *) view->data - (const char *) gdn->data) != + (size_t) snap_off_elems * sizeof(float)) { + continue; + } + *write_rows = set_rows->src[1]; *state_dst = set_rows->src[2]; *fused_set_rows = set_rows; diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index a8eb28d9eecc..b3d9863762e0 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -2899,7 +2899,18 @@ ggml_tensor * llm_graph_context::build_rs_cache_view( ggml_tensor * states = ggml_reshape_2d(ctx0, s, state_size, s->ne[1]); // same cache hygiene as build_rs, minus the main gather (the consumer reads - // per-seq rows via inp->s_copy_main directly) + // per-seq rows via inp->s_copy_main directly, inside the GDN op). + // + // KNOWN LIMITATION (tracked follow-up): build_rs gathers the main rows + // BEFORE this extra relocation, so an overlapping main row is read before + // being overwritten. rows mode defers the main read into the consumer, and + // s_copy() maps a main row to an arbitrary cache slot (idx*size + src0), + // which can fall inside the extra destination [rs_head+n_seqs, rs_head+n_rs) + // during a cache reorder -- so this relocation could clobber a main row the + // consumer will later read. Not reachable on the current single-sequence + // decode path, but it is a real multi-sequence hazard; the correct fix is + // to order the relocation AFTER the GDN read (build_rs's read-before-write + // ordering), which is a graph-dependency refactor left as follow-up. ggml_tensor * state_zero = ggml_view_1d(ctx0, states, state_size*(rs_zero >= 0), rs_zero*states->nb[1]*(rs_zero >= 0)); ggml_build_forward_expand(gf, ggml_scale_inplace(ctx0, state_zero, 0)); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 18b87d63c26f..7512e341f7e1 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -3866,16 +3866,20 @@ struct test_gated_delta_net : public test_case { const bool kda; const int64_t K; // snapshot slot count: 1 = final-only, >1 = last K states const bool rows_mode; // rows-indexed state read from a 2D cache view (src[6]) + const int64_t cache_rows; // rows-mode cache row count (-1 => n_seqs + 3) std::string vars() override { - return VARS_TO_STR10(type, head_count, head_size, n_seq_tokens, n_seqs, v_repeat, permuted, kda, K, rows_mode); + return VARS_TO_STR11(type, head_count, head_size, n_seq_tokens, n_seqs, v_repeat, permuted, kda, K, rows_mode, cache_rows); } + int64_t n_cache_rows() const { return cache_rows > 0 ? cache_rows : n_seqs + 3; } + test_gated_delta_net(ggml_type type = GGML_TYPE_F32, int64_t head_count = 4, int64_t head_size = 16, int64_t n_seq_tokens = 1, int64_t n_seqs = 1, - int v_repeat = 1, bool permuted = false, bool kda = false, int64_t K = 1, bool rows_mode = false) + int v_repeat = 1, bool permuted = false, bool kda = false, int64_t K = 1, bool rows_mode = false, + int64_t cache_rows = -1) : type(type), head_count(head_count), head_size(head_size), n_seq_tokens(n_seq_tokens), n_seqs(n_seqs), - v_repeat(v_repeat), permuted(permuted), kda(kda), K(K), rows_mode(rows_mode) {} + v_repeat(v_repeat), permuted(permuted), kda(kda), K(K), rows_mode(rows_mode), cache_rows(cache_rows) {} ggml_tensor * build_graph(ggml_context * ctx) override { ggml_tensor * q; @@ -3907,7 +3911,7 @@ struct test_gated_delta_net : public test_case { // 2D cache view with more rows than sequences; per-seq state rows // are picked via the I32 rows tensor (see initialize_tensors) const int64_t D = head_size * v_repeat * head_size * head_count; - ggml_tensor * states = ggml_new_tensor_2d(ctx, type, D, n_seqs + 3); + ggml_tensor * states = ggml_new_tensor_2d(ctx, type, D, n_cache_rows()); ggml_tensor * rows = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_seqs); ggml_set_name(states, "state"); ggml_set_name(rows, "rows"); @@ -3933,7 +3937,7 @@ struct test_gated_delta_net : public test_case { // deterministic, distinct, in-range cache rows (stride 2 over n_seqs+3) std::vector idx(t->ne[0]); for (int64_t i = 0; i < t->ne[0]; i++) { - idx[i] = (int32_t) ((i*2 + 1) % (t->ne[0] + 3)); + idx[i] = (int32_t) ((i*2 + 1) % n_cache_rows()); } ggml_backend_tensor_set(t, idx.data(), 0, idx.size()*sizeof(int32_t)); } else { @@ -9172,6 +9176,9 @@ 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)); // rows mode: state read directly from a 2D cache view at rows[seq] (src[6]) + test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 1, 1, 1, false, false, /*K=*/1, /*rows=*/true)); // rows-mode K==1 final-state branch + test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 1, 1, 1, false, false, /*K=*/2, /*rows=*/true, /*cache_rows=*/1)); // 1-row cache + K>1: CPU workspace-sizing regression + test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 1, 2, 1, false, false, /*K=*/1, /*rows=*/true)); // rows-mode K==1, multi-seq test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 32, 1, 1, 1, false, false, /*K=*/2, /*rows=*/true)); test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 4, 2, 1, false, false, /*K=*/4, /*rows=*/true)); test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 8, 128, 4, 1, 1, false, false, /*K=*/4, /*rows=*/true));