From 8a27b212d4e15814ac504c7ef066a2478421a245 Mon Sep 17 00:00:00 2001 From: abdulaziz Date: Sat, 19 Sep 2026 01:55:08 +0300 Subject: [PATCH 01/15] gdn: in-place cache-writing GDN op for CPU recurrent decode (K == 1 gate) New op form ggml_gated_delta_net_cache (op_params[2] = 1, src[6] = rows_in I32[n_seqs], src[9] = rows_out I64[min(T,K)*n_seqs]): the CPU kernel reads each sequence's state at cache row rows_in[seq], updates the live row rows_out[seq] in place (no copy when the rows coincide) and writes snapshot slots 1..min(T,K)-1 straight into the cache; the result carries the attention scores only. No get_rows gather, no dst state planes, no write-back cpy / set_rows per layer: 6 -> 2 state-sized moves per seq. Graph side: llama_memory_recurrent_context::rs_inplace_ok(n_seqs), a side-effect-free predicate (never calls s_copy, which consumes rs_idx): !is_full && n == n_seqs && cells[head+i].src0 == head+i. Stored in llm_graph_input_rs::rs_inplace and added to can_reuse_rs so a reused graph can never be in the wrong mode. build_rs_cache_view(inplace) asserts the extra relocation range is empty, which makes the documented rows-mode ordering hazard unreachable for the in-place consumer. Model side (qwen35): gdn_state_inplace_dev_ok (no GPU/IGPU device), the gate gdn_state_rows_env && inplace_dev_ok && inp->rs_inplace && fused_ok && cparams.n_rs_seq == 0 (Stage 7A: K == 1 only), the Metal rows+SET_ROWS mode kept for Metal-only machines; build_recurrent_attn gains the state_inplace branch, every legacy branch is untouched and serves the transition ubatches (join/leave/seq_cp/reorder) and the env switch. GGML_GDN_STATE_GATHER=1 still forces the legacy gathered path (A/B). Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01QrdPnP5KL449GsoWBZx52R --- ggml/include/ggml.h | 25 ++++++++++++ ggml/src/ggml-cpu/ggml-cpu.c | 4 +- ggml/src/ggml-cpu/ops.cpp | 69 +++++++++++++++++++++++++++------- ggml/src/ggml.c | 69 ++++++++++++++++++++++++++++++++++ src/llama-graph.cpp | 39 +++++++++++++------ src/llama-graph.h | 9 ++++- src/llama-memory-recurrent.cpp | 27 +++++++++++++ src/llama-memory-recurrent.h | 8 ++++ src/models/delta-net-base.cpp | 47 ++++++++++++++++++++++- src/models/models.h | 12 ++++-- src/models/qwen35.cpp | 41 +++++++++++++++----- 11 files changed, 309 insertions(+), 41 deletions(-) diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index d875a15ed066..7d518edf64c9 100644 --- a/ggml/include/ggml.h +++ b/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 diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index 9cd2c2d39a5a..3a6e21591542 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/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: diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 4fc7514bb3fb..83a616e101a6 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/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 + diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index fa2a9c44f5cf..c233176541b3 100644 --- a/ggml/src/ggml.c +++ b/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, diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 99f263e32bd7..002c599eddb2 100644 --- a/src/llama-graph.cpp +++ b/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 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)); diff --git a/src/llama-graph.h b/src/llama-graph.h index e85c1aa20d1f..c3b7e13af6ba 100644 --- a/src/llama-graph.h +++ b/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 diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index e2990972ef7e..7ac65695b8ff 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -1270,3 +1270,30 @@ int32_t llama_memory_recurrent_context::s_copy(int i) const { } return (int32_t)(idx * mem->size) + src0; } + +bool llama_memory_recurrent_context::rs_inplace_ok(uint32_t n_seqs) const { + // the reserve context spans the whole cache (head = 0, n_rs = size) + if (is_full) { + return false; + } + + // the extra relocation range [head + n_seqs, head + n_rs) must be empty + if (mem->n != n_seqs) { + return false; + } + + // every ubatch cell must read its own state. After find_slot this holds for + // every cell that already owned its state (src0 = src = own index) and for + // the one-fresh-seq case where rs_z is the seq's own cell (the graph zeroes + // that row before the op reads it). It fails on a seq_cp fan-out, on a cell + // reorder and when two fresh seqs share rs_z -- those ubatches take the + // gathered path. + for (uint32_t i = 0; i < n_seqs; ++i) { + const uint32_t cell_idx = mem->head + i; + if (mem->cells[cell_idx].src0 != (int32_t) cell_idx) { + return false; + } + } + + return true; +} diff --git a/src/llama-memory-recurrent.h b/src/llama-memory-recurrent.h index b13b7b748f5e..83717a67e3a9 100644 --- a/src/llama-memory-recurrent.h +++ b/src/llama-memory-recurrent.h @@ -173,6 +173,14 @@ class llama_memory_recurrent_context : public llama_memory_context_i { int32_t s_copy(int i) const; + // true when the GDN op may update this ubatch's cache rows in place: not the + // full (reserve) context, the cell range is exactly the ubatch (n == n_seqs) + // and every cell reads its own state (src0 == own index; a pending rollback + // still reads a plane of the same cell). Side-effect free on purpose: unlike + // s_copy() it never consumes the per-seq rollback index, so it can be + // evaluated at graph build and again in can_reuse. + bool rs_inplace_ok(uint32_t n_seqs) const; + private: const llama_memory_status status; diff --git a/src/models/delta-net-base.cpp b/src/models/delta-net-base.cpp index ae67400c6131..26799679f6af 100644 --- a/src/models/delta-net-base.cpp +++ b/src/models/delta-net-base.cpp @@ -539,7 +539,8 @@ ggml_tensor * llm_build_delta_net_base::build_recurrent_attn( ggml_tensor * b, ggml_tensor * s, int il, - ggml_tensor * state_rows) { + ggml_tensor * state_rows, + bool state_inplace) { const auto * mctx_cur = inp->mctx; const auto kv_head = mctx_cur->get_head(); const uint32_t mem_size = mctx_cur->get_size(); @@ -553,7 +554,49 @@ ggml_tensor * llm_build_delta_net_base::build_recurrent_attn( const bool keep = cparams.n_rs_seq > 0; - GGML_ASSERT(state_rows == nullptr || keep); // rows mode is a ring-path optimization + GGML_ASSERT(state_rows == nullptr || keep || state_inplace); // Metal rows mode is a ring-path optimization + GGML_ASSERT(!state_inplace || state_rows != nullptr); + + if (state_inplace) { + // in-place mode (CPU): the cache-writing op reads each seq's state at + // cache row state_rows[seq] (rows_in = inp->s_copy_main: the live row, or + // the rollback plane of the same cell), updates the live row + // write_rows[seq] in place and writes snapshot slots 1..min(T, K)-1 + // straight into cache rows write_rows[slot*n_seqs + seq] -- the same + // rows the SET_ROWS / strided cpy below would have written. The result + // carries the attention scores only: no gather, no dst state planes, no + // write-back copy. Same float instruction sequence on the same inputs as + // the paths below, so the outputs and the resulting cache rows are + // bit-identical (tests/test-gdn-inplace.cpp). + const int64_t K = cparams.n_rs_seq + 1; + + const bool raw = gdn_raw_beta && gdn_raw_alpha && gdn_raw_dt_bias && gdn_raw_a; + ggml_tensor * gg = raw ? gdn_raw_alpha : g; + ggml_tensor * bb = raw ? gdn_raw_beta : b; + + ggml_tensor * write_rows = build_rs_write_rows(inp, K, n_seq_tokens, n_seqs); + + ggml_tensor * gdn_out = ggml_gated_delta_net_cache(ctx0, q, k, v, gg, bb, s, state_rows, write_rows, + (int) K, (int) mem_size); + if (raw) { + ggml_gated_delta_net_set_raw_gates(gdn_out, gdn_raw_dt_bias, gdn_raw_a); + } + if (n_seq_tokens > 1) { + res->add_fused_node({LLM_FUSED_OP_GDN_CH, gdn_out, il}); + } else { + res->add_fused_node({LLM_FUSED_OP_GDN_AR, gdn_out, il}); + } + + ggml_tensor * output = ggml_view_4d(ctx0, 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); + cb(output, "attn_output", il); + + return output; + } if (!keep) { auto attn_out = build_delta_net(q, k, v, g, b, s, il); diff --git a/src/models/models.h b/src/models/models.h index a82133099d5c..1ae68a28edea 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -84,6 +84,10 @@ struct llm_build_delta_net_base : public llm_graph_context { // live state directly at cache row state_rows[seq] (inp->s_copy_main) -- // no gathered scratch; the snapshot write becomes a SET_ROWS the Metal // backend can fold into the fused op's epilogue. + // state_inplace (CPU only, requires state_rows): the cache-writing op + // ggml_gated_delta_net_cache updates the live rows in place and writes the + // snapshot slots straight into the cache -- no dst state planes, no SET_ROWS, + // no write-back cpy. Only valid on ubatches with inp->rs_inplace. // set per layer before build_recurrent_attn: the fused GDN op then receives the // pre-activation beta / alpha and folds sigmoid / softplus into its prologue // (ggml_gated_delta_net_set_raw_gates); the non-fused paths keep the activated g / b @@ -102,7 +106,8 @@ struct llm_build_delta_net_base : public llm_graph_context { ggml_tensor * b, ggml_tensor * s, int il, - ggml_tensor * state_rows = nullptr); + ggml_tensor * state_rows = nullptr, + bool state_inplace = false); }; struct llm_build_rwkv6_base : public llm_graph_context { @@ -2267,8 +2272,9 @@ struct llama_model_qwen35 : public llama_model_base { graph(const llama_model & model, const llm_graph_params & params); // device-dependent path choices, scanned once per graph build (not per layer) - bool gdn_state_rows_dev_ok = true; // every GPU device is Metal: fused GDN may read state rows in place - bool gdn_raw_gates_dev_ok = true; // every device is CPU/Metal/CUDA/ROCm/MUSA: fused GDN takes raw gates + bool gdn_state_rows_dev_ok = true; // every GPU device is Metal: fused GDN may read state rows in place + bool gdn_state_inplace_dev_ok = true; // no GPU/IGPU device at all: the in-place (cache-writing) GDN is CPU-only + bool gdn_raw_gates_dev_ok = true; // every device is CPU/Metal/CUDA/ROCm/MUSA: fused GDN takes raw gates private: ggml_tensor * build_layer_attn( llm_graph_input_attn_kv * inp_attn, diff --git a/src/models/qwen35.cpp b/src/models/qwen35.cpp index c5f816e96268..89ab318a4e89 100644 --- a/src/models/qwen35.cpp +++ b/src/models/qwen35.cpp @@ -149,13 +149,18 @@ llama_model_qwen35::graph::graph(const llama_model & model, const llm_graph_para ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(ldev.dev); const char * reg_name = reg ? ggml_backend_reg_name(reg) : nullptr; if (reg_name == nullptr) { - gdn_state_rows_dev_ok = false; - gdn_raw_gates_dev_ok = false; + gdn_state_rows_dev_ok = false; + gdn_state_inplace_dev_ok = false; + gdn_raw_gates_dev_ok = false; break; } // integrated GPUs (e.g. unified-memory CUDA devices) report IGPU, not GPU const bool is_gpu = ggml_backend_dev_type(ldev.dev) == GGML_BACKEND_DEVICE_TYPE_GPU || ggml_backend_dev_type(ldev.dev) == GGML_BACKEND_DEVICE_TYPE_IGPU; + if (is_gpu) { + // the cache-writing GDN op mutates one of its sources; only the CPU backend implements it + gdn_state_inplace_dev_ok = false; + } if (is_gpu && strcmp(reg_name, "MTL") != 0) { gdn_state_rows_dev_ok = false; } @@ -472,18 +477,34 @@ ggml_tensor * llama_model_qwen35::graph::build_layer_attn_linear( // ring path: read per-seq live state directly from the cache inside the // fused GDN op (rows mode) instead of a gather per layer. - // GGML_GDN_STATE_GATHER=1 restores the legacy gathered path (A/B). - // rows mode (the src[6] variant) is implemented on CPU and Metal only; - // other GPU backends reject it in supports_op, which would silently move - // the whole recurrent op to CPU -- keep the gathered form unless every - // GPU device in the model is Metal. + // GGML_GDN_STATE_GATHER=1 restores the legacy gathered path (A/B) for both + // rows modes below. + // - in-place (CPU only, ggml_gated_delta_net_cache): the op updates the + // cache rows themselves -- no gather, no dst state planes, no write-back. + // Taken on every ubatch whose sequences all read their own cell + // (inp->rs_inplace); transition ubatches (a seq joins / leaves / is + // copied / reordered) fall back to the gathered path, which is exact and + // ordering-safe. Part of the graph reuse key (can_reuse_rs). + // - Metal rows + SET_ROWS (the src[6] variant, ring path only): kept for + // machines whose GPU devices are all Metal; other GPU backends reject it + // in supports_op, which would silently move the whole recurrent op to CPU. static const bool gdn_state_rows_env = getenv("GGML_GDN_STATE_GATHER") == nullptr; - const bool gdn_state_rows = gdn_state_rows_env && gdn_state_rows_dev_ok && cparams.n_rs_seq > 0; + // the in-place op always runs the fused kernel: only take it where the legacy + // path runs the fused kernel too (the ring path always does; the K == 1 path + // only with the fused_gdn_* probes on), so the A/B stays bit-identical + const bool fused_ok = cparams.n_rs_seq > 0 || (n_seq_tokens == 1 ? cparams.fused_gdn_ar : cparams.fused_gdn_ch); + + // Stage 7A gate: in-place only without speculation (K == 1) + const bool gdn_state_inplace = gdn_state_rows_env && gdn_state_inplace_dev_ok && inp->rs_inplace && fused_ok && + cparams.n_rs_seq == 0; + const bool gdn_state_rows_metal = gdn_state_rows_env && gdn_state_rows_dev_ok && !gdn_state_inplace_dev_ok && + cparams.n_rs_seq > 0; + const bool gdn_state_rows = gdn_state_inplace || gdn_state_rows_metal; ggml_tensor * state; if (gdn_state_rows) { - state = build_rs_cache_view(inp, ssm_states_all, hparams.n_embd_s(), n_seqs); + state = build_rs_cache_view(inp, ssm_states_all, hparams.n_embd_s(), n_seqs, gdn_state_inplace); cb(state, "state_cache_view", il); } else { state = build_rs(inp, ssm_states_all, hparams.n_embd_s(), n_seqs); @@ -560,7 +581,7 @@ ggml_tensor * llama_model_qwen35::graph::build_layer_attn_linear( cb(v_conv, "v_conv_predelta", il); ggml_tensor * output = build_recurrent_attn(inp, ssm_states_all, q_conv, k_conv, v_conv, gate, beta, state, il, - gdn_state_rows ? inp->s_copy_main : nullptr); + gdn_state_rows ? inp->s_copy_main : nullptr, gdn_state_inplace); // z: [head_dim, n_heads, n_tokens, n_seqs] -> [n_heads * n_tokens * n_seqs, head_dim] ggml_tensor * z_2d = ggml_reshape_4d(ctx0, z, head_v_dim, num_v_heads, n_seq_tokens, n_seqs); From 40e926557865c366738dc6d42a22bb024110f648 Mon Sep 17 00:00:00 2001 From: abdulaziz Date: Sat, 19 Sep 2026 01:59:31 +0300 Subject: [PATCH 02/15] gdn: take the in-place path in MTP speculative mode too (K = n_rs_seq + 1) Drop the cparams.n_rs_seq == 0 restriction of the qwen35 gate: with --spec-type draft-mtp the cache-writing op now updates slot 0 in place on the live row and writes slots 1..min(T, K)-1 straight into the snapshot planes of the same cell (rows_out = the existing s_write_rows mapping slot*mem_size + head + seq); a pending rollback reads plane rs_idx of the same cell (rows_in != rows_out -> copy-then-in-place in the kernel). No dst state planes, no SET_ROWS, no per-thread state scratch. Transition ubatches still take the gathered K-slot path. GGML_GDN_INPLACE_K1_ONLY=1 restores the previous commit's K == 1 gate so the two can be bisected on the same binary; GGML_GDN_STATE_GATHER=1 still forces the legacy gathered path for everything. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01QrdPnP5KL449GsoWBZx52R --- src/models/qwen35.cpp | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/models/qwen35.cpp b/src/models/qwen35.cpp index 89ab318a4e89..073a6e95b63e 100644 --- a/src/models/qwen35.cpp +++ b/src/models/qwen35.cpp @@ -495,9 +495,15 @@ ggml_tensor * llama_model_qwen35::graph::build_layer_attn_linear( // only with the fused_gdn_* probes on), so the A/B stays bit-identical const bool fused_ok = cparams.n_rs_seq > 0 || (n_seq_tokens == 1 ? cparams.fused_gdn_ar : cparams.fused_gdn_ch); - // Stage 7A gate: in-place only without speculation (K == 1) + // GGML_GDN_INPLACE_K1_ONLY=1 restricts the in-place path to K == 1 (no + // speculation); by default the MTP speculative mode (K = n_rs_seq + 1) takes + // it too: slot 0 is updated in place, slots 1..K-1 are written straight into + // the snapshot planes of the same cell, and a pending rollback reads plane + // rs_idx of the same cell (rows_in != rows_out: copy-then-in-place). + static const bool gdn_inplace_k1_only_env = getenv("GGML_GDN_INPLACE_K1_ONLY") != nullptr; + const bool gdn_state_inplace = gdn_state_rows_env && gdn_state_inplace_dev_ok && inp->rs_inplace && fused_ok && - cparams.n_rs_seq == 0; + (cparams.n_rs_seq == 0 || !gdn_inplace_k1_only_env); const bool gdn_state_rows_metal = gdn_state_rows_env && gdn_state_rows_dev_ok && !gdn_state_inplace_dev_ok && cparams.n_rs_seq > 0; const bool gdn_state_rows = gdn_state_inplace || gdn_state_rows_metal; From 8c1b37a7880c1d24e101e538dc58908737ad7bdc Mon Sep 17 00:00:00 2001 From: abdulaziz Date: Sat, 19 Sep 2026 01:59:31 +0300 Subject: [PATCH 03/15] tests: byte-for-byte check of the cache-writing GDN op vs the legacy paths tests/test-gdn-inplace.cpp (llama_build_and_test) runs, on identical synthetic inputs and in three separate ggml contexts, the gathered path (get_rows -> ggml_gated_delta_net(K) -> write-back cpy, exactly as build_recurrent_attn emits it for K == 1 and K > 1), the rows path (ggml_gated_delta_net_rows -> ggml_set_rows) and the new ggml_gated_delta_net_cache, then memcmp's the attention scores and the WHOLE cache tensor (touched rows identical, untouched rows untouched) and checks the in-place op actually modified the cache. Matrix: K in {1, 3}, T in {1, 3}, n_seqs in {1, 4, 8}, rows_in in plane 0 (identity) and in every rollback plane 1..K-1 of the same cell, raw gates on/off, H_k == H_v and H_k == H_v/2, 1 and 4 threads: 192 cases, 0 bytes differ. CPU only, no model. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01QrdPnP5KL449GsoWBZx52R --- tests/CMakeLists.txt | 5 + tests/test-gdn-inplace.cpp | 375 +++++++++++++++++++++++++++++++++++++ 2 files changed, 380 insertions(+) create mode 100644 tests/test-gdn-inplace.cpp diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 67ff970be0b3..17aca7d3bdb4 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -169,6 +169,11 @@ endif () llama_build(test-recurrent-state-rollback.cpp) +# byte-for-byte check of the cache-writing GDN op (ggml_gated_delta_net_cache) +# against the legacy gathered and rows+set_rows paths; CPU only, tiny synthetic +# shapes, no model needed +llama_build_and_test(test-gdn-inplace.cpp) + if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) # these tests are disabled on Windows because they use internal functions not exported with LLAMA_API (when building with shared libraries) llama_build_and_test(test-unicode.cpp) diff --git a/tests/test-gdn-inplace.cpp b/tests/test-gdn-inplace.cpp new file mode 100644 index 000000000000..c52eeb008af7 --- /dev/null +++ b/tests/test-gdn-inplace.cpp @@ -0,0 +1,375 @@ +// Byte-for-byte check of the cache-writing GDN op (ggml_gated_delta_net_cache) +// against the two legacy paths the graph used to emit, on identical inputs: +// +// A (gathered): ggml_get_rows(cache, rows_in) -> ggml_gated_delta_net(K) +// -> write-back cpy of the min(T, K) state planes into the +// cache rows slot*mem_size + head + seq +// (src/models/delta-net-base.cpp, !keep and keep/gathered) +// B (rows): ggml_gated_delta_net_rows(cache view, rows_in, K) +// -> ggml_set_rows(cache, snapshots, write_rows) +// (the Metal rows mode) +// C (in-place): ggml_gated_delta_net_cache(cache view, rows_in, write_rows, K) +// -- the state never leaves the cache +// +// For every case the attention scores of C must equal A and B byte for byte +// and the WHOLE cache tensor after C must equal the cache after A and after +// B byte for byte (touched rows identical, untouched rows untouched). +// +// Matrix: K in {1, 3}, T in {1, 3}, n_seqs in {1, 4, 8}, rows_in in plane 0 +// (identity: rows_in == rows_out) and in plane r in 1..K-1 (rollback read of +// the same cell), raw gates on/off, H_k == H_v and H_k == H_v/2, 1 and 4 +// threads. KDA (vector gate) off, as in the qwen35 graph. +// +// CPU only; runs on tiny synthetic shapes (S_v = 32, H_v = 4) and needs no +// model. + +#include "ggml.h" +#include "ggml-cpu.h" + +#include +#include +#include +#include +#include +#include +#include + +struct gdn_case { + int64_t S_v; + int64_t H_v; + int64_t H_k; + int64_t T; + int64_t n_seqs; + int64_t K; + int64_t mem_size; + int64_t head; + int64_t r_in; // plane read for every seq: rows_in[i] = r_in*mem_size + head + i + bool raw_gates; + int n_threads; +}; + +// deterministic input set shared by the three paths +struct gdn_inputs { + std::vector q, k, v, g, beta, dt_bias, a, cache; + std::vector rows_in; + std::vector rows_out; +}; + +static gdn_inputs make_inputs(const gdn_case & c, uint32_t seed) { + std::mt19937 rng(seed); + auto fill = [&](std::vector & dst, size_t n, float lo, float hi) { + std::uniform_real_distribution d(lo, hi); + dst.resize(n); + for (size_t i = 0; i < n; ++i) { + dst[i] = d(rng); + } + }; + + const int64_t D = c.S_v * c.S_v * c.H_v; + const int64_t n_rows = c.mem_size * c.K; + + gdn_inputs in; + fill(in.q, c.S_v * c.H_k * c.T * c.n_seqs, -1.0f, 1.0f); + fill(in.k, c.S_v * c.H_k * c.T * c.n_seqs, -1.0f, 1.0f); + fill(in.v, c.S_v * c.H_v * c.T * c.n_seqs, -1.0f, 1.0f); + if (c.raw_gates) { + // pre-activation: g -> a[h]*softplus(g + dt_bias[h]), beta -> sigmoid(beta) + fill(in.g, c.H_v * c.T * c.n_seqs, -2.0f, 2.0f); + fill(in.beta, c.H_v * c.T * c.n_seqs, -3.0f, 3.0f); + fill(in.dt_bias, c.H_v, -1.0f, 1.0f); + fill(in.a, c.H_v, -1.0f, -0.1f); + } else { + // activated: g = log decay in (-2, 0), beta in (0, 1) + fill(in.g, c.H_v * c.T * c.n_seqs, -2.0f, 0.0f); + fill(in.beta, c.H_v * c.T * c.n_seqs, 0.0f, 1.0f); + } + // every row of every plane is random so a wrong row read/write shows up + fill(in.cache, D * n_rows, -1.0f, 1.0f); + + in.rows_in.resize(c.n_seqs); + for (int64_t i = 0; i < c.n_seqs; ++i) { + in.rows_in[i] = (int32_t) (c.r_in * c.mem_size + c.head + i); + } + + // set_input_rs mapping: row of (slot s, seq i) = s*mem_size + head + i + const int64_t n_write = std::min(c.T, c.K); + in.rows_out.resize(n_write * c.n_seqs); + for (int64_t r = 0; r < n_write * c.n_seqs; ++r) { + in.rows_out[r] = (r / c.n_seqs) * c.mem_size + c.head + (r % c.n_seqs); + } + + return in; +} + +struct gdn_tensors { + ggml_tensor * q; + ggml_tensor * k; + ggml_tensor * v; + ggml_tensor * g; + ggml_tensor * beta; + ggml_tensor * dt_bias; + ggml_tensor * a; + ggml_tensor * cache; // 2D [D, mem_size*K] + ggml_tensor * rows_in; // I32 [n_seqs] + ggml_tensor * rows_out; // I64 [min(T,K)*n_seqs] +}; + +static gdn_tensors make_tensors(ggml_context * ctx, const gdn_case & c, const gdn_inputs & in) { + const int64_t D = c.S_v * c.S_v * c.H_v; + const int64_t n_rows = c.mem_size * c.K; + + gdn_tensors t; + t.q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, c.S_v, c.H_k, c.T, c.n_seqs); + t.k = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, c.S_v, c.H_k, c.T, c.n_seqs); + t.v = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, c.S_v, c.H_v, c.T, c.n_seqs); + t.g = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 1, c.H_v, c.T, c.n_seqs); + t.beta = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 1, c.H_v, c.T, c.n_seqs); + t.dt_bias = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, c.H_v); + t.a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, c.H_v); + t.cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, D, n_rows); + t.rows_in = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, c.n_seqs); + t.rows_out = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, (int64_t) in.rows_out.size()); + + memcpy(t.q->data, in.q.data(), ggml_nbytes(t.q)); + memcpy(t.k->data, in.k.data(), ggml_nbytes(t.k)); + memcpy(t.v->data, in.v.data(), ggml_nbytes(t.v)); + memcpy(t.g->data, in.g.data(), ggml_nbytes(t.g)); + memcpy(t.beta->data, in.beta.data(), ggml_nbytes(t.beta)); + if (c.raw_gates) { + memcpy(t.dt_bias->data, in.dt_bias.data(), ggml_nbytes(t.dt_bias)); + memcpy(t.a->data, in.a.data(), ggml_nbytes(t.a)); + } else { + memset(t.dt_bias->data, 0, ggml_nbytes(t.dt_bias)); + memset(t.a->data, 0, ggml_nbytes(t.a)); + } + memcpy(t.cache->data, in.cache.data(), ggml_nbytes(t.cache)); + memcpy(t.rows_in->data, in.rows_in.data(), ggml_nbytes(t.rows_in)); + memcpy(t.rows_out->data, in.rows_out.data(), ggml_nbytes(t.rows_out)); + + return t; +} + +static ggml_context * new_ctx() { + ggml_init_params params = { + /*.mem_size =*/ 64u * 1024u * 1024u, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ false, + }; + ggml_context * ctx = ggml_init(params); + if (!ctx) { + fprintf(stderr, "ggml_init failed\n"); + exit(1); + } + return ctx; +} + +struct path_result { + std::vector attn; // attention scores, S_v*H_v*T*n_seqs floats + std::vector cache; // whole cache tensor after the graph ran +}; + +static void snapshot(path_result & r, const ggml_tensor * attn_src, size_t attn_bytes, const ggml_tensor * cache) { + r.attn.resize(attn_bytes); + memcpy(r.attn.data(), attn_src->data, attn_bytes); + r.cache.resize(ggml_nbytes(cache)); + memcpy(r.cache.data(), cache->data, ggml_nbytes(cache)); +} + +// A: gathered legacy path, exactly as build_recurrent_attn's !keep (K == 1) +// and keep/gathered (K > 1) branches emit it +static path_result run_gathered(const gdn_case & c, const gdn_inputs & in) { + ggml_context * ctx = new_ctx(); + gdn_tensors t = make_tensors(ctx, c, in); + + const int64_t D = c.S_v * c.S_v * c.H_v; + const int64_t attn_score_elems = c.S_v * c.H_v * c.T * c.n_seqs; + const int64_t state_size_per_snap = D * c.n_seqs; + const int64_t n_written = std::min(c.T, c.K); + const size_t row_size = D * sizeof(float); + + ggml_cgraph * gf = ggml_new_graph(ctx); + + // build_rs main gather (identity or rollback plane) -> 4D state + ggml_tensor * gathered = ggml_get_rows(ctx, t.cache, t.rows_in); + ggml_tensor * s4d = ggml_reshape_4d(ctx, gathered, c.S_v, c.S_v, c.H_v, c.n_seqs); + + ggml_tensor * gdn_out = ggml_gated_delta_net(ctx, t.q, t.k, t.v, t.g, t.beta, s4d, c.K); + if (c.raw_gates) { + ggml_gated_delta_net_set_raw_gates(gdn_out, t.dt_bias, t.a); + } + + if (c.K == 1) { + // build_delta_net_fused new_state view + ggml_cpy into rows [head, head + n_seqs) + ggml_tensor * new_state = ggml_view_4d(ctx, gdn_out, + c.S_v, c.S_v, c.H_v, c.n_seqs, + ggml_row_size(gdn_out->type, c.S_v), + ggml_row_size(gdn_out->type, c.S_v * c.S_v), + ggml_row_size(gdn_out->type, c.S_v * c.S_v * c.H_v), + ggml_row_size(gdn_out->type, c.S_v * c.H_v * c.T * c.n_seqs)); + ggml_build_forward_expand(gf, + ggml_cpy(ctx, new_state, + ggml_view_2d(ctx, t.cache, D, c.n_seqs, t.cache->nb[1], c.head * row_size))); + } else { + // strided 3D cpy: snapshot slot s -> plane s of the cache at head + ggml_tensor * src = ggml_view_3d(ctx, gdn_out, + D, c.n_seqs, n_written, + ggml_row_size(gdn_out->type, D), + ggml_row_size(gdn_out->type, state_size_per_snap), + ggml_row_size(gdn_out->type, attn_score_elems)); + ggml_tensor * dst = ggml_view_3d(ctx, t.cache, + D, c.n_seqs, n_written, + t.cache->nb[1], + (size_t) c.mem_size * row_size, + (size_t) c.head * row_size); + ggml_build_forward_expand(gf, ggml_cpy(ctx, src, dst)); + } + + if (ggml_graph_compute_with_ctx(ctx, gf, c.n_threads) != GGML_STATUS_SUCCESS) { + fprintf(stderr, "gathered graph compute failed\n"); + exit(1); + } + + path_result r; + snapshot(r, gdn_out, attn_score_elems * sizeof(float), t.cache); + ggml_free(ctx); + return r; +} + +// B: rows mode + SET_ROWS, exactly as the keep/rows branch emits it +static path_result run_rows(const gdn_case & c, const gdn_inputs & in) { + ggml_context * ctx = new_ctx(); + gdn_tensors t = make_tensors(ctx, c, in); + + const int64_t D = c.S_v * c.S_v * c.H_v; + const int64_t attn_score_elems = c.S_v * c.H_v * c.T * c.n_seqs; + const int64_t n_written = std::min(c.T, c.K); + + ggml_cgraph * gf = ggml_new_graph(ctx); + + ggml_tensor * gdn_out = ggml_gated_delta_net_rows(ctx, t.q, t.k, t.v, t.g, t.beta, t.cache, t.rows_in, (int) c.K); + if (c.raw_gates) { + ggml_gated_delta_net_set_raw_gates(gdn_out, t.dt_bias, t.a); + } + + ggml_tensor * snaps = ggml_view_2d(ctx, gdn_out, + D, n_written * c.n_seqs, + ggml_row_size(gdn_out->type, D), + ggml_row_size(gdn_out->type, attn_score_elems)); + + ggml_build_forward_expand(gf, ggml_set_rows(ctx, t.cache, snaps, t.rows_out)); + + if (ggml_graph_compute_with_ctx(ctx, gf, c.n_threads) != GGML_STATUS_SUCCESS) { + fprintf(stderr, "rows graph compute failed\n"); + exit(1); + } + + path_result r; + snapshot(r, gdn_out, attn_score_elems * sizeof(float), t.cache); + ggml_free(ctx); + return r; +} + +// C: the cache-writing op +static path_result run_inplace(const gdn_case & c, const gdn_inputs & in) { + ggml_context * ctx = new_ctx(); + gdn_tensors t = make_tensors(ctx, c, in); + + const int64_t attn_score_elems = c.S_v * c.H_v * c.T * c.n_seqs; + + ggml_cgraph * gf = ggml_new_graph(ctx); + + ggml_tensor * gdn_out = ggml_gated_delta_net_cache(ctx, t.q, t.k, t.v, t.g, t.beta, t.cache, + t.rows_in, t.rows_out, (int) c.K, (int) c.mem_size); + if (c.raw_gates) { + ggml_gated_delta_net_set_raw_gates(gdn_out, t.dt_bias, t.a); + } + + if (ggml_nelements(gdn_out) != attn_score_elems) { + fprintf(stderr, "in-place result carries %lld elements, expected scores only (%lld)\n", + (long long) ggml_nelements(gdn_out), (long long) attn_score_elems); + exit(1); + } + + ggml_build_forward_expand(gf, gdn_out); + + if (ggml_graph_compute_with_ctx(ctx, gf, c.n_threads) != GGML_STATUS_SUCCESS) { + fprintf(stderr, "in-place graph compute failed\n"); + exit(1); + } + + path_result r; + snapshot(r, gdn_out, attn_score_elems * sizeof(float), t.cache); + ggml_free(ctx); + return r; +} + +static size_t count_diff(const std::vector & a, const std::vector & b) { + if (a.size() != b.size()) { + return std::max(a.size(), b.size()); + } + size_t n = 0; + for (size_t i = 0; i < a.size(); ++i) { + n += a[i] != b[i]; + } + return n; +} + +int main() { + const int64_t S_v = 32; + const int64_t H_v = 4; + + int n_cases = 0; + int n_failed = 0; + uint32_t seed = 1234; + + for (int64_t K : { (int64_t) 1, (int64_t) 3 }) { + for (int64_t T : { (int64_t) 1, (int64_t) 3 }) { + for (int64_t n_seqs : { (int64_t) 1, (int64_t) 4, (int64_t) 8 }) { + for (int64_t r_in = 0; r_in < K; ++r_in) { + for (int raw = 0; raw < 2; ++raw) { + for (int64_t H_k : { H_v, H_v / 2 }) { + for (int n_threads : { 1, 4 }) { + gdn_case c; + c.S_v = S_v; + c.H_v = H_v; + c.H_k = H_k; + c.T = T; + c.n_seqs = n_seqs; + c.K = K; + c.mem_size = n_seqs + 3; // cells before and after the ubatch stay untouched + c.head = 2; + c.r_in = r_in; + c.raw_gates = raw != 0; + c.n_threads = n_threads; + + const gdn_inputs in = make_inputs(c, seed++); + + const path_result A = run_gathered(c, in); + const path_result B = run_rows(c, in); + const path_result C = run_inplace(c, in); + + const size_t d_attn_A = count_diff(C.attn, A.attn); + const size_t d_attn_B = count_diff(C.attn, B.attn); + const size_t d_cache_A = count_diff(C.cache, A.cache); + const size_t d_cache_B = count_diff(C.cache, B.cache); + + // the check is not vacuous: the in-place op must have changed the live rows + std::vector cache0((const uint8_t *) in.cache.data(), + (const uint8_t *) in.cache.data() + in.cache.size() * sizeof(float)); + const bool touched = count_diff(C.cache, cache0) != 0; + + const bool ok = d_attn_A == 0 && d_attn_B == 0 && d_cache_A == 0 && d_cache_B == 0 && touched; + + printf("K=%lld T=%lld n_seqs=%lld r_in=%lld raw=%d H_k=%lld nth=%d : attn diff A/B = %zu/%zu, " + "cache diff A/B = %zu/%zu, cache touched = %d -> %s\n", + (long long) K, (long long) T, (long long) n_seqs, (long long) r_in, raw, (long long) H_k, n_threads, + d_attn_A, d_attn_B, d_cache_A, d_cache_B, (int) touched, ok ? "OK" : "FAIL"); + + n_cases += 1; + n_failed += !ok; + }}}}}}} + + printf("%d cases, %d failed\n", n_cases, n_failed); + + return n_failed == 0 ? 0 : 1; +} From c875ea463d607e4c2d2f5743cca2a8813610520f Mon Sep 17 00:00:00 2001 From: abdulaziz Date: Sat, 19 Sep 2026 02:01:20 +0300 Subject: [PATCH 04/15] memory-recurrent: lazily zeroed anonymous-mmap RS buffer on the CPU host The recurrent cache is mem_size*(1+n_rs_seq) rows per layer and was memset at construction (ggml_backend_buffer_clear), so every cell of every snapshot plane was resident from the start: --parallel 16 with n_rs_seq = 2 pinned ~4.7 GiB of the 15 GiB here while one active slot needs ~0.3 GiB, evicting the mmap'd weights that decode streams on every token. For the CPU buffer type on Linux the buffer is now one anonymous private mmap wrapped with ggml_backend_cpu_buffer_from_ptr (tensors placed with a ggml_tallocr exactly as ggml_backend_alloc_ctx_tensors_from_buft would), with no eager memset: the kernel hands out zero pages on first touch, so residency follows the cells actually written. clear(true) drops the pages with madvise(MADV_DONTNEED) instead of a memset; the destructor unmaps after freeing the buffers (from_ptr buffers do not own their memory). Non-host bufts, non-Linux hosts and LLAMA_RS_EAGER_ZERO=1 keep today's alloc + clear. Exactness argument (in the ctor comment): every row the graph reads is either the rs_z row zeroed in-graph before the read, a row written by an earlier step (src0 = src of a used cell), a snapshot plane written by the same sequence before a rollback reads it (rollback <= n_rs_seq, tail kept in one ubatch, slots 0..min(T,K)-1 written), or a row written by state_read_data; padding is never read. The startup line keeps its format ('RS buffer size = N MiB') and appends '(anonymous mmap, lazily zeroed)' when the lazy path is active. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01QrdPnP5KL449GsoWBZx52R --- src/llama-memory-recurrent.cpp | 135 +++++++++++++++++++++++++++++++-- src/llama-memory-recurrent.h | 14 +++- 2 files changed, 143 insertions(+), 6 deletions(-) diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index 7ac65695b8ff..dc70452295e8 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -1,5 +1,6 @@ #include "llama-memory-recurrent.h" +#include "ggml-alloc.h" #include "ggml-backend.h" #include "llama-impl.h" #include "llama-io.h" @@ -8,15 +9,66 @@ #include #include +#include +#include #include #include #include #include +#ifdef __linux__ +#include +#endif + // // llama_memory_recurrent // +#ifdef __linux__ +// One anonymous, private mapping for every tensor of ctx, placed exactly as +// ggml_backend_alloc_ctx_tensors_from_buft would place them (same padded sizes, +// same alignment, same order) but wrapped with ggml_backend_cpu_buffer_from_ptr, +// which does not own the memory: the caller munmaps it after the buffer is +// freed. Returns nullptr (and touches nothing) when the mapping cannot be made. +static ggml_backend_buffer_t rs_alloc_lazy_zero_host_buffer( + ggml_context * ctx, ggml_backend_buffer_type_t buft, void ** ptr_out, size_t * size_out) { + const size_t size = ggml_backend_alloc_ctx_tensors_from_buft_size(ctx, buft); + if (size == 0) { + return nullptr; + } + + void * ptr = mmap(nullptr, size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); + if (ptr == MAP_FAILED) { + LLAMA_LOG_WARN("%s: anonymous mmap of %zu bytes failed (%s), falling back to the eager clear\n", + __func__, size, strerror(errno)); + return nullptr; + } + + // mmap returns page-aligned memory, which satisfies the CPU buffer alignment + ggml_backend_buffer_t buf = ggml_backend_cpu_buffer_from_ptr(ptr, size); + if (!buf) { + munmap(ptr, size); + return nullptr; + } + + ggml_tallocr talloc = ggml_tallocr_new(buf); + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { + if (t->data == nullptr && t->view_src == nullptr) { + if (ggml_tallocr_alloc(&talloc, t) != GGML_STATUS_SUCCESS) { + ggml_backend_buffer_free(buf); + munmap(ptr, size); + return nullptr; + } + } + } + + *ptr_out = ptr; + *size_out = size; + + return buf; +} +#endif + llama_memory_recurrent::llama_memory_recurrent( const llama_model & model, ggml_type type_r, @@ -106,13 +158,62 @@ llama_memory_recurrent::llama_memory_recurrent( } // allocate tensors and initialize the buffers to avoid NaNs in the padding + // + // Host (CPU) buffers, Linux: by default the whole RS buffer is one anonymous + // private mmap (rs_alloc_lazy_zero_host_buffer) and is NOT memset. The + // kernel hands out zero pages on first touch, so the "no NaN padding" + // guarantee of the eager clear holds without making every cell of every + // snapshot plane resident at construction: the buffer is + // mem_size*(1+n_rs_seq) rows per layer, so --parallel 16 with n_rs_seq = 2 + // pinned ~4.7 GiB of the 15 GiB here (evicting the mmap'd weights that + // decode streams on every token) while one active slot needs ~0.3 GiB. + // Residency now follows the cells that were actually written. + // + // Why skipping the memset is exact -- every row the graph ever reads is + // zero by construction or was written before it is read: + // - a fresh sequence reads the rs_z row, which the graph zeroes in place + // (ggml_scale_inplace by 0 in build_rs / build_rs_cache_view) before any + // consumer reads it; 0 * (a kernel zero page) == 0 either way; + // - every other main row is src0 = src of a used cell (find_slot), i.e. a + // row written by an earlier step's write-back / in-place update, by the + // extra-row relocation, or by state_read_data; + // - a rollback read of snapshot plane r (1 <= r <= n_rs_seq) only happens + // after a ubatch of that sequence wrote slot r of the same cell: seq_rm + // only accepts rollback <= n_rs_seq, split_equal keeps the trailing + // n_rs_seq + 1 tokens of a sequence in one ubatch and the op writes slots + // 0..min(T, K)-1, so the plane was written before it is read; + // - the alignment padding between tensors is never read. + // clear(true) drops the pages with madvise(MADV_DONTNEED) instead of a + // memset (a private anonymous mapping then reads back as zero). + // LLAMA_RS_EAGER_ZERO=1 restores the eager alloc + memset path. +#ifdef __linux__ + static const bool rs_eager_zero = getenv("LLAMA_RS_EAGER_ZERO") != nullptr; +#endif + for (auto & [buft, ctx] : ctx_map) { - ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft); + ggml_backend_buffer_t buf = nullptr; + bool lazy = false; + +#ifdef __linux__ + if (!rs_eager_zero && buft == ggml_backend_cpu_buffer_type()) { + void * ptr = nullptr; + size_t size = 0; + buf = rs_alloc_lazy_zero_host_buffer(ctx.get(), buft, &ptr, &size); + if (buf) { + mmaps.push_back({ buf, ptr, size }); + lazy = true; + } + } +#endif if (!buf) { - throw std::runtime_error("failed to allocate buffer for rs cache"); + buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft); + if (!buf) { + throw std::runtime_error("failed to allocate buffer for rs cache"); + } + ggml_backend_buffer_clear(buf, 0); } - ggml_backend_buffer_clear(buf, 0); - LLAMA_LOG_INFO("%s: %10s RS buffer size = %8.2f MiB\n", __func__, ggml_backend_buffer_name(buf), ggml_backend_buffer_get_size(buf)/1024.0/1024.0); + LLAMA_LOG_INFO("%s: %10s RS buffer size = %8.2f MiB%s\n", __func__, ggml_backend_buffer_name(buf), ggml_backend_buffer_get_size(buf)/1024.0/1024.0, + lazy ? " (anonymous mmap, lazily zeroed)" : ""); ctxs_bufs.emplace_back(std::move(ctx), buf); } @@ -127,6 +228,17 @@ llama_memory_recurrent::llama_memory_recurrent( } } +llama_memory_recurrent::~llama_memory_recurrent() { + // the from_ptr buffers do not own their memory: free them first, then unmap + ctxs_bufs.clear(); +#ifdef __linux__ + for (const auto & m : mmaps) { + munmap(m.ptr, m.size); + } +#endif + mmaps.clear(); +} + void llama_memory_recurrent::clear(bool data) { for (int32_t i = 0; i < (int32_t) size; ++i) { cells[i].pos = -1; @@ -140,7 +252,20 @@ void llama_memory_recurrent::clear(bool data) { if (data) { for (auto & [_, buf] : ctxs_bufs) { - ggml_backend_buffer_clear(buf.get(), 0); + bool dropped = false; +#ifdef __linux__ + for (const auto & m : mmaps) { + if (m.buf == buf.get()) { + // drop the pages instead of memset: a private anonymous + // mapping reads back as zero and stays non-resident + dropped = madvise(m.ptr, m.size, MADV_DONTNEED) == 0; + break; + } + } +#endif + if (!dropped) { + ggml_backend_buffer_clear(buf.get(), 0); + } } } diff --git a/src/llama-memory-recurrent.h b/src/llama-memory-recurrent.h index 83717a67e3a9..78aa31dccc4c 100644 --- a/src/llama-memory-recurrent.h +++ b/src/llama-memory-recurrent.h @@ -26,7 +26,7 @@ class llama_memory_recurrent : public llama_memory_i { uint32_t n_rs_seq, const layer_filter_cb & filter); - ~llama_memory_recurrent() = default; + ~llama_memory_recurrent(); // // llama_memory_i @@ -121,6 +121,18 @@ class llama_memory_recurrent : public llama_memory_i { // ggml contexts for the KV cache along with the allocated backend buffers: std::vector> ctxs_bufs; + // host RS buffers backed by an anonymous private mmap (kernel zero pages, + // resident only where written) and wrapped with + // ggml_backend_cpu_buffer_from_ptr, which does not own the memory: unmapped + // in the destructor after the buffers are freed. Empty when + // LLAMA_RS_EAGER_ZERO=1 or on non-Linux hosts (eager alloc + memset). + struct rs_mmap_region { + ggml_backend_buffer_t buf; + void * ptr; + size_t size; + }; + std::vector mmaps; + size_t total_size() const; size_t size_r_bytes() const; From 89f8a273efb1c818e40779d285e7eec803b4883f Mon Sep 17 00:00:00 2001 From: abdulaziz Date: Sat, 19 Sep 2026 05:24:05 +0300 Subject: [PATCH 05/15] server: honor the per-request speculative.n_max (0 = no drafting for the request) The "speculative.n_max" request field was parsed inside an `#if 0` block and nothing in the server read task->params.speculative, so every generating slot drafted the global --spec-draft-n-max tokens whether or not the request asked for it. With draft-mtp the per-slot token counts then differed step by step (1 + d, d in 0..n_max), which split_equal(n_keep_tail = n_rs_seq + 1) cannot batch: N "opted-out" slots decoded as up to N ubatches per step instead of one. Move only the n_max field out of the `#if 0` (n_min/p_min/type/ngram_* stay disabled) and record the policy per slot at task launch: - spec_enabled = can_speculate() && n_max != 0 -> get_n_draft_max() returns 0 for an opted-out slot: dp.drafting stays false, no spec checkpoint update, handle_last_sampled_token adds exactly one token, post_decode samples it normally and the verify loop never sees the slot. N such slots with consecutive ids form one N x 1 ubatch (remaining == 0 satisfies the tail rule), the same call shape the no-spec server makes. - spec_n_max clamps the draft of a request that sends 0 < n_max; a request without the field carries the server-wide value, at which the MTP loop stops anyway, so its draft, batch, checkpoints and rollbacks are unchanged byte for byte. Mixed steps split as today: a drafting slot contributes 1 + d tokens and is emitted in its own equal-length ubatch (tail rule), the opted-out run forms a T = 1 ubatch. LLAMA_SERVER_SPEC_MULTISLOT=0 restores the previous policy (every slot drafts, the request field is ignored). Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01QrdPnP5KL449GsoWBZx52R --- tools/server/server-context.cpp | 41 +++++++++++++++++++++++++++++++++ tools/server/server-schema.cpp | 12 ++++++---- 2 files changed, 49 insertions(+), 4 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index a9edbd7be8b4..680359df807b 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -212,6 +212,11 @@ struct server_slot { common_prompt_checkpoint spec_ckpt; bool spec_is_replay = false; + // per-request drafting policy, set in launch_slot_with_task from "speculative.n_max" + // (LLAMA_SERVER_SPEC_MULTISLOT=0 restores the server-wide policy: every slot drafts, no cap) + bool spec_enabled = false; // false: the slot never drafts -> one token per step, single state plane + int spec_n_max = -1; // per-request draft cap; < 0: no per-request limit + // TODO: move members that belong to the task (such as `generated_text`, `has_new_line`) to task_results_state // see https://github.com/ggml-org/llama.cpp/pull/18283#issuecomment-3710175837 std::unique_ptr task; @@ -441,6 +446,11 @@ struct server_slot { return 0; } + // the request opted out of drafting ("speculative.n_max": 0) + if (!spec_enabled) { + return 0; + } + // determine the max draft that fits the current slot state // note: slot.prompt is not yet expanded with the `id` token sampled above // also, need to leave space for 1 extra token to allow context shifts @@ -450,6 +460,12 @@ struct server_slot { n_draft_max = std::min(n_draft_max, n_remaining() - 1); } + // per-request cap; a request without the field carries the server-wide n_max, at which the + // draft implementation stops anyway, so the clamp is a no-op for it + if (spec_n_max > 0) { + n_draft_max = std::min(n_draft_max, spec_n_max); + } + SLT_DBG(*this, "max possible draft: %d\n", n_draft_max); return n_draft_max; @@ -862,6 +878,11 @@ struct server_context_impl { int slots_debug = 0; // env: LLAMA_SERVER_SLOTS_DEBUG int slots_n_diff = 0; // env: LLAMA_SERVER_SLOTS_N_DIFF + // env: LLAMA_SERVER_SPEC_MULTISLOT (default 1). 0 restores the server-wide speculative policy: + // every slot drafts with the global n_max, the per-request "speculative.n_max" is ignored and + // every sequence keeps the full rollback snapshot budget + bool spec_multislot = true; + int n_empty_consecutive = 0; std::unique_ptr prompt_cache; @@ -1258,6 +1279,15 @@ struct server_context_impl { } } + { + const char * LLAMA_SERVER_SPEC_MULTISLOT = getenv("LLAMA_SERVER_SPEC_MULTISLOT"); + spec_multislot = LLAMA_SERVER_SPEC_MULTISLOT ? atoi(LLAMA_SERVER_SPEC_MULTISLOT) != 0 : true; + + if (!spec_multislot) { + SRV_WRN("%s\n", "LLAMA_SERVER_SPEC_MULTISLOT = 0: per-request \"speculative.n_max\" is ignored, every slot drafts"); + } + } + // the update_slots() logic will always submit a maximum of n_batch or n_parallel tokens // note that n_batch can be > n_ctx (e.g. for non-causal attention models such as BERT where the KV cache is not used) { @@ -1719,6 +1749,17 @@ struct server_context_impl { slot.task = std::make_unique(std::move(task)); + // per-request drafting policy: "speculative.n_max": 0 opts the request out of speculative + // decoding, so the slot adds exactly one token per step and N such slots decode as one + // N x 1 ubatch (the shape the no-spec server decodes). A request without the field keeps the + // server-wide n_max -> the unchanged drafting path. The schema rejects negative values. + { + const int n_max_req = slot.task->params.speculative.draft.n_max; + + slot.spec_enabled = slot.can_speculate() && (!spec_multislot || n_max_req != 0); + slot.spec_n_max = spec_multislot ? n_max_req : -1; + } + slot.state = slot.task->is_child() ? SLOT_STATE_WAIT_OTHER // wait for the parent to process prompt : SLOT_STATE_STARTED; diff --git a/tools/server/server-schema.cpp b/tools/server/server-schema.cpp index 64b9251295ce..dca2dbee0a67 100644 --- a/tools/server/server-schema.cpp +++ b/tools/server/server-schema.cpp @@ -194,13 +194,17 @@ std::vector> make_llama_cmpl_schema(const common_params & // Speculative decoding params // - // TODO: to keep things simple, we disable speculative parameter adjustments for now -#if 0 - // TODO: for now, be able to adjust only the draft-model based speculative parameters + // per-request draft cap (draft-model based methods). 0 opts the request out of drafting: the + // slot then adds exactly one token per step and batches with the other non-drafting slots + // (server_slot::spec_enabled). Without the field the request inherits the server-wide + // --spec-draft-n-max (params_base), i.e. the unchanged drafting path. add((new field_num("speculative.n_max", params.speculative.draft.n_max)) ->set_hard_limits(0, INT32_MAX) - ->set_desc("Maximum number of tokens to draft during speculative decoding")); + ->set_desc("Maximum number of tokens to draft during speculative decoding (0 = no drafting for this request)")); + // TODO: to keep things simple, we disable the remaining speculative parameter adjustments for now +#if 0 + // TODO: for now, be able to adjust only the draft-model based speculative parameters add((new field_num("speculative.n_min", params.speculative.draft.n_min)) ->set_hard_limits(0, INT32_MAX) ->set_desc("Minimum number of draft tokens to use for speculative decoding"); From 44e4b09d9263e839a946009464cc312e2067449b Mon Sep 17 00:00:00 2001 From: abdulaziz Date: Sat, 19 Sep 2026 05:27:49 +0300 Subject: [PATCH 06/15] llama: per-sequence rollback snapshot budget (llama_memory_seq_rs_snapshots) The recurrent memory widens its R/S tensors to (1 + n_rs_seq) planes and every ubatch writes min(T, n_rs_seq + 1) of them for every sequence, whether or not that sequence will ever roll back. Add a per-seq budget rs_n_snap in [0, n_rs_seq], default n_rs_seq (every existing caller keeps today's behaviour), settable through the new public API and forwarded by the hybrid memory: - llama_memory_recurrent::seq_rm accepts a partial tail removal only up to the seq's budget; a budget-0 seq refuses exactly like a memory with n_rs_seq == 0, so a caller sees the FULL-type behaviour for that sequence. - llama_memory_recurrent_context::get_n_snap() = max budget over the seqs of the ubatch apply() just placed (n_rs_seq for the reserve context so the compute buffers fit the largest graph, 0 without snapshots). Side-effect free, so it is valid at graph build and again in can_reuse. Nothing resets the budget but the API (seq_rm / clear / state_read leave it), and no graph consumes it yet. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01QrdPnP5KL449GsoWBZx52R --- include/llama.h | 10 +++++++ src/llama-context.cpp | 11 +++++++ src/llama-memory-hybrid.cpp | 4 +++ src/llama-memory-hybrid.h | 1 + src/llama-memory-recurrent.cpp | 54 ++++++++++++++++++++++++++++++++-- src/llama-memory-recurrent.h | 14 +++++++++ src/llama-memory.h | 5 ++++ 7 files changed, 97 insertions(+), 2 deletions(-) diff --git a/include/llama.h b/include/llama.h index deb1e12828c6..781a4251f2e1 100644 --- a/include/llama.h +++ b/include/llama.h @@ -757,6 +757,16 @@ extern "C" { llama_pos p0, llama_pos p1); + // Set the number of recurrent-state rollback snapshot planes (0..llama_n_rs_seq) that the memory + // may write for and roll back on the specified sequence. Default: llama_n_rs_seq (every plane). + // A sequence with budget 0 keeps a single state plane; partial removals of its tail (see + // llama_memory_seq_rm) are refused, exactly as with a context created with n_rs_seq = 0. + // Returns false if the memory has no recurrent state or the arguments are out of range + LLAMA_API bool llama_memory_seq_rs_snapshots( + llama_memory_t mem, + llama_seq_id seq_id, + uint32_t n_snap); + // Copy all tokens that belong to the specified sequence to another sequence // p0 < 0 : [0, p1] // p1 < 0 : [p0, inf) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index cc8c004eb680..a06fef35c15a 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -4232,6 +4232,17 @@ bool llama_memory_seq_rm( return mem->seq_rm(seq_id, p0, p1); } +bool llama_memory_seq_rs_snapshots( + llama_memory_t mem, + llama_seq_id seq_id, + uint32_t n_snap) { + if (!mem) { + return false; + } + + return mem->seq_rs_snapshots(seq_id, n_snap); +} + void llama_memory_seq_cp( llama_memory_t mem, llama_seq_id seq_id_src, diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index 42c7381a9e6f..099422df09fd 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -149,6 +149,10 @@ bool llama_memory_hybrid::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1 return mem_attn->seq_rm(seq_id, p0, p1); } +bool llama_memory_hybrid::seq_rs_snapshots(llama_seq_id seq_id, uint32_t n_snap) { + return mem_recr->seq_rs_snapshots(seq_id, n_snap); +} + void llama_memory_hybrid::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) { mem_attn->seq_cp(seq_id_src, seq_id_dst, p0, p1); mem_recr->seq_cp(seq_id_src, seq_id_dst, p0, p1); diff --git a/src/llama-memory-hybrid.h b/src/llama-memory-hybrid.h index 484eafb74991..78fba3ddb1d6 100644 --- a/src/llama-memory-hybrid.h +++ b/src/llama-memory-hybrid.h @@ -61,6 +61,7 @@ class llama_memory_hybrid : public llama_memory_i { void clear(bool data) override; bool seq_rm (llama_seq_id seq_id, llama_pos p0, llama_pos p1) override; + bool seq_rs_snapshots(llama_seq_id seq_id, uint32_t n_snap) override; void seq_cp (llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) override; void seq_keep(llama_seq_id seq_id) override; void seq_add (llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) override; diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index dc70452295e8..d2e01b480749 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -86,6 +86,7 @@ llama_memory_recurrent::llama_memory_recurrent( this->n_rs_seq = n_rs_seq; rs_idx.assign(n_seq_max, 0); + rs_n_snap.assign(n_seq_max, n_rs_seq); // every seq keeps all planes unless the API lowers its budget cells.clear(); cells.resize(mem_size); @@ -304,12 +305,13 @@ bool llama_memory_recurrent::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos if (tail_id >= 0) { auto & cell = cells[tail_id]; - // partial rollback via per-token snapshot index (bounded by n_rs_seq) + // partial rollback via per-token snapshot index, bounded by the seq's snapshot budget + // (rs_n_snap[seq_id] <= n_rs_seq; a budget of 0 refuses like a memory with n_rs_seq == 0) if (0 < p0 && p0 <= cell.pos && p1 > cell.pos) { const llama_pos rollback = cell.pos - (p0 - 1); // pending rollback is single-use const bool pending = rs_idx[seq_id] != 0; - if (!pending && rollback >= 1 && rollback <= (llama_pos) n_rs_seq) { + if (!pending && rollback >= 1 && rollback <= (llama_pos) rs_n_snap[seq_id]) { set_rs_idx(seq_id, (uint32_t) rollback); cell.pos = p0 - 1; return true; @@ -531,6 +533,24 @@ void llama_memory_recurrent::set_rs_idx(llama_seq_id seq_id, uint32_t idx) { rs_idx[seq_id] = idx; } +bool llama_memory_recurrent::seq_rs_snapshots(llama_seq_id seq_id, uint32_t n_snap) { + if (seq_id < 0 || (uint32_t) seq_id >= n_seq_max) { + LLAMA_LOG_ERROR("%s: invalid seq_id (%d) - larger than n_seq_max (%d)\n", __func__, seq_id, n_seq_max); + return false; + } + + if (n_snap > n_rs_seq) { + LLAMA_LOG_ERROR("%s: snapshot budget (%u) larger than n_rs_seq (%u)\n", __func__, n_snap, n_rs_seq); + return false; + } + + assert(n_seq_max == rs_n_snap.size()); + + rs_n_snap[seq_id] = n_snap; + + return true; +} + std::map llama_memory_recurrent::memory_breakdown() const { std::map ret; for (const auto & [_, buf] : ctxs_bufs) { @@ -1422,3 +1442,33 @@ bool llama_memory_recurrent_context::rs_inplace_ok(uint32_t n_seqs) const { return true; } + +uint32_t llama_memory_recurrent_context::get_n_snap() const { + // no snapshot planes at all (n_rs_seq == 0): K = 1 regardless of the budgets + if (mem->n_rs_seq == 0) { + return 0; + } + + // the reserve context sizes the compute buffers: keep the largest graph + if (is_full || ubatches.empty()) { + return mem->n_rs_seq; + } + + // ubatches[i_next] is the ubatch apply() -> find_slot just placed; for an equal_seqs ubatch the + // seq ids of sequence set s are those of its first token + const auto & ubatch = ubatches[i_next]; + + uint32_t res = 0; + + for (uint32_t s = 0; s < ubatch.n_seqs; ++s) { + const uint32_t i = s*ubatch.n_seq_tokens; + for (int32_t j = 0; j < ubatch.n_seq_id[i]; ++j) { + const llama_seq_id seq_id = ubatch.seq_id[i][j]; + if (seq_id >= 0 && (size_t) seq_id < mem->rs_n_snap.size()) { + res = std::max(res, mem->rs_n_snap[seq_id]); + } + } + } + + return res; +} diff --git a/src/llama-memory-recurrent.h b/src/llama-memory-recurrent.h index 78aa31dccc4c..91473a6474c5 100644 --- a/src/llama-memory-recurrent.h +++ b/src/llama-memory-recurrent.h @@ -44,6 +44,7 @@ class llama_memory_recurrent : public llama_memory_i { void clear(bool data) override; bool seq_rm (llama_seq_id seq_id, llama_pos p0, llama_pos p1) override; + bool seq_rs_snapshots(llama_seq_id seq_id, uint32_t n_snap) override; void seq_cp (llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) override; void seq_keep(llama_seq_id seq_id) override; void seq_add (llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) override; @@ -78,6 +79,12 @@ class llama_memory_recurrent : public llama_memory_i { void set_rs_idx(llama_seq_id seq_id, uint32_t idx); + // per-seq rollback snapshot budget in [0, n_rs_seq] (llama_memory_seq_rs_snapshots): the number + // of snapshot planes a ubatch may write for the seq -- the graph writes K = 1 + the max over the + // ubatch's seqs (llama_memory_recurrent_context::get_n_snap) -- and the largest partial rollback + // seq_rm accepts for it. Defaults to n_rs_seq; only the API changes it (not seq_rm/clear/state_read) + std::vector rs_n_snap; + // computed before each graph build uint32_t n = 0; @@ -193,6 +200,13 @@ class llama_memory_recurrent_context : public llama_memory_context_i { // evaluated at graph build and again in can_reuse. bool rs_inplace_ok(uint32_t n_seqs) const; + // rollback snapshot planes of the current ubatch: the GDN paths write K = 1 + get_n_snap() + // state planes per seq. It is the max of the per-seq budgets (mem->rs_n_snap) over the + // ubatch's seqs, n_rs_seq for the full (reserve) context so the compute buffers fit the + // largest graph, and 0 when the memory keeps no snapshots. Side-effect free (same rule as + // rs_inplace_ok): evaluated at graph build and again in can_reuse (part of the reuse key). + uint32_t get_n_snap() const; + private: const llama_memory_status status; diff --git a/src/llama-memory.h b/src/llama-memory.h index db825396645e..6407dc4abceb 100644 --- a/src/llama-memory.h +++ b/src/llama-memory.h @@ -108,6 +108,11 @@ struct llama_memory_i { virtual void clear(bool data) = 0; virtual bool seq_rm (llama_seq_id seq_id, llama_pos p0, llama_pos p1) = 0; + + // number of rollback snapshot planes the memory may write for / roll back on this sequence + // (see llama_memory_seq_rs_snapshots); memories without recurrent state return false + virtual bool seq_rs_snapshots(llama_seq_id /*seq_id*/, uint32_t /*n_snap*/) { return false; } + virtual void seq_cp (llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) = 0; virtual void seq_keep(llama_seq_id seq_id) = 0; virtual void seq_add (llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) = 0; From 724056660af60fe882ef9e4f2e9d2205504d79e1 Mon Sep 17 00:00:00 2001 From: abdulaziz Date: Sat, 19 Sep 2026 05:27:49 +0300 Subject: [PATCH 07/15] gdn: derive the snapshot plane count K of a ubatch from the seqs' budgets The delta-net state write paths took K = cparams.n_rs_seq + 1 for every ubatch. Take K = 1 + llama_memory_recurrent_context::get_n_snap() instead, recorded once per graph in llm_graph_input_rs::rs_n_snap (build_rs_inp_impl, so both the recurrent and the hybrid callers get it) and made part of the graph reuse key: at T = 1 the s_write_rows length is identical for K = 1 and K = 3 while the op params (ggml_gated_delta_net_cache), the number of conv-state copies and the fused-gate decision differ. A ubatch that contains a drafting seq (budget n_rs_seq) gets exactly today's K on every path -- in-place, gathered, rows -- identical topology, identical s_write_rows contents, identical op params: bit-identical scores and cache rows. A ubatch of opted-out seqs (budget 0) gets K = 1: one conv-state copy, a single S plane on the in-place path (the kernel writes no snapshot rows for K == 1), the !keep single-row copy on the gathered path and n_written = 1 on the rows path -- also on the prompt ubatch and on the gathered transition ubatches, which is what bounds the per-slot residency to one plane. qwen35 keys fused_ok / the K1-only env / the Metal rows mode on K as well, so an opted-out ubatch builds node for node the graph the no-spec context builds (n_rs_seq == 0): the no-spec server becomes a usable md5 oracle for those requests. qwen35moe / qwen3next share build_recurrent_attn and keep today's K unless the new API is called. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01QrdPnP5KL449GsoWBZx52R --- src/llama-graph.cpp | 5 +++++ src/llama-graph.h | 6 ++++++ src/models/delta-net-base.cpp | 24 ++++++++++++++++-------- src/models/qwen35.cpp | 13 ++++++++++--- 4 files changed, 37 insertions(+), 11 deletions(-) diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 002c599eddb2..569d13c5bfd5 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -422,6 +422,10 @@ bool llm_graph_input_rs::can_reuse_rs(const llama_memory_recurrent_context * mct // was already updated by apply() -> find_slot before can_reuse runs res &= rs_inplace == mctx_cur->rs_inplace_ok(params.ubatch.n_seqs); + // the snapshot plane count K = rs_n_snap + 1 is baked into the op params + // (and the s_write_rows check above cannot see it at T = 1) + res &= rs_n_snap == mctx_cur->get_n_snap(); + return res; } @@ -3579,6 +3583,7 @@ static std::unique_ptr build_rs_inp_impl( inp->rs_z = mctx_cur->get_rs_z(); inp->rs_inplace = mctx_cur->rs_inplace_ok(n_seqs); + inp->rs_n_snap = mctx_cur->get_n_snap(); return inp; } diff --git a/src/llama-graph.h b/src/llama-graph.h index c3b7e13af6ba..d87d3c788ec3 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -350,6 +350,12 @@ class llm_graph_input_rs : public llm_graph_input_i { // update the cache rows in place (rs_inplace_ok); part of the reuse key bool rs_inplace = false; + // rollback snapshot planes of this graph's ubatch: the GDN paths write + // K = rs_n_snap + 1 state planes per seq (get_n_snap). Part of the reuse + // key: at T = 1 the s_write_rows length is the same for every K while the + // op params, the conv-state copies and the fused-gate choice differ + uint32_t rs_n_snap = 0; + const llama_memory_recurrent_context * mctx; // used in view offsets, need to match for valid graph reuse diff --git a/src/models/delta-net-base.cpp b/src/models/delta-net-base.cpp index 26799679f6af..4aef1caa24f7 100644 --- a/src/models/delta-net-base.cpp +++ b/src/models/delta-net-base.cpp @@ -481,7 +481,12 @@ ggml_tensor * llm_build_delta_net_base::build_conv_state( const size_t row_size = ggml_row_size(conv_states_all->type, row_count); - if (cparams.n_rs_seq == 0) { + // snapshot planes of this ubatch: K = 1 + the rollback budget of its seqs + // (n_rs_seq + 1 when a drafting seq is present, 1 for a ubatch of opted-out + // seqs and for a memory without snapshots -- the graph the no-spec context builds) + const int64_t K = 1 + (int64_t) inp->rs_n_snap; + + if (K == 1) { const int64_t s_idx = conv_input->ne[0] - conv_states->ne[0]; const int64_t s_slot = 0; @@ -503,8 +508,7 @@ ggml_tensor * llm_build_delta_net_base::build_conv_state( // [TAG_RECURRENT_ROLLBACK_SPLITS] // this logic assumes that the last (n_rs_seq + 1) tokens of a sequence in a batch are inside // the same ubatch, which `split_equal()` guarantees via its n_keep_tail argument - - const int64_t K = (int64_t) cparams.n_rs_seq + 1; + // (K <= n_rs_seq + 1, so the guarantee covers every budget) for (int64_t t = 1; t <= K; ++t) { const int64_t s_idx = std::max(0, conv_input->ne[0] - conv_states->ne[0] - K + t); @@ -552,7 +556,13 @@ ggml_tensor * llm_build_delta_net_base::build_recurrent_attn( const int64_t n_seqs = v->ne[3]; const int64_t n_seq_tokens = q->ne[2]; - const bool keep = cparams.n_rs_seq > 0; + // snapshot planes of this ubatch (see build_conv_state): every state write + // path below derives its plane count from K, so a ubatch of opted-out seqs + // (K == 1) touches plane 0 only -- also on the prompt ubatch and on the + // gathered transition path + const int64_t K = 1 + (int64_t) inp->rs_n_snap; + + const bool keep = K > 1; GGML_ASSERT(state_rows == nullptr || keep || state_inplace); // Metal rows mode is a ring-path optimization GGML_ASSERT(!state_inplace || state_rows != nullptr); @@ -567,9 +577,8 @@ ggml_tensor * llm_build_delta_net_base::build_recurrent_attn( // carries the attention scores only: no gather, no dst state planes, no // write-back copy. Same float instruction sequence on the same inputs as // the paths below, so the outputs and the resulting cache rows are - // bit-identical (tests/test-gdn-inplace.cpp). - const int64_t K = cparams.n_rs_seq + 1; - + // bit-identical (tests/test-gdn-inplace.cpp). With K == 1 the kernel + // writes no snapshot rows at all. const bool raw = gdn_raw_beta && gdn_raw_alpha && gdn_raw_dt_bias && gdn_raw_a; ggml_tensor * gg = raw ? gdn_raw_alpha : g; ggml_tensor * bb = raw ? gdn_raw_beta : b; @@ -614,7 +623,6 @@ ggml_tensor * llm_build_delta_net_base::build_recurrent_attn( } const int64_t D = S_v * S_v * H_v; - const int64_t K = cparams.n_rs_seq + 1; const bool raw = gdn_raw_beta && gdn_raw_alpha && gdn_raw_dt_bias && gdn_raw_a; ggml_tensor * gg = raw ? gdn_raw_alpha : g; diff --git a/src/models/qwen35.cpp b/src/models/qwen35.cpp index 073a6e95b63e..38552a40f1cf 100644 --- a/src/models/qwen35.cpp +++ b/src/models/qwen35.cpp @@ -490,10 +490,17 @@ ggml_tensor * llama_model_qwen35::graph::build_layer_attn_linear( // in supports_op, which would silently move the whole recurrent op to CPU. static const bool gdn_state_rows_env = getenv("GGML_GDN_STATE_GATHER") == nullptr; + // snapshot planes of this ubatch: K = n_rs_seq + 1 when a drafting seq is in + // it, 1 for a ubatch of opted-out seqs (llama_memory_seq_rs_snapshots) and + // for a context without snapshots. Keyed on K rather than cparams.n_rs_seq + // so an opted-out ubatch builds node for node the graph the no-spec context + // builds (single conv-state copy, plane 0 only, fused gates by the probes). + const int64_t K = 1 + (int64_t) inp->rs_n_snap; + // the in-place op always runs the fused kernel: only take it where the legacy // path runs the fused kernel too (the ring path always does; the K == 1 path // only with the fused_gdn_* probes on), so the A/B stays bit-identical - const bool fused_ok = cparams.n_rs_seq > 0 || (n_seq_tokens == 1 ? cparams.fused_gdn_ar : cparams.fused_gdn_ch); + const bool fused_ok = K > 1 || (n_seq_tokens == 1 ? cparams.fused_gdn_ar : cparams.fused_gdn_ch); // GGML_GDN_INPLACE_K1_ONLY=1 restricts the in-place path to K == 1 (no // speculation); by default the MTP speculative mode (K = n_rs_seq + 1) takes @@ -503,9 +510,9 @@ ggml_tensor * llama_model_qwen35::graph::build_layer_attn_linear( static const bool gdn_inplace_k1_only_env = getenv("GGML_GDN_INPLACE_K1_ONLY") != nullptr; const bool gdn_state_inplace = gdn_state_rows_env && gdn_state_inplace_dev_ok && inp->rs_inplace && fused_ok && - (cparams.n_rs_seq == 0 || !gdn_inplace_k1_only_env); + (K == 1 || !gdn_inplace_k1_only_env); const bool gdn_state_rows_metal = gdn_state_rows_env && gdn_state_rows_dev_ok && !gdn_state_inplace_dev_ok && - cparams.n_rs_seq > 0; + K > 1; const bool gdn_state_rows = gdn_state_inplace || gdn_state_rows_metal; ggml_tensor * state; From fd9a0af1ae185052fa8b3fa51c5f2858e109c851 Mon Sep 17 00:00:00 2001 From: abdulaziz Date: Sat, 19 Sep 2026 05:30:51 +0300 Subject: [PATCH 08/15] server: give an opted-out slot a rollback snapshot budget of 0 At task launch, right after the drafting policy is fixed, set the slot's sequence budget on the target memory: llama_n_rs_seq planes for a drafting slot (today's K on every ubatch it takes part in), 0 for a slot whose request sent "speculative.n_max": 0. Such a slot decodes with K = 1 -- one conv-state copy, plane 0 only on the prompt ubatch, the generation ubatches and the gathered transition ubatches -- node for node the graph the no-spec server builds, so its residency is one S plane plus the prompt-end checkpoint that both servers keep. The refusal that a budget of 0 adds to llama_memory_seq_rm is unreachable from the server for these slots: every partial tail removal it issues comes from the verify loop (non-empty spec_draft only); the prompt-reuse path resets or restores a checkpoint and then removes from cell.pos + 1, the cache-reuse and context-shift ranges end at or before cell.pos, and rm_all never enters the partial branch -- the same calls the FULL-type no-spec server makes for every slot. A slot that inherits a pending rollback index from an earlier drafting request reads a fully written plane (covered by test-gdn-inplace). The MTP draft context on qwen35 holds a plain KV cache and needs no call. LLAMA_SERVER_SPEC_MULTISLOT=0 keeps every budget at n_rs_seq. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01QrdPnP5KL449GsoWBZx52R --- tools/server/server-context.cpp | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 680359df807b..db0c0e85996a 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -1758,6 +1758,13 @@ struct server_context_impl { slot.spec_enabled = slot.can_speculate() && (!spec_multislot || n_max_req != 0); slot.spec_n_max = spec_multislot ? n_max_req : -1; + + // rollback snapshot budget of the slot's sequence (recurrent memories; a no-op for the + // others): every plane for a drafting slot (K = n_rs_seq + 1, the unchanged graph), none + // for an opted-out slot (K = 1: plane 0 only, the graph the no-spec server builds). Set + // before the slot's first ubatch. The MTP draft context holds a plain KV cache, so the + // call goes to ctx_tgt only. + llama_memory_seq_rs_snapshots(llama_get_memory(ctx_tgt), slot.id, slot.spec_enabled ? llama_n_rs_seq(ctx_tgt) : 0); } slot.state = slot.task->is_child() From e4be22f4670f54de6a03a7142ff45e61f80edac3 Mon Sep 17 00:00:00 2001 From: abdulaziz Date: Sat, 19 Sep 2026 05:30:51 +0300 Subject: [PATCH 09/15] tests: gdn in-place op with a K = 1 rollback read and a T = 1 cross-K check Two families on top of the existing 192 cases (whose inputs, order and output lines are unchanged), for the per-sequence snapshot budget: - K = 1 with rows_in in plane r in {1, 2} of an (r + 1)-plane cache: an opted-out request launched on a slot whose seq still carries a pending rollback index from a previous drafting request. The cache-writing op must read plane r and write plane 0 only, byte for byte like the gathered and rows paths (96 cases). - T = 1 cross-K on a 3-plane cache: the op with K = 1 and with K = 3 on identical inputs must leave byte-identical scores and cache -- the mixed T = 1 ubatch where a drafting seq with an empty draft shares the ubatch with opted-out seqs (K = 3) gives every seq exactly the K = 1 result (72 cases). The cache now allocates n_planes rows per cell (== K for the original matrix). 360 cases, 0 failed. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01QrdPnP5KL449GsoWBZx52R --- tests/test-gdn-inplace.cpp | 157 ++++++++++++++++++++++++++++++++----- 1 file changed, 136 insertions(+), 21 deletions(-) diff --git a/tests/test-gdn-inplace.cpp b/tests/test-gdn-inplace.cpp index c52eeb008af7..5e1cc5a81e75 100644 --- a/tests/test-gdn-inplace.cpp +++ b/tests/test-gdn-inplace.cpp @@ -20,6 +20,20 @@ // the same cell), raw gates on/off, H_k == H_v and H_k == H_v/2, 1 and 4 // threads. KDA (vector gate) off, as in the qwen35 graph. // +// Two extra families cover the per-sequence snapshot budget +// (llama_memory_seq_rs_snapshots), where a ubatch of opted-out seqs runs with +// K = 1 while the cache still holds n_rs_seq + 1 planes: +// +// - K = 1 with rows_in in plane r in {1, 2} of a (r + 1)-plane cache: an +// opted-out request launched on a slot that inherited a pending rollback +// index from a previous drafting request. C must read plane r and write +// plane 0 only, byte for byte like A and B. +// - T = 1 cross-K: on identical inputs and a 3-plane cache, C with K = 1 and +// C with K = 3 must leave byte-identical scores and cache -- a T = 1 +// ubatch writes plane 0 only, so an opted-out seq sharing a T = 1 ubatch +// with a drafting seq (K = 3, the mixed step with an empty draft) gets +// exactly the K = 1 result. +// // CPU only; runs on tiny synthetic shapes (S_v = 32, H_v = 4) and needs no // model. @@ -41,6 +55,7 @@ struct gdn_case { int64_t T; int64_t n_seqs; int64_t K; + int64_t n_planes; // planes allocated in the cache (>= K and > r_in; the memory holds n_rs_seq + 1 whatever K is) int64_t mem_size; int64_t head; int64_t r_in; // plane read for every seq: rows_in[i] = r_in*mem_size + head + i @@ -66,7 +81,7 @@ static gdn_inputs make_inputs(const gdn_case & c, uint32_t seed) { }; const int64_t D = c.S_v * c.S_v * c.H_v; - const int64_t n_rows = c.mem_size * c.K; + const int64_t n_rows = c.mem_size * c.n_planes; gdn_inputs in; fill(in.q, c.S_v * c.H_k * c.T * c.n_seqs, -1.0f, 1.0f); @@ -109,14 +124,14 @@ struct gdn_tensors { ggml_tensor * beta; ggml_tensor * dt_bias; ggml_tensor * a; - ggml_tensor * cache; // 2D [D, mem_size*K] + ggml_tensor * cache; // 2D [D, mem_size*n_planes] ggml_tensor * rows_in; // I32 [n_seqs] ggml_tensor * rows_out; // I64 [min(T,K)*n_seqs] }; static gdn_tensors make_tensors(ggml_context * ctx, const gdn_case & c, const gdn_inputs & in) { const int64_t D = c.S_v * c.S_v * c.H_v; - const int64_t n_rows = c.mem_size * c.K; + const int64_t n_rows = c.mem_size * c.n_planes; gdn_tensors t; t.q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, c.S_v, c.H_k, c.T, c.n_seqs); @@ -126,7 +141,7 @@ static gdn_tensors make_tensors(ggml_context * ctx, const gdn_case & c, const gd t.beta = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 1, c.H_v, c.T, c.n_seqs); t.dt_bias = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, c.H_v); t.a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, c.H_v); - t.cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, D, n_rows); + t.cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, D, n_rows); // 2D [D, mem_size*n_planes] t.rows_in = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, c.n_seqs); t.rows_out = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, (int64_t) in.rows_out.size()); @@ -314,6 +329,60 @@ static size_t count_diff(const std::vector & a, const std::vector cache0((const uint8_t *) in.cache.data(), + (const uint8_t *) in.cache.data() + in.cache.size() * sizeof(float)); + return count_diff(r.cache, cache0) != 0; +} + +// C vs A and B on one input set +static bool run_abc(const gdn_case & c, const gdn_inputs & in, const char * tag) { + const path_result A = run_gathered(c, in); + const path_result B = run_rows(c, in); + const path_result C = run_inplace(c, in); + + const size_t d_attn_A = count_diff(C.attn, A.attn); + const size_t d_attn_B = count_diff(C.attn, B.attn); + const size_t d_cache_A = count_diff(C.cache, A.cache); + const size_t d_cache_B = count_diff(C.cache, B.cache); + + const bool touched = cache_touched(in, C); + + const bool ok = d_attn_A == 0 && d_attn_B == 0 && d_cache_A == 0 && d_cache_B == 0 && touched; + + printf("%sK=%lld T=%lld n_seqs=%lld r_in=%lld raw=%d H_k=%lld nth=%d : attn diff A/B = %zu/%zu, " + "cache diff A/B = %zu/%zu, cache touched = %d -> %s\n", + tag, (long long) c.K, (long long) c.T, (long long) c.n_seqs, (long long) c.r_in, (int) c.raw_gates, + (long long) c.H_k, c.n_threads, + d_attn_A, d_attn_B, d_cache_A, d_cache_B, (int) touched, ok ? "OK" : "FAIL"); + + return ok; +} + +// C with K = 1 vs C with K = 3 on one input set (T = 1, 3-plane cache) +static bool run_cross_k(const gdn_case & c1, const gdn_inputs & in) { + gdn_case c3 = c1; + c3.K = 3; + + const path_result C1 = run_inplace(c1, in); + const path_result C3 = run_inplace(c3, in); + + const size_t d_attn = count_diff(C3.attn, C1.attn); + const size_t d_cache = count_diff(C3.cache, C1.cache); + + const bool touched = cache_touched(in, C1); + + const bool ok = d_attn == 0 && d_cache == 0 && touched; + + printf("[cross-K T=1] K=1 vs K=3 n_seqs=%lld r_in=%lld raw=%d H_k=%lld nth=%d : attn diff = %zu, " + "cache diff = %zu, cache touched = %d -> %s\n", + (long long) c1.n_seqs, (long long) c1.r_in, (int) c1.raw_gates, (long long) c1.H_k, c1.n_threads, + d_attn, d_cache, (int) touched, ok ? "OK" : "FAIL"); + + return ok; +} + int main() { const int64_t S_v = 32; const int64_t H_v = 4; @@ -336,6 +405,7 @@ int main() { c.T = T; c.n_seqs = n_seqs; c.K = K; + c.n_planes = K; c.mem_size = n_seqs + 3; // cells before and after the ubatch stay untouched c.head = 2; c.r_in = r_in; @@ -344,30 +414,75 @@ int main() { const gdn_inputs in = make_inputs(c, seed++); - const path_result A = run_gathered(c, in); - const path_result B = run_rows(c, in); - const path_result C = run_inplace(c, in); + const bool ok = run_abc(c, in, ""); - const size_t d_attn_A = count_diff(C.attn, A.attn); - const size_t d_attn_B = count_diff(C.attn, B.attn); - const size_t d_cache_A = count_diff(C.cache, A.cache); - const size_t d_cache_B = count_diff(C.cache, B.cache); + n_cases += 1; + n_failed += !ok; + }}}}}}} - // the check is not vacuous: the in-place op must have changed the live rows - std::vector cache0((const uint8_t *) in.cache.data(), - (const uint8_t *) in.cache.data() + in.cache.size() * sizeof(float)); - const bool touched = count_diff(C.cache, cache0) != 0; + // K = 1 reading a pending rollback plane r in {1, 2} of an (r + 1)-plane cache: + // an opted-out request (budget 0 -> K = 1) launched on a slot whose seq still + // carries rs_idx from a previous drafting request (the new prompt extends the + // cached one, so no reset). Plane r must be read and plane 0 alone written. + for (int64_t T : { (int64_t) 1, (int64_t) 3 }) { + for (int64_t n_seqs : { (int64_t) 1, (int64_t) 4, (int64_t) 8 }) { + for (int64_t r_in : { (int64_t) 1, (int64_t) 2 }) { + for (int raw = 0; raw < 2; ++raw) { + for (int64_t H_k : { H_v, H_v / 2 }) { + for (int n_threads : { 1, 4 }) { + gdn_case c; + c.S_v = S_v; + c.H_v = H_v; + c.H_k = H_k; + c.T = T; + c.n_seqs = n_seqs; + c.K = 1; + c.n_planes = std::max(c.K, r_in + 1); + c.mem_size = n_seqs + 3; + c.head = 2; + c.r_in = r_in; + c.raw_gates = raw != 0; + c.n_threads = n_threads; - const bool ok = d_attn_A == 0 && d_attn_B == 0 && d_cache_A == 0 && d_cache_B == 0 && touched; + const gdn_inputs in = make_inputs(c, seed++); - printf("K=%lld T=%lld n_seqs=%lld r_in=%lld raw=%d H_k=%lld nth=%d : attn diff A/B = %zu/%zu, " - "cache diff A/B = %zu/%zu, cache touched = %d -> %s\n", - (long long) K, (long long) T, (long long) n_seqs, (long long) r_in, raw, (long long) H_k, n_threads, - d_attn_A, d_attn_B, d_cache_A, d_cache_B, (int) touched, ok ? "OK" : "FAIL"); + const bool ok = run_abc(c, in, "[K=1 rollback read] "); n_cases += 1; n_failed += !ok; - }}}}}}} + }}}}}} + + // T = 1 cross-K on a 3-plane cache: the mixed T = 1 ubatch (a drafting seq with + // an empty draft next to opted-out seqs) runs with K = 3 and must leave every + // seq exactly what K = 1 would have: scores and the whole cache byte-identical. + // r_in in {0, 1, 2} covers the steady state and the pending-rollback read. + for (int64_t n_seqs : { (int64_t) 1, (int64_t) 4, (int64_t) 8 }) { + for (int64_t r_in : { (int64_t) 0, (int64_t) 1, (int64_t) 2 }) { + for (int raw = 0; raw < 2; ++raw) { + for (int64_t H_k : { H_v, H_v / 2 }) { + for (int n_threads : { 1, 4 }) { + gdn_case c; + c.S_v = S_v; + c.H_v = H_v; + c.H_k = H_k; + c.T = 1; + c.n_seqs = n_seqs; + c.K = 1; + c.n_planes = 3; + c.mem_size = n_seqs + 3; + c.head = 2; + c.r_in = r_in; + c.raw_gates = raw != 0; + c.n_threads = n_threads; + + // T = 1: rows_out holds min(T, K)*n_seqs = n_seqs rows for K = 1 and K = 3 alike + const gdn_inputs in = make_inputs(c, seed++); + + const bool ok = run_cross_k(c, in); + + n_cases += 1; + n_failed += !ok; + }}}}} printf("%d cases, %d failed\n", n_cases, n_failed); From 80164f75d9112725fe5ef9546e5931a27f7d0dd5 Mon Sep 17 00:00:00 2001 From: abdulaziz Date: Sat, 19 Sep 2026 05:30:51 +0300 Subject: [PATCH 10/15] server: document the per-request speculative.n_max Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01QrdPnP5KL449GsoWBZx52R --- tools/server/README.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tools/server/README.md b/tools/server/README.md index 93736c3edfa9..ec60c1b3e687 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -596,6 +596,8 @@ These words will not be included in the completion, so make sure to add them to `lora`: A list of LoRA adapters to be applied to this specific request. Each object in the list must contain `id` and `scale` fields. For example: `[{"id": 0, "scale": 0.5}, {"id": 1, "scale": 1.1}]`. If a LoRA adapter is not specified in the list, its scale will default to `0.0`. Please note that requests with different LoRA configurations will not be batched together, which may result in performance degradation. +`speculative.n_max`: Maximum number of tokens to draft for this request when the server runs a draft-model based speculative decoding method (e.g. `--spec-type draft-mtp`). `0` disables drafting for the request: the slot adds one token per step, decodes in one ubatch together with the other non-drafting slots and keeps a single recurrent-state plane (no rollback snapshots). `0 < n_max < --spec-draft-n-max` truncates the draft; larger values have no effect (the draft stops at the server-wide limit). `LLAMA_SERVER_SPEC_MULTISLOT=0` in the server environment restores the previous policy (the field is ignored and every slot drafts). Default: the server `--spec-draft-n-max` + **Response format** - Note: In streaming mode (`stream`), only `content`, `tokens` and `stop` will be returned until end of completion. Responses are sent using the [Server-sent events](https://html.spec.whatwg.org/multipage/server-sent-events.html) standard. Note: the browser's `EventSource` interface cannot be used due to its lack of `POST` request support. From b6d3cb7f5ff8ccf3cdb9b62ac62be29f4de152ef Mon Sep 17 00:00:00 2001 From: abdulaziz Date: Sat, 19 Sep 2026 05:33:29 +0300 Subject: [PATCH 11/15] memory-recurrent: only roll back onto planes the seq's last ubatch wrote A partial seq_rm accepted any rollback in 1..n_rs_seq and pointed the next read at that snapshot plane, whether or not the ubatch that set cell.pos had written it: a T = 1 ubatch writes plane 0 only, so a rollback of 1 after it would read the state an older ubatch left in plane 1 (reachable today from llama_context::decode's failed-ubatch cleanup, which removes [pos_min, inf) of a T = 1 ubatch). Track rs_n_valid per seq -- min(T, K) - 1, recorded when the ubatch is applied (not in find_slot, which prepare() dry-runs), copied by seq_cp, 0 after rm_all / clear / state_read -- and refuse a rollback beyond it. Never refuses a rollback the server asks for: a drafting slot's verify ubatch has T = 1 + n_draft and K = n_rs_seq + 1, so rs_n_valid = min(n_draft, n_rs_seq) >= n_rollback (accepted >= 1, and n_rollback > n_rs_seq already takes the checkpoint path). The failed-ubatch cleanup now returns false (its result is ignored) instead of installing a stale read. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01QrdPnP5KL449GsoWBZx52R --- src/llama-memory-recurrent.cpp | 48 ++++++++++++++++++++++++++++++++-- src/llama-memory-recurrent.h | 9 +++++++ 2 files changed, 55 insertions(+), 2 deletions(-) diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index d2e01b480749..a9749a0befc1 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -87,6 +87,7 @@ llama_memory_recurrent::llama_memory_recurrent( this->n_rs_seq = n_rs_seq; rs_idx.assign(n_seq_max, 0); rs_n_snap.assign(n_seq_max, n_rs_seq); // every seq keeps all planes unless the API lowers its budget + rs_n_valid.assign(n_seq_max, 0); cells.clear(); cells.resize(mem_size); @@ -271,6 +272,7 @@ void llama_memory_recurrent::clear(bool data) { } std::fill(rs_idx.begin(), rs_idx.end(), 0); + std::fill(rs_n_valid.begin(), rs_n_valid.end(), 0); } bool llama_memory_recurrent::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) { @@ -292,6 +294,7 @@ bool llama_memory_recurrent::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos const bool rm_all = p0 == 0 && p1 == std::numeric_limits::max(); if (rm_all) { set_rs_idx(seq_id, 0); + set_rs_n_valid(seq_id, 0); } // models like Mamba or RWKV can't have a state partially erased at the end @@ -307,11 +310,14 @@ bool llama_memory_recurrent::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos // partial rollback via per-token snapshot index, bounded by the seq's snapshot budget // (rs_n_snap[seq_id] <= n_rs_seq; a budget of 0 refuses like a memory with n_rs_seq == 0) + // and by the planes its last ubatch actually wrote (rs_n_valid: beyond it a plane holds + // the state of an older ubatch, so reading it would not be a rollback) if (0 < p0 && p0 <= cell.pos && p1 > cell.pos) { const llama_pos rollback = cell.pos - (p0 - 1); // pending rollback is single-use const bool pending = rs_idx[seq_id] != 0; - if (!pending && rollback >= 1 && rollback <= (llama_pos) rs_n_snap[seq_id]) { + if (!pending && rollback >= 1 && + rollback <= (llama_pos) rs_n_snap[seq_id] && rollback <= (llama_pos) rs_n_valid[seq_id]) { set_rs_idx(seq_id, (uint32_t) rollback); cell.pos = p0 - 1; return true; @@ -397,6 +403,13 @@ void llama_memory_recurrent::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id tail_dst.tail = tail_src.tail; } } + + // until its first ubatch the destination reads the source cell's planes (src0), so it inherits + // what the source's last ubatch wrote + if (seq_id_src >= 0 && seq_id_dst >= 0 && + (size_t) seq_id_src < rs_n_valid.size() && (size_t) seq_id_dst < rs_n_valid.size()) { + rs_n_valid[seq_id_dst] = rs_n_valid[seq_id_src]; + } } void llama_memory_recurrent::seq_keep(llama_seq_id seq_id) { @@ -551,6 +564,20 @@ bool llama_memory_recurrent::seq_rs_snapshots(llama_seq_id seq_id, uint32_t n_sn return true; } +void llama_memory_recurrent::set_rs_n_valid(llama_seq_id seq_id, uint32_t n) { + if (seq_id < 0) { + std::fill(rs_n_valid.begin(), rs_n_valid.end(), n); + return; + } + + assert(n_seq_max == rs_n_valid.size()); + + GGML_ASSERT((uint32_t) seq_id < n_seq_max); + GGML_ASSERT(n <= n_rs_seq); + + rs_n_valid[seq_id] = n; +} + std::map llama_memory_recurrent::memory_breakdown() const { std::map ret; for (const auto & [_, buf] : ctxs_bufs) { @@ -994,6 +1021,7 @@ void llama_memory_recurrent::state_read(llama_io_read_i & io, llama_seq_id seq_i if (n_rs_seq != 0) { set_rs_idx(seq_id, 0); + set_rs_n_valid(seq_id, 0); // a restored cell holds plane 0 only } } @@ -1357,7 +1385,23 @@ bool llama_memory_recurrent_context::apply() { return true; } - mem->find_slot(ubatches[i_next]); + const auto & ubatch = ubatches[i_next]; + + mem->find_slot(ubatch); + + // once the graph ran, planes 1..n_valid of every seq in the ubatch hold the states 1..n_valid + // tokens back (min(T, K) planes are written, K = 1 + get_n_snap()). Recorded here rather than + // in find_slot because prepare() dry-runs find_slot for every ubatch + if (mem->n_rs_seq > 0) { + const uint32_t n_valid = std::min(ubatch.n_seq_tokens, 1 + get_n_snap()) - 1; + + for (uint32_t s = 0; s < ubatch.n_seqs; ++s) { + const uint32_t i = s*ubatch.n_seq_tokens; + for (int32_t j = 0; j < ubatch.n_seq_id[i]; ++j) { + mem->set_rs_n_valid(ubatch.seq_id[i][j], n_valid); + } + } + } return true; } diff --git a/src/llama-memory-recurrent.h b/src/llama-memory-recurrent.h index 91473a6474c5..21676fa5c2a6 100644 --- a/src/llama-memory-recurrent.h +++ b/src/llama-memory-recurrent.h @@ -85,6 +85,15 @@ class llama_memory_recurrent : public llama_memory_i { // seq_rm accepts for it. Defaults to n_rs_seq; only the API changes it (not seq_rm/clear/state_read) std::vector rs_n_snap; + // per-seq count of snapshot planes actually written by the seq's last ubatch: planes + // 1..rs_n_valid hold the states 1..rs_n_valid tokens back from cell.pos, an older ubatch + // left whatever lies beyond. The largest partial rollback seq_rm can honour without reading + // a stale plane. Recorded when a ubatch is applied (min(T, K) - 1), copied by seq_cp, + // 0 after rm_all / clear / state_read + std::vector rs_n_valid; + + void set_rs_n_valid(llama_seq_id seq_id, uint32_t n); // seq_id < 0: every seq + // computed before each graph build uint32_t n = 0; From 0fd8b1d072f96c1b6879cbe173351777c9954d76 Mon Sep 17 00:00:00 2001 From: abdulaziz Date: Sat, 19 Sep 2026 06:30:45 +0300 Subject: [PATCH 12/15] llama: document the recurrent rollback invariant on llama_memory_seq_rm A micro-batch snapshots the state after each of its last min(T, K) - 1 tokens and never the state it started from (ops.cpp slot = n_tokens - 1 - t, in-place and gathered alike, upstream 1a07bfa included), so a partial removal of a sequence's tail is a rollback of n positions only while 1 <= n <= min(T - 1, n_snap) with T the token count of the sequence's last micro-batch: N single-token decodes followed by N positions back never had a snapshot to read and are refused since 317e4d0 instead of reading a plane an older micro-batch wrote. State it on the API. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01QrdPnP5KL449GsoWBZx52R --- include/llama.h | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/include/llama.h b/include/llama.h index 781a4251f2e1..20fd0157aa14 100644 --- a/include/llama.h +++ b/include/llama.h @@ -751,6 +751,14 @@ extern "C" { // seq_id < 0 : match any sequence [TAG_LLAMA_SEQ_ID_NEG] // p0 < 0 : [0, p1] // p1 < 0 : [p0, inf) + // Recurrent memory with rollback snapshots (n_rs_seq > 0): removing the tail [p0, inf) of a + // sequence is honoured as a rollback of n = pos_max - p0 + 1 positions only when + // 1 <= n <= min(T - 1, n_snap), where T is the token count of the sequence's last micro-batch + // and n_snap its snapshot budget (llama_memory_seq_rs_snapshots, default n_rs_seq): a micro-batch + // snapshots the state after each of its last min(T, n_rs_seq + 1) - 1 tokens and never the state + // it started from, so a rollback across earlier micro-batches (e.g. N single-token decodes, then + // N positions back) is refused rather than read from a stale plane. The rollback stays pending + // until the sequence's next decode consumes it; a second partial removal before that is refused LLAMA_API bool llama_memory_seq_rm( llama_memory_t mem, llama_seq_id seq_id, From cbb4c557be86fe6842ea56393175713e55f3b1f7 Mon Sep 17 00:00:00 2001 From: abdulaziz Date: Sat, 19 Sep 2026 06:30:45 +0300 Subject: [PATCH 13/15] tests: roll back only onto planes the last ubatch wrote in test-recurrent-state-rollback Two steps rolled back across earlier micro-batches, which the memory never snapshots (see llama_memory_seq_rm): the partial-flags checkpoint step replayed three positions one token at a time and then rolled the same three back -- on a6536c3 that passed only because plane 3, untouched by the T = 1 replays, still held the prompt ubatch's state for the identical tokens; since 317e4d0 it is refused (rs_n_valid == 0) and the test exited with 'partial rollback failed'. The multi-seq sub-test rolled back exactly the 3 tokens of a 3-token ubatch and read plane 3 of the 16-token prefill ubatch (the state after position 12, not 15) on every revision since the upstream merge. Partial step: extend both contexts with the same n_rollback + 1 tokens in ONE ubatch and roll back n_rollback, the largest rollback that ubatch supports; replay positions ext_pos + 1 .. ext_pos + 3 (still crossing DSV4's ratio-4 boundary). The 'full' replay logits are kept for the dirty-context check. Multi-seq: ctx_roll decodes [0, 15) then [15, 19) as one 4-token ubatch and rolls back 3 (plane 3 = state after 15); ctx_ref is restored from each seq's checkpoint at that point, which resolves the pending plane into a plain state, so both contexts replay from bit-identical bytes and the 1e-7 tolerance stays valid by construction. A forward-only reference cannot be compared bitwise: an honest rollback target sits inside a multi-token ubatch of ctx_roll and the CPU matmul kernels depend on the ubatch size (sgemm rejects small n). Not run here (model binaries are off limits in this session); established by construction from the kernel slot mapping and the memory's plane indexing. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01QrdPnP5KL449GsoWBZx52R --- tests/test-recurrent-state-rollback.cpp | 90 +++++++++++++++++++------ 1 file changed, 69 insertions(+), 21 deletions(-) diff --git a/tests/test-recurrent-state-rollback.cpp b/tests/test-recurrent-state-rollback.cpp index c6f599e584ca..31602a9a1dc9 100644 --- a/tests/test-recurrent-state-rollback.cpp +++ b/tests/test-recurrent-state-rollback.cpp @@ -17,10 +17,11 @@ static llama_context * make_ctx(const common_params & params, llama_model * mode return llama_init_from_model(model, cparams); } -static bool decode_tokens(llama_context * ctx, const std::vector & tokens, uint32_t count) { +// decode tokens[0, count) at positions [pos0, pos0 + count) of seq 0 as one batch, logits for the last +static bool decode_tokens(llama_context * ctx, const std::vector & tokens, uint32_t count, llama_pos pos0 = 0) { llama_batch batch = llama_batch_init(count, 0, 1); - for (uint32_t pos = 0; pos < count; ++pos) { - common_batch_add(batch, tokens[pos], pos, { 0 }, pos + 1 == count); + for (uint32_t i = 0; i < count; ++i) { + common_batch_add(batch, tokens[i], pos0 + (llama_pos) i, { 0 }, i + 1 == count); } const bool ok = llama_decode(ctx, batch) == 0; llama_batch_free(batch); @@ -39,14 +40,22 @@ static bool decode_one(llama_context * ctx, llama_token tok, llama_pos pos) { // per-seq token count exceeds n_ubatch: each seq's replay spans several // ubatches while its rollback restore is still pending. Compared against a // reference context that never advanced past the rollback point and decodes -// the identical replay batch. +// the identical replay batch. The reference is restored from each seq's +// checkpoint taken at the rollback point, so both contexts start the replay +// from bit-identical bytes and differ only in how the memory reads them +// (pending snapshot plane vs plane 0). A forward-only reference could not be +// compared bitwise: an honest rollback target always sits inside a multi-token +// ubatch of ctx_roll (a ubatch snapshots the state after each of its last +// min(T, K) - 1 tokens, never the state it started from), and the CPU matmul +// kernels depend on the ubatch size. static bool test_multi_seq_split_replay(const common_params & params, llama_model * model, const int n_vocab) { constexpr uint32_t n_seqs = 2; constexpr uint32_t n_ubatch = 16; constexpr uint32_t n_prompt = 19; constexpr uint32_t n_rollback = 3; constexpr uint32_t n_replay = 40; // > n_ubatch so each seq spans multiple ubatches - constexpr llama_pos p0 = n_prompt - n_rollback; + constexpr llama_pos p0 = n_prompt - n_rollback; // replay start; the rollback lands on the state after p0 - 1 + constexpr llama_pos p_tail = p0 - 1; // the last pre-rollback ubatch [p_tail, n_prompt) has n_rollback + 1 tokens const auto make_ctx_multi = [&]() { auto cparams = common_context_params_to_llama(params); @@ -83,18 +92,21 @@ static bool test_multi_seq_split_replay(const common_params & params, llama_mode bool ok = true; - // both contexts decode the identical [0, p0) prefill; only ctx_roll decodes - // the tail, which is then rolled back so its restore is pending at replay + // ctx_roll decodes [0, p_tail) and then [p_tail, n_prompt) as one ubatch of + // n_rollback + 1 tokens, rolled back by n_rollback -- the largest rollback + // that ubatch supports: plane n_rollback is the state after position p_tail, + // and the restore is pending at replay. ctx_ref receives that same state + // through a checkpoint (state_write resolves the pending plane into a plain + // plane-0 state) for (uint32_t s = 0; s < n_seqs && ok; ++s) { llama_batch batch = llama_batch_init(n_prompt, 0, 1); - for (llama_pos pos = 0; pos < (llama_pos) p0; ++pos) { + for (llama_pos pos = 0; pos < p_tail; ++pos) { common_batch_add(batch, tok(s, pos), pos, { (llama_seq_id) s }, false); } ok = ok && llama_decode(ctx_roll, batch) == 0; - ok = ok && llama_decode(ctx_ref, batch) == 0; common_batch_clear(batch); - for (llama_pos pos = p0; pos < (llama_pos) n_prompt; ++pos) { + for (llama_pos pos = p_tail; pos < (llama_pos) n_prompt; ++pos) { common_batch_add(batch, tok(s, pos), pos, { (llama_seq_id) s }, false); } ok = ok && llama_decode(ctx_roll, batch) == 0; @@ -104,6 +116,12 @@ static bool test_multi_seq_split_replay(const common_params & params, llama_mode // a second partial removal while one is pending must be refused ok = ok && !llama_memory_seq_rm(llama_get_memory(ctx_roll), (llama_seq_id) s, p0 - 1, -1); + + if (ok) { + common_prompt_checkpoint ckpt; + ckpt.update_tgt(ctx_roll, (llama_seq_id) s, 0); + ckpt.load_tgt(ctx_ref, (llama_seq_id) s, 0); + } } if (!ok) { fprintf(stderr, "%s : multi-seq prefill/rollback failed\n", __func__); @@ -274,7 +292,9 @@ int main(int argc, char ** argv) { const uint32_t n_tokens = tokens.size(); const llama_pos rollback_pos = (llama_pos) n_tokens - n_rollback; - // Decode the full prompt on the source, then roll back three positions. + // Decode the full prompt on the source as one ubatch (n_ubatch >= n_tokens), then roll back + // three positions: a ubatch snapshots the state after each of its last min(T, K) - 1 tokens, + // so plane 3 holds the state after position rollback_pos - 1. // Replaying them crosses DSV4's ratio-4 compressor boundary. // Rollback leaves the recurrent memory in a snapshot state (rs_idx != 0). if (!decode_tokens(ctx_src, tokens, n_tokens)) { @@ -292,12 +312,14 @@ int main(int argc, char ** argv) { ckpt.load_tgt(ctx_dst, 0, 0); constexpr float eps = 1e-5f; - std::vector> logits_src_replay(n_rollback); - const auto replay_and_compare = [&](const char * mode) { + // replay toks[0, n_rollback) one token at a time at positions [pos0, pos0 + n_rollback) on both + // contexts and compare the logits; logits_out (optional) keeps the source logits of each step + const auto replay_and_compare = [&](const char * mode, const std::vector & toks, llama_pos pos0, + std::vector> * logits_out) { for (uint32_t i = 0; i < n_rollback; ++i) { - const llama_pos pos = rollback_pos + i; - if (!decode_one(ctx_src, tokens[pos], pos) || - !decode_one(ctx_dst, tokens[pos], pos)) { + const llama_pos pos = pos0 + (llama_pos) i; + if (!decode_one(ctx_src, toks[i], pos) || + !decode_one(ctx_dst, toks[i], pos)) { fprintf(stderr, "%s : %s replay failed at position %d\n", __func__, mode, pos); return false; } @@ -309,7 +331,9 @@ int main(int argc, char ** argv) { return false; } - logits_src_replay[i].assign(logits_src, logits_src + n_vocab); + if (logits_out != nullptr) { + (*logits_out)[i].assign(logits_src, logits_src + n_vocab); + } for (int token = 0; token < n_vocab; ++token) { if (std::fabs(logits_src[token] - logits_dst[token]) > eps) { fprintf(stderr, "%s : %s logits mismatch at position %d, token %d (%g != %g)\n", @@ -320,12 +344,35 @@ int main(int argc, char ** argv) { } return true; }; - if (!replay_and_compare("full")) { + + std::vector> logits_src_replay(n_rollback); + const std::vector toks_replay(tokens.begin() + rollback_pos, tokens.end()); + if (!replay_and_compare("full", toks_replay, rollback_pos, &logits_src_replay)) { + return 1; + } + + // Partial-flags checkpoint of a pending rollback. The three single-token replays above wrote + // plane 0 only: a ubatch snapshots the state after each of its last min(T, K) - 1 tokens and + // never the state it started from, so rolling those positions back again would have to read + // planes an older ubatch wrote, and llama_memory_seq_rm refuses that. Extend both sequences + // with the same n_rollback + 1 tokens in ONE ubatch instead and roll back n_rollback, the + // largest rollback that ubatch supports: plane n_rollback is the state after its first token. + // The replay of positions ext_pos + 1 .. ext_pos + n_rollback again crosses DSV4's ratio-4 + // compressor boundary (n_tokens = n_rs_seq + 1 = 9, n_rollback = 3). + const llama_pos ext_pos = (llama_pos) n_tokens; + const uint32_t n_ext = n_rollback + 1; + std::vector ext(n_ext); + for (uint32_t i = 0; i < n_ext; ++i) { + ext[i] = tokens[(1 + i) % n_tokens]; + } + if (!decode_tokens(ctx_src, ext, n_ext, ext_pos) || !decode_tokens(ctx_dst, ext, n_ext, ext_pos)) { + fprintf(stderr, "%s : extension decode failed\n", __func__); return 1; } - if (!llama_memory_seq_rm(llama_get_memory(ctx_src), 0, rollback_pos, -1) || - !llama_memory_seq_rm(llama_get_memory(ctx_dst), 0, rollback_pos, -1)) { + const llama_pos rollback_pos_ext = ext_pos + 1; + if (!llama_memory_seq_rm(llama_get_memory(ctx_src), 0, rollback_pos_ext, -1) || + !llama_memory_seq_rm(llama_get_memory(ctx_dst), 0, rollback_pos_ext, -1)) { fprintf(stderr, "%s : partial rollback failed\n", __func__); return 1; } @@ -335,7 +382,8 @@ int main(int argc, char ** argv) { ckpt_partial.update_tgt(ctx_src, 0, partial_flags); ckpt_partial.load_tgt(ctx_dst, 0, partial_flags); - if (!replay_and_compare("partial")) { + const std::vector toks_ext(ext.begin() + 1, ext.end()); + if (!replay_and_compare("partial", toks_ext, rollback_pos_ext, nullptr)) { return 1; } From 2d2c3c3ffcd98198956a4111c139746fb8c9debc Mon Sep 17 00:00:00 2001 From: abdulaziz Date: Sat, 19 Sep 2026 06:30:45 +0300 Subject: [PATCH 14/15] examples: rs-rollback decodes the drafts as one verify ubatch so the rollback stays inside it Both harnesses over-generated with single-token decodes and then rolled back rewind (+1) positions: with fixed snapshot planes those decodes write plane 0 only, so the rollback read a plane the prompt ubatch had written (a wrong state, silently) and is refused since 317e4d0 ('ROLLBACK REFUSED'). Decode the accepted token together with rewind draft tokens as ONE ubatch (the server's verify shape, logits for the accepted token's row only) and drop the drafts: with T = rewind + 1 a rollback of rewind positions is the largest that ubatch supports and lands on the state after the accepted token; needs n_rs_seq >= rewind (was rewind + 1). The multi-seq harness keeps the other sequences in lockstep with an ordinary single-token step at that iteration. Header comments state the invariant and that a mismatch right at a rollback step can be matmul rounding between the (R + 1)-token and single-token kernels. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01QrdPnP5KL449GsoWBZx52R --- examples/rs-rollback/rs-rollback-multi.cpp | 97 +++++++++++++++------- examples/rs-rollback/rs-rollback.cpp | 80 ++++++++++-------- 2 files changed, 112 insertions(+), 65 deletions(-) diff --git a/examples/rs-rollback/rs-rollback-multi.cpp b/examples/rs-rollback/rs-rollback-multi.cpp index 2c810cf47452..a65f7eacd524 100644 --- a/examples/rs-rollback/rs-rollback-multi.cpp +++ b/examples/rs-rollback/rs-rollback-multi.cpp @@ -6,11 +6,20 @@ // that class of bug at all. // // pass A (reference): run S sequences concurrently, no rollback. -// pass B (rollback): identical, but sequence 0 periodically over-generates R -// tokens which are then seq_rm()'d away and re-generated. +// pass B (rollback): identical, but sequence 0 periodically decodes its next +// token together with R draft tokens as ONE ubatch (the +// server's verify shape) and then seq_rm()s the R draft +// positions away. // // PASS requires every sequence -- including the untouched ones -- to emit // exactly the reference token stream. +// +// The rollback stays inside the verify ubatch on purpose: a ubatch snapshots the +// state after each of its last min(T, K) - 1 tokens and never the state it +// started from, so a rollback across single-token steps is refused. At a +// rollback step sequence 0's logits come from an (R + 1)-token ubatch and the +// other sequences decode without it, so a mismatch exactly there can in rare +// near-tie cases be matmul rounding rather than a snapshot bug. #include "llama.h" #include @@ -153,8 +162,62 @@ int main(int argc, char ** argv) { out[s].push_back(greedy_ith(ctx, vocab, batch.n_tokens - 1)); } + const int n_vocab_run = llama_vocab_n_tokens(vocab); + + // invariant at the top of every iteration: out[s].back() is sampled-but-not-yet-decoded int since = 0; while ((int) out[0].size() < n_predict) { + if (with_rollback && since >= stride && (int) out[0].size() + rewind + 2 < n_predict) { + n_rollbacks++; + // seq 0 alone takes a verify-style step: its pending token plus `rewind` + // draft tokens as ONE ubatch (logits for the pending token's row only), + // then the drafts are dropped. A ubatch snapshots the state after each + // of its last min(T, K) - 1 tokens and never the state it started from, + // so with T = rewind + 1 the rollback of `rewind` positions is the + // largest this ubatch supports and lands on the state after the pending + // token -- needs n_rs_seq >= rewind. The drafts' content is irrelevant. + const int npos0 = npos[0]; + batch.n_tokens = 0; + for (int r = 0; r <= rewind; r++) { + const llama_token t = r == 0 ? out[0].back() : (llama_token) ((out[0].back() + 7*r + 1) % n_vocab_run); + batch.token[r] = t; + batch.pos[r] = npos0 + r; + batch.n_seq_id[r] = 1; + batch.seq_id[r][0] = 0; + batch.logits[r] = r == 0; + batch.n_tokens++; + } + if (llama_decode(ctx, batch)) { + fprintf(stderr, "decode(verify) failed\n"); + return false; + } + + const llama_pos keep = npos0 + 1; + if (!llama_memory_seq_rm(mem, 0, keep, -1)) { + fprintf(stderr, "ROLLBACK REFUSED: keep=%d npos0=%d rewind=%d\n", keep, npos0 + rewind + 1, rewind); + return false; + } + npos[0] = keep; + out[0].push_back(greedy_ith(ctx, vocab, 0)); + + // the other sequences take their ordinary single-token step + std::vector> others; + for (int s = 1; s < n_seqs; s++) { + others.push_back({ s, out[s].back() }); + } + if (!others.empty()) { + if (!submit(others)) { + fprintf(stderr, "decode(step) failed\n"); + return false; + } + for (int s = 1; s < n_seqs; s++) { + out[s].push_back(greedy_ith(ctx, vocab, s - 1)); + } + } + since = 0; + continue; + } + // one token per sequence, all in one batch std::vector> items; for (int s = 0; s < n_seqs; s++) { @@ -168,36 +231,6 @@ int main(int argc, char ** argv) { out[s].push_back(greedy_ith(ctx, vocab, s)); } since++; - - if (with_rollback && since >= stride && (int) out[0].size() + rewind + 2 < n_predict) { - n_rollbacks++; - // seq 0 alone over-generates, then discards - for (int r = 0; r < rewind; r++) { - std::vector> one = { - { 0, out[0].back() } - }; - if (!submit(one)) { - fprintf(stderr, "decode(draft) failed\n"); - return false; - } - out[0].push_back(greedy_ith(ctx, vocab, 0)); - } - for (int r = 0; r < rewind; r++) { - out[0].pop_back(); - } - - // invariant: out[s].back() is sampled-but-not-yet-decoded. The - // drafts occupied positions [npos0 - rewind, npos0), and that - // range also covers the decode of the still-pending token, so - // discarding it restores the exact pre-draft state -- no replay. - const llama_pos keep = npos[0] - rewind; - if (!llama_memory_seq_rm(mem, 0, keep, -1)) { - fprintf(stderr, "ROLLBACK REFUSED: keep=%d npos0=%d rewind=%d\n", keep, npos[0], rewind); - return false; - } - npos[0] = keep; - since = 0; - } } llama_batch_free(batch); diff --git a/examples/rs-rollback/rs-rollback.cpp b/examples/rs-rollback/rs-rollback.cpp index c03a72a69e98..0f9be8270310 100644 --- a/examples/rs-rollback/rs-rollback.cpp +++ b/examples/rs-rollback/rs-rollback.cpp @@ -1,16 +1,25 @@ -// Correctness harness for the recurrent-state snapshot ring (n_rs_seq > 0). +// Correctness harness for the recurrent-state snapshot planes (n_rs_seq > 0). // -// A forward-only generation never reads the snapshot slots back, so it cannot +// A forward-only generation never reads the snapshot planes back, so it cannot // catch a bad snapshot write. This does: // // pass A (reference): greedily generate N tokens straight through. -// pass B (rollback): generate the same N tokens, but every S tokens -// over-generate by R, then llama_memory_seq_rm() those R -// positions away and re-generate them. +// pass B (rollback): generate the same N tokens, but every S accepted tokens +// decode the next accepted token together with R draft +// tokens as ONE ubatch (the server's verify shape), then +// llama_memory_seq_rm() the R draft positions away. // // Pass B only lands on the same tokens if the state restored from the snapshot -// ring is exactly the state that produced them the first time. Any error in the +// plane is exactly the state after the accepted token. Any error in the // conv-window or SSM snapshot write shows up as a token mismatch. +// +// The rollback stays inside the verify ubatch on purpose: a ubatch snapshots the +// state after each of its last min(T, K) - 1 tokens and never the state it +// started from, so a rollback across single-token steps is refused. Right after +// a rollback step the accepted token's logits come from an (R + 1)-token ubatch +// instead of a single-token one; a mismatch exactly there can in rare near-tie +// cases be matmul rounding rather than a snapshot bug -- the first-mismatch +// index tells. #include "llama.h" #include @@ -39,7 +48,7 @@ int main(int argc, char ** argv) { int n_predict = 128; int stride = 16; // roll back every `stride` accepted tokens int rewind = 4; // how many tokens to discard and re-generate - int n_rs_seq = 8; // snapshot ring depth; must be >= rewind + 1 + int n_rs_seq = 8; // snapshot planes per seq; must be >= rewind for (int i = 1; i < argc; i++) { if (!strcmp(argv[i], "-m") && i + 1 < argc) { @@ -98,49 +107,54 @@ int main(int argc, char ** argv) { } int n_pos = n_prompt; + const int n_vocab = llama_vocab_n_tokens(vocab); + int since = 0; while ((int) out.size() < n_predict) { llama_token tok = greedy(ctx, vocab); out.push_back(tok); - - batch = llama_batch_get_one(&tok, 1); - if (llama_decode(ctx, batch)) { - fprintf(stderr, "decode failed\n"); - return false; - } - n_pos++; since++; if (with_rollback && since >= stride && (int) out.size() + rewind < n_predict) { n_rollbacks++; - // over-generate `rewind` tokens, then throw them away - for (int r = 0; r < rewind; r++) { - llama_token t = greedy(ctx, vocab); - batch = llama_batch_get_one(&t, 1); - if (llama_decode(ctx, batch)) { - fprintf(stderr, "decode(draft) failed\n"); - return false; - } - n_pos++; + // Verify-style step: decode the accepted token together with `rewind` + // draft tokens as ONE ubatch (logits for the accepted token's row only), + // then drop the drafts. A ubatch snapshots the state after each of its + // last min(T, K) - 1 tokens and never the state it started from, so with + // T = rewind + 1 a rollback of `rewind` positions is the largest this + // ubatch supports and lands on the state after the accepted token -- + // needs n_rs_seq >= rewind. The drafts' content is irrelevant. + llama_batch vb = llama_batch_init(rewind + 1, 0, 1); + for (int r = 0; r <= rewind; r++) { + vb.token[r] = r == 0 ? tok : (llama_token) ((tok + 7*r + 1) % n_vocab); + vb.pos[r] = n_pos + r; + vb.n_seq_id[r] = 1; + vb.seq_id[r][0] = 0; + vb.logits[r] = r == 0; + } + vb.n_tokens = rewind + 1; + const bool ok_verify = llama_decode(ctx, vb) == 0; + llama_batch_free(vb); + if (!ok_verify) { + fprintf(stderr, "decode(verify) failed\n"); + return false; } - // Drop the drafts AND the last accepted token in one go, then - // re-decode the accepted token so the logits are valid again. - // This forces the state to be restored from a snapshot taken - // (rewind + 1) steps back -- needs n_rs_seq >= rewind + 1. - const llama_pos keep = n_pos - rewind - 1; + + const llama_pos keep = n_pos + 1; if (!llama_memory_seq_rm(mem, 0, keep, -1)) { - fprintf(stderr, "ROLLBACK REFUSED: keep=%d n_pos=%d rewind=%d\n", keep, n_pos, rewind); + fprintf(stderr, "ROLLBACK REFUSED: keep=%d n_pos=%d rewind=%d\n", keep, n_pos + rewind + 1, rewind); return false; } n_pos = keep; - - batch = llama_batch_get_one(&out.back(), 1); + since = 0; + // the next greedy() reads the accepted token's logits, the batch's only output row + } else { + batch = llama_batch_get_one(&tok, 1); if (llama_decode(ctx, batch)) { - fprintf(stderr, "decode(replay) failed\n"); + fprintf(stderr, "decode failed\n"); return false; } n_pos++; - since = 0; } } llama_free(ctx); From 3ed7c4340d114c4e05c3aa5b7035b578e1bb7f8e Mon Sep 17 00:00:00 2001 From: abdulaziz Date: Sat, 19 Sep 2026 06:30:45 +0300 Subject: [PATCH 15/15] speculative: no rollback snapshot planes in the draft context The draft context inherited n_rs_seq = draft.n_max from common_context_params_to_llama, so common_context_can_seq_rm classified a recurrent drafter as RS and the server / speculative-simple removed the drafted tail with a partial seq_rm (server-context.cpp:2997, speculative-simple.cpp:209) instead of restoring the checkpoint. Every drafter decodes its draft one token per step (speculative.cpp draft loop), so the memory never held that snapshot: on a6536c3 the draft context continued from a stale plane (garbage drafts after the first round, silently); since 317e4d0 the refusal would GGML_ABORT at server-context.cpp:2998. With n_rs_seq = 0 a recurrent drafter is classified FULL and takes the existing checkpoint path (update_dft before drafting, load_dft before the now no-op seq_rm), the only correct one; a KV-cache drafter (transformer, the qwen35 / nemotron-h MTP contexts, dspark) is classified PART by the same probe it already ran and ignores n_rs_seq, and the target keeps its window. Also drops K planes of recurrent state per cell from the draft context. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01QrdPnP5KL449GsoWBZx52R --- common/speculative.cpp | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index 4de75a399494..f8cdc3741617 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -3148,7 +3148,14 @@ common_speculative_init_result::common_speculative_init_result( // the draft context holds as many tokens per sequence as the target context cparams.n_ctx = llama_n_ctx(ctx_tgt); - // n_rs_seq stays as common_context_params_to_llama set it: the draft context needs the same rollback window as the target, with n_rs_seq == 0 its seq_rm fails silently on partial acceptance and keeps stale positions + // no rollback snapshot planes in the draft context: every drafter decodes its draft one token per + // step (or drives a plain KV cache), and a recurrent memory snapshots only the tokens of a + // sequence's last micro-batch (see llama_memory_seq_rm), so the partial removal of a draft can + // never be honoured from a snapshot -- with planes it used to read a stale one; without them + // common_context_can_seq_rm classifies a recurrent drafter as FULL and the server / + // speculative-simple take the checkpoint path, the only correct one. A KV-cache drafter is + // classified PART as before; the target keeps its window (common_context_params_to_llama) + cparams.n_rs_seq = 0; cparams.ctx_other = ctx_tgt; std::string model_path;