llama: in-place recurrent (GDN) state update for multi-sequence decode + lazily zeroed RS buffer

ggml/include/ggml.h

@@ -2620,6 +2620,31 @@ extern "C" {
             struct ggml_tensor  * rows,
             int                   n_snap_slots);
 
+    // cache-writing form (CPU backend only): like ggml_gated_delta_net_rows, but
+    // the op MUTATES `states` instead of returning state planes -- the result is
+    // the attention scores only, [S_v*H_v, n_tokens*n_seqs]. For sequence i the
+    // kernel reads cache row rows_in[i] (I32, n_seqs entries), updates the live
+    // state in place in row rows_out[0*n_seqs + i] (I64, min(n_tokens, K)*n_seqs
+    // entries; no copy at all when rows_in[i] == rows_out[i]) and writes snapshot
+    // slot s (state s tokens back) for s in 1..min(n_tokens, K)-1 straight into
+    // row rows_out[s*n_seqs + i]. The caller guarantees that rows_in[i] and every
+    // rows_out[s*n_seqs + i] are planes of ONE cell (row % mem_size equal, checked
+    // by the kernel) and that the cells of different sequences are pairwise
+    // distinct, so no work unit ever reads a row another unit writes. Removes the
+    // gather, the dst state planes and the write-back copy from recurrent decode.
+    GGML_API struct ggml_tensor * ggml_gated_delta_net_cache(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * q,
+            struct ggml_tensor  * k,
+            struct ggml_tensor  * v,
+            struct ggml_tensor  * g,
+            struct ggml_tensor  * beta,
+            struct ggml_tensor  * states,   // 2D cache view [S_v*S_v*H_v, n_rows], mutated
+            struct ggml_tensor  * rows_in,  // I32 [n_seqs]
+            struct ggml_tensor  * rows_out, // I64 [min(n_tokens, K)*n_seqs]
+            int                   n_snap_slots,
+            int                   mem_size);
+
     // fold the per-head gate activations into a gated_delta_net op (scalar gate only):
     //   beta -> sigmoid(beta),  g -> a[h] * softplus(g + dt_bias[h])
     // dt_bias and a are F32 with H_v elements; removes four elementwise ops per layer

ggml/src/ggml-cpu/ggml-cpu.c

@@ -3011,7 +3011,9 @@ struct ggml_cplan ggml_graph_plan(
                     {
                         const int64_t S_v = node->src[2]->ne[0];
                         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);
+                        // cache-write mode (op_params[2]) computes on the cache row: no state scratch
+                        const bool cache_write = ggml_get_op_params_i32(node, 2) != 0;
+                        const int64_t per_thread = S_v + ((K > 1 && !cache_write) ? S_v * S_v : 0);
                         cur = per_thread * sizeof(float) * n_tasks;
                     } break;
                 case GGML_OP_COUNT:

ggml/src/ggml-cpu/ops.cpp

@@ -10819,11 +10819,26 @@ static void ggml_compute_forward_gated_delta_net_one_chunk(
     const int64_t state_seq_stride = src_rows ? 0 : (int64_t) (src_state->nb[3] / sizeof(float));
     const int64_t state_row_size   = src_rows ? (int64_t) (src_state->nb[1] / sizeof(float)) : 0;
 
-    const int64_t per_thread = S_v + (K > 1 ? S_v * S_v : 0);
+    // cache-write mode (op_params[2], ggml_gated_delta_net_cache): the live state is
+    // updated in place in cache row rows_out[seq] and snapshot slots 1..K-1 are
+    // written straight into cache rows rows_out[slot*n_seqs + seq]; the result holds
+    // the attention scores only (state_out_base is unused in this mode).
+    const bool      cache_write = ggml_get_op_params_i32(dst, 2) != 0;
+    GGML_ASSERT(!cache_write || state_rows_idx != nullptr);
+    const int64_t * rows_out = cache_write ? (const int64_t *) dst->src[9]->data : nullptr;
+    float *         cache    = cache_write ? (float *) src_state->data : nullptr;
+    const int64_t   mem_size = cache_write ? ggml_get_op_params_i32(dst, 3) : 1;
+
+    // the legacy K > 1 path stages the state through a per-thread scratch; the
+    // cache-write path never does (it computes on the cache row itself).
+    // ggml-cpu.c (graph plan work size) mirrors this expression.
+    const bool use_state_work = K > 1 && !cache_write;
+
+    const int64_t per_thread = S_v + (use_state_work ? S_v * S_v : 0);
     const int ith = params->ith;
 
     float * delta       = (float *)params->wdata + ith * per_thread + CACHE_LINE_SIZE_F32;
-    float * state_work  = K > 1 ? (delta + S_v) : nullptr;
+    float * state_work  = use_state_work ? (delta + S_v) : nullptr;
 
     // output layout: [attn_scores | new_states]
     // attn_scores: S_v * H * n_tokens * n_seqs    floats
@@ -10855,19 +10870,34 @@ static void ggml_compute_forward_gated_delta_net_one_chunk(
         const int64_t iq3 = iv3 / rq3;
         const int64_t ik3 = iv3 / rk3;
 
-        // For K=1, write directly to the single output slot to avoid an extra memcpy at the end.
-        // For K>1, work in scratch and copy out per-token when the slot is in range.
-        float * s_out = (K > 1)
-            ? state_work
-            : state_out_base + (iv3 * H + iv1) * S_v * S_v;
-
-        // copy input state into the working buffer and operate in-place.
-        // scratch mode: state layout [S_v, S_v, H, n_seqs], seq iv3 starts at
-        // iv3 * state_seq_stride. rows mode: cache row state_rows_idx[iv3].
+        // input state: scratch mode: state layout [S_v, S_v, H, n_seqs], seq iv3
+        // starts at iv3 * state_seq_stride. rows mode: cache row state_rows_idx[iv3].
         const float * s_in = state_rows_idx
             ? state_in_base + (int64_t) state_rows_idx[iv3] * state_row_size + iv1 * S_v * S_v
             : state_in_base + iv3 * state_seq_stride + iv1 * S_v * S_v;
-        memcpy(s_out, s_in, S_v * S_v * sizeof(float));
+
+        float * s_out;
+        if (cache_write) {
+            // operate directly on the live cache row rows_out[iv3]. rows_in[iv3] is
+            // either that same row (steady state: no copy at all) or another plane
+            // of the same cell (pending rollback / freshly zeroed cell): copy it in
+            // first -- distinct rows, so plain memcpy is safe.
+            const int64_t row_out = rows_out[iv3];
+            GGML_ASSERT(row_out % mem_size == (int64_t) state_rows_idx[iv3] % mem_size);
+            s_out = cache + row_out * state_row_size + iv1 * S_v * S_v;
+            if (s_out != s_in) {
+                memcpy(s_out, s_in, S_v * S_v * sizeof(float));
+            }
+        } else {
+            // For K=1, write directly to the single output slot to avoid an extra memcpy at the end.
+            // For K>1, work in scratch and copy out per-token when the slot is in range.
+            s_out = (K > 1)
+                ? state_work
+                : state_out_base + (iv3 * H + iv1) * S_v * S_v;
+
+            // copy input state into the working buffer and operate in-place.
+            memcpy(s_out, s_in, S_v * S_v * sizeof(float));
+        }
 
         // attn output pointer for first token of this (head, seq)
         float * attn_data = attn_out_base + (iv3 * n_tokens * H + iv1) * S_v;
@@ -10924,7 +10954,20 @@ static void ggml_compute_forward_gated_delta_net_one_chunk(
 
             attn_data += S_v * H; // advance to next token
 
-            if (K > 1) {
+            if (cache_write) {
+                // snapshot slot s (state s tokens back) -> cache row rows_out[s*n_seqs + seq].
+                // slot 0 is s_out itself (already in place). slot <= n_tokens-1 and
+                // slot < K keep the index below min(n_tokens, K)*n_seqs. Within this
+                // unit the rollback source plane was fully copied into s_out before
+                // the token loop, so overwriting it here is read-before-write safe.
+                if (K > 1) {
+                    const int64_t slot = n_tokens - 1 - t;
+                    if (slot >= 1 && slot < K) {
+                        float * snap = cache + rows_out[slot * n_seqs + iv3] * state_row_size + iv1 * S_v * S_v;
+                        memcpy(snap, s_out, S_v * S_v * sizeof(float));
+                    }
+                }
+            } else if (K > 1) {
                 const int64_t target_slot = n_tokens - 1 - t;
                 if (target_slot >= 0 && target_slot < K) {
                     float * curr_state_o = state_out_base + target_slot * state_size_per_snap +

ggml/src/ggml.c

@@ -6423,6 +6423,75 @@ struct ggml_tensor * ggml_gated_delta_net_rows(
     return result;
 }
 
+struct ggml_tensor * ggml_gated_delta_net_cache(
+        struct ggml_context * ctx,
+        struct ggml_tensor  * q,
+        struct ggml_tensor  * k,
+        struct ggml_tensor  * v,
+        struct ggml_tensor  * g,
+        struct ggml_tensor  * beta,
+        struct ggml_tensor  * states,
+        struct ggml_tensor  * rows_in,
+        struct ggml_tensor  * rows_out,
+        int                   n_snap_slots,
+        int                   mem_size) {
+    GGML_ASSERT(ggml_is_contiguous_rows(q));
+    GGML_ASSERT(ggml_is_contiguous_rows(k));
+    GGML_ASSERT(ggml_is_contiguous_rows(v));
+    GGML_ASSERT(ggml_is_contiguous(g));
+    GGML_ASSERT(ggml_is_contiguous(beta));
+    GGML_ASSERT(ggml_is_contiguous(states));
+    GGML_ASSERT(ggml_is_contiguous(rows_in));
+    GGML_ASSERT(ggml_is_contiguous(rows_out));
+
+    GGML_ASSERT(q->type == GGML_TYPE_F32);
+    GGML_ASSERT(k->type == GGML_TYPE_F32);
+    GGML_ASSERT(v->type == GGML_TYPE_F32);
+    GGML_ASSERT(g->type == GGML_TYPE_F32);
+    GGML_ASSERT(beta->type == GGML_TYPE_F32);
+    GGML_ASSERT(states->type == GGML_TYPE_F32);
+    GGML_ASSERT(rows_in->type == GGML_TYPE_I32);
+    GGML_ASSERT(rows_out->type == GGML_TYPE_I64);
+
+    const int64_t S_v      = v->ne[0];
+    const int64_t H        = v->ne[1];
+    const int64_t n_tokens = v->ne[2];
+    const int64_t n_seqs   = v->ne[3];
+
+    GGML_ASSERT(g->ne[0] == 1 || g->ne[0] == S_v);
+    GGML_ASSERT(beta->ne[0] == 1);
+
+    // states is a 2D cache view (D, n_rows); each row is one sequence's full state
+    GGML_ASSERT(states->ne[0] == S_v * S_v * H);
+    GGML_ASSERT(rows_in->ne[0] == n_seqs);
+
+    const int64_t K = n_snap_slots;
+    GGML_ASSERT(K >= 1);
+    GGML_ASSERT(mem_size >= 1);
+    GGML_ASSERT(rows_out->ne[0] == MIN(n_tokens, K) * n_seqs);
+
+    // the state stays in the cache: the result carries the attention scores only
+    const int64_t ne[4] = { S_v * H, n_tokens * n_seqs, 1, 1 };
+    struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne);
+
+    ggml_set_op_params_i32(result, 0, (int32_t) K);
+    ggml_set_op_params_i32(result, 2, 1);        // cache-write flag (op_params[1] is the raw-gates bit)
+    ggml_set_op_params_i32(result, 3, mem_size); // kernel same-cell check: row % mem_size
+
+    result->op     = GGML_OP_GATED_DELTA_NET;
+    result->src[0] = q;
+    result->src[1] = k;
+    result->src[2] = v;
+    result->src[3] = g;
+    result->src[4] = beta;
+    result->src[5] = states;
+    result->src[6] = rows_in;
+    // src[7] and src[8] are reserved for ggml_gated_delta_net_set_raw_gates
+    result->src[9] = rows_out;
+
+    return result;
+}
+
 void ggml_gated_delta_net_set_raw_gates(
         struct ggml_tensor  * gdn,
         struct ggml_tensor  * dt_bias,

src/llama-graph.cpp

@@ -417,6 +417,11 @@ bool llm_graph_input_rs::can_reuse_rs(const llama_memory_recurrent_context * mct
     res &= head == mctx_cur->get_head();
     res &= rs_z == mctx_cur->get_rs_z();
 
+    // a graph built for the in-place GDN path must not be reused by a ubatch
+    // whose cells were reordered / freshly zeroed (and vice versa); the memory
+    // was already updated by apply() -> find_slot before can_reuse runs
+    res &= rs_inplace == mctx_cur->rs_inplace_ok(params.ubatch.n_seqs);
+
     return res;
 }
 
@@ -3573,6 +3578,8 @@ static std::unique_ptr<llm_graph_input_rs> build_rs_inp_impl(
     inp->head = mctx_cur->get_head();
     inp->rs_z = mctx_cur->get_rs_z();
 
+    inp->rs_inplace = mctx_cur->rs_inplace_ok(n_seqs);
+
     return inp;
 }
 
@@ -3601,28 +3608,38 @@ ggml_tensor * llm_graph_context::build_rs_cache_view(
         llm_graph_input_rs * inp,
         ggml_tensor * s,
             int32_t   state_size,
-            int32_t   n_seqs) const {
+            int32_t   n_seqs,
+               bool   inplace) const {
     const auto * kv_state = inp->mctx;
 
     const uint32_t n_rs     = kv_state->get_n_rs();
     const uint32_t rs_head  = kv_state->get_head();
     const  int32_t rs_zero  = kv_state->get_rs_z();
 
+    // the in-place consumer mutates the cache rows of the ubatch: it is only
+    // entered when every main row is the identity and the extra range below is
+    // empty (rs_inplace_ok), which makes the ordering hazard described next
+    // unreachable for it
+    GGML_ASSERT(!inplace || (inp->rs_inplace && n_rs == (uint32_t) n_seqs));
+
     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, 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.
+    // KNOWN LIMITATION of the Metal rows mode (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. The CPU in-place
+    // mode never enters this function with a non-empty extra range (see the
+    // assert above; transition ubatches take the gathered build_rs path), so the
+    // hazard only remains for the Metal rows + SET_ROWS mode on a cache reorder;
+    // the correct fix there is to order the relocation AFTER the GDN read
+    // (build_rs's read-before-write ordering), 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));
 

src/llama-graph.h

@@ -346,6 +346,10 @@ class llm_graph_input_rs : public llm_graph_input_i {
     ggml_tensor * s_write_rows = nullptr;
     int64_t       s_write_K    = 0;
 
+    // every ubatch seq reads its own cell and n_rs == n_seqs: the GDN op may
+    // update the cache rows in place (rs_inplace_ok); part of the reuse key
+    bool rs_inplace = false;
+
     const llama_memory_recurrent_context * mctx;
 
     // used in view offsets, need to match for valid graph reuse
@@ -1437,11 +1441,14 @@ struct llm_graph_context {
     // (state_size, n_rows) cache view. For consumers that read per-seq state
     // rows directly via inp->s_copy_main (e.g. ggml_gated_delta_net_rows),
     // saving a get_rows per layer per decode.
+    // inplace: the consumer is the cache-writing op (ggml_gated_delta_net_cache);
+    // requires inp->rs_inplace, i.e. an empty extra range and identity main rows.
     ggml_tensor * build_rs_cache_view(
             llm_graph_input_rs * inp,
             ggml_tensor * s,
                 int32_t   state_size,
-                int32_t   n_seqs) const;
+                int32_t   n_seqs,
+                   bool   inplace = false) const;
 
     // I64 cache-row indices for the rows-mode snapshot scatter: row of
     // (slot s, seq i) = s*mem_size + head + i, matching the strided-cpy