metal gdn: review follow-ups that missed #61 (CPU workspace overflow, write-fold guards, K==1 tests)

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;

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;

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

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<int32_t> 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<std::unique_ptr<test_case>> 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));