train: one rule for the recurrent states a graph keeps, so training runs in a context with rollback slots

src/llama-context-opt.cpp

@@ -317,6 +317,14 @@ bool llama_context::opt_walk::train_chunk(uint32_t pos_ctx, uint32_t n_chunk) {
                 }
             }
             for (const auto & ws : res->t_walk_state) {
+                // a state read without its write reporting the exit (build_walk_state_exit) would drop
+                // the gradient later chunks put on it: a wrong step, silently. Refused by name.
+                if (ws.tap && ws.gs == nullptr && lctx.cparams.walk_state_surrogate) {
+                    ggml_free(ctx_compute_opt);
+                    lctx.opt_fail(std::string("this model's recurrent state write (") + ws.cache->name
+                        + ") does not report its exit to the walk (llm_graph_context::build_walk_state_exit): its gradient across chunks would be dropped");
+                    return false;
+                }
                 if (ws.tap) {
                     ggml_opt_accumulate_grad(opt_ctx, ws.tap, acc.state.at(ws.cache));
                 }
@@ -406,12 +414,13 @@ void llama_context::opt_walk::run_exact() {
     // Each window starts from the REQUESTED horizon: a halving that fit one window's device budget
     // must not shrink every later window and epoch, which may be shorter (Fable).
     lctx.opt_walk_horizon = lctx.opt_walk_horizon_req;
-    // A recurrent context with rollback slots keeps several states per sequence; the walk carries
-    // exactly one. Refused by name, never asserted: training runs inside the serving process, and
-    // an assert there takes serving down (2026-10-07 04:09Z on the 5090).
-    if (recr != nullptr && lctx.cparams.n_rs_seq > 0) {
-        lctx.opt_fail("the exact walk needs a training context with no recurrent rollback slots (n_rs_seq = 0, it has "
-            + std::to_string(lctx.cparams.n_rs_seq) + "): create the training context with n_rs_seq = 0");
+    // A serving context keeps rollback slots for drafts; the walk's graphs keep one state, the live
+    // one (llm_graph_context::rs_keep), so they train in that context. Only a PENDING undo reads a
+    // snapshot instead of the live state: that alone is refused, by name, never asserted, since
+    // training runs inside the serving process (2026-10-07 04:09Z on the 5090).
+    if (recr != nullptr && recr->rollback_pending(0)) {
+        lctx.opt_fail("the exact walk cannot train a sequence with a pending recurrent rollback (n_rs_seq = "
+            + std::to_string(lctx.cparams.n_rs_seq) + "): finish or drop the draft first");
         return;
     }
     plan.mark_exact(recr != nullptr, lctx.opt_walk_horizon);

src/llama-graph.cpp

@@ -3690,6 +3690,11 @@ ggml_tensor * llm_graph_context::build_rs(
                     get_state_rows);
 }
 
+int64_t llm_graph_context::rs_keep() const {
+    const bool training = cparams.training && cparams.walk_exact;
+    return cparams.n_rs_seq > 0 && !training ? (int64_t) cparams.n_rs_seq + 1 : 1;
+}
+
 void llm_graph_context::build_walk_state_exit(ggml_tensor * cache, ggml_tensor * exit) const {
     if (!(cparams.training && cparams.walk_exact && cparams.walk_state_surrogate)) {
         return;

src/llama-graph.h

@@ -1367,6 +1367,15 @@ struct llm_graph_context {
                 int32_t   n_seqs,
             const llm_graph_get_rows_fn & get_state_rows = ggml_get_rows) const;
 
+    // THE ONE RULE for how many recurrent states a graph keeps per sequence, read by every family's
+    // state write (delta-net, mamba, lfm2, bailing) and by the ops that produce snapshots. Serving with
+    // rollback slots keeps K = n_rs_seq + 1 (group g = the state g tokens back, so a draft can be undone);
+    // a context without them, and a TRAINING graph, keep one: the live state, group 0, which is the same
+    // bytes in both layouts. So training reads and writes the state serving reads, in the serving
+    // context; the undo groups it leaves stale are rewritten by the next serving decode, and the walk
+    // refuses a sequence with an undo pending (llama_memory_recurrent::rollback_pending).
+    int64_t rs_keep() const;
+
     // THE EXACT WALK: the state a chunk leaves (`exit`, what is written back to `cache`) meets the
     // gradient later chunks put on it, through the surrogate <exit, gs>. A no-op outside the walk.
     void build_walk_state_exit(ggml_tensor * cache, ggml_tensor * exit) const;

src/llama-memory-recurrent.h

@@ -78,6 +78,12 @@ class llama_memory_recurrent : public llama_memory_i {
 
     void set_rs_idx(llama_seq_id seq_id, uint32_t idx);
 
+    // an undo of a draft is pending on this sequence: its next read takes a snapshot group, not the
+    // live state (group 0) a training graph reads and writes (llm_graph_context::rs_keep)
+    bool rollback_pending(llama_seq_id seq_id) const {
+        return n_rs_seq > 0 && seq_id >= 0 && (size_t) seq_id < rs_idx.size() && rs_idx[seq_id] != 0;
+    }
+
     // computed before each graph build
     uint32_t n = 0;
 

src/models/bailingmoe3.cpp

@@ -183,7 +183,7 @@ static ggml_tensor * bailingmoe3_causal_conv1d(
         int64_t n_tokens,
         int64_t cache_head,
         uint32_t mem_size,
-        uint32_t n_rs_seq) {
+        int64_t K) { // states kept per sequence: llm_graph_context::rs_keep
     const int64_t d_inner = head_dim * n_head;
     const int64_t conv_state_size = (d_conv - 1) * d_inner;
     const int64_t total_state_size = 3 * conv_state_size;
@@ -197,7 +197,6 @@ static ggml_tensor * bailingmoe3_causal_conv1d(
     x_proj = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs);
     ggml_tensor * conv_x = ggml_concat(ctx0, conv_state, ggml_transpose(ctx0, x_proj), 0);
 
-    const int64_t K = (int64_t) n_rs_seq + 1;
     const int64_t n_written = std::min<int64_t>(n_seq_tokens, K);
 
     for (int64_t slot = 0; slot < n_written; ++slot) {
@@ -262,13 +261,13 @@ llama_model_bailingmoe3::graph::graph(const llama_model & model, const llm_graph
 
             ggml_tensor * q = bailingmoe3_causal_conv1d(
                     gf, ctx0, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv,
-                    d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq);
+                    d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, rs_keep());
             ggml_tensor * k = bailingmoe3_causal_conv1d(
                     gf, ctx0, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv,
-                    d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq);
+                    d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, rs_keep());
             ggml_tensor * v = bailingmoe3_causal_conv1d(
                     gf, ctx0, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv,
-                    d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq);
+                    d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, rs_keep());
 
             ggml_tensor * gate = ggml_mul_mat(ctx0, layer.ssm_f_a, cur);
             gate = ggml_add(ctx0, gate, layer.ssm_dt_b);

src/models/delta-net-base.cpp

@@ -476,7 +476,7 @@ 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) {
+    if (rs_keep() == 1) {
         const int64_t s_idx  = conv_input->ne[0] - conv_states->ne[0];
         const int64_t s_slot = 0;
 
@@ -496,12 +496,12 @@ ggml_tensor * llm_build_delta_net_base::build_conv_state(
         ggml_build_forward_expand(gf, ggml_cpy(ctx0, conv_state_last, conv_state_update));
         build_walk_state_exit(conv_states_all, conv_state_last);
     } else {
-        GGML_ASSERT(!cparams.walk_exact && "the exact walk keeps one recurrent state per sequence (n_rs_seq = 0)");
+        GGML_ASSERT(!cparams.walk_exact && "a training graph keeps one recurrent state per sequence (rs_keep)");
         // [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;
+        const int64_t K = rs_keep();
 
         for (int64_t t = 1; t <= K; ++t) {
             const int64_t s_idx  = std::max<int64_t>(0, conv_input->ne[0] - conv_states->ne[0] - K + t);
@@ -545,7 +545,7 @@ ggml_tensor * llm_build_delta_net_base::build_recurrent_attn(
     const int64_t n_seqs       = s->ne[3];
     const int64_t n_seq_tokens = q->ne[2];
 
-    const bool keep = cparams.n_rs_seq > 0;
+    const bool keep = rs_keep() > 1;
 
     if (!keep) {
         auto attn_out = build_delta_net(q, k, v, g, b, s, il);
@@ -564,7 +564,7 @@ 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 int64_t K = rs_keep();
 
     // state s is 4D [S_v, S_v, H_v, n_seqs]; K snapshot slots are written into the output.
     ggml_tensor * gdn_out = ggml_gated_delta_net(ctx0, q, k, v, g, b, s, K);

src/models/lfm2.cpp

@@ -205,7 +205,7 @@ llama_model_lfm2::graph<iswa>::graph(const llama_model & model, const llm_graph_
         GGML_ASSERT(bx->ne[0] > conv->ne[0]);
 
         // write conv states: slot 0 = the final state, slot s = the state s tokens back (partial rollback)
-        const int64_t K         = hparams.causal_attn && cparams.n_rs_seq > 0 ? (int64_t) cparams.n_rs_seq + 1 : 1;
+        const int64_t K         = hparams.causal_attn ? rs_keep() : 1;
         const int64_t n_written = std::min<int64_t>(n_seq_tokens, K);
         const auto    mem_size  = mctx_cur->get_size();
         const size_t  row_size  = ggml_row_size(conv_state->type, (int64_t) d_conv * n_embd);

src/models/mamba-base.cpp

@@ -167,7 +167,7 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp,
     const int64_t n_seqs   = ubatch.n_seqs;
 
     const int64_t n_seq_tokens = ubatch.n_seq_tokens;
-    const int64_t K            = cparams.n_rs_seq > 0 ? (int64_t) cparams.n_rs_seq + 1 : 1;
+    const int64_t K            = rs_keep();
 
     GGML_ASSERT(n_seqs != 0);
     GGML_ASSERT(ubatch.equal_seqs());

tests/test-walk-exact.cpp

@@ -463,36 +463,21 @@ int main(int argc, char ** argv) {
         }
     }
     if (llama_model_is_hybrid(model)) {
-        // 4. regression for the 5090 crash (2026-10-07 04:09Z): a recurrent context with rollback
-        // slots (serving's, inherited by the training context) made the exact walk ASSERT inside
-        // the serving process. It must refuse by name, and the adapter stays untouched.
-        llama_adapter_lora * adapter = llama_adapter_lora_init(model, init.c_str());
-        const std::vector<float> before = adapter_params(adapter);
+        // 4. TRAINING IN THE SERVING CONTEXT (Joel, 2026-10-13: "training is basically llama inference with
+        // backprop"). Serving keeps recurrent rollback slots for drafts; a training graph keeps one state, the
+        // live one (llm_graph_context::rs_keep), so a context WITH rollback slots takes the same step as one
+        // without. Was a refusal (the 5090 crash, 2026-10-07 04:09Z), which forced a second training
+        // context beside serving.
+        std::vector<uint8_t> all(WINDOW + 1, 1);
+        const auto without = step_delta(params, model, init, tokens, all, WINDOW / 4, true);
         g_n_rs_seq = 4;
-        llama_context * ctx = make_ctx(params, model, WINDOW / 4);
+        const auto with    = step_delta(params, model, init, tokens, all, WINDOW / 4, true);
         g_n_rs_seq = 0;
-        float scale = 1.0f;
-        GGML_ASSERT(llama_set_adapters_lora(ctx, &adapter, 1, &scale) == 0);
-        llama_opt_params lopt{};
-        lopt.param_filter   = llama_opt_param_filter_all;
-        lopt.get_opt_pars   = sgd_pars;
-        lopt.optimizer_type = GGML_OPT_OPTIMIZER_TYPE_SGD;
-        lopt.adapter        = adapter;
-        lopt.walk_exact     = true;
-        llama_opt_init(ctx, model, lopt);
-        std::vector<std::vector<llama_token>> seqs = { std::vector<llama_token>(tokens.begin(), tokens.begin() + WINDOW + 1) };
-        std::vector<std::vector<uint8_t>>     loss = { std::vector<uint8_t>(WINDOW + 1, 1) };
-        ggml_opt_dataset_t dataset = common_opt_dataset_init_masked(WINDOW, seqs, loss, tokens[0]);
-        llama_opt_epoch(ctx, dataset, nullptr, nullptr, /*idata_split =*/ 1, nullptr, nullptr);
-        const std::string why = llama_opt_failure(ctx);
-        printf("  rollback slots: failed=%d, \"%s\"\n", (int) llama_opt_failed(ctx), why.c_str());
-        if (!llama_opt_failed(ctx) || why.find("rollback") == std::string::npos || adapter_params(adapter) != before) {
-            fprintf(stderr, "FAILED: a recurrent context with rollback slots was not refused by name\n");
+        if (!same_step(compare("rollback slots: exact walk in a context with n_rs_seq 4 vs 0", with, without))) {
+            fprintf(stderr, "FAILED: rollback slots changed the exact walk's step\n");
+            worst_tensors(with, without, 24);
             ++failures;
         }
-        ggml_opt_dataset_free(dataset);
-        llama_free(ctx);
-        llama_adapter_lora_free(adapter);
     }
 
     {