train: the training context keeps no recurrent rollback slots; the exact walk refuses them by name, never asserts

src/llama-context.cpp

@@ -3758,6 +3758,17 @@ void llama_context::opt_epoch_iter(
         // 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).
         opt_walk_horizon = 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 && cparams.n_rs_seq > 0) {
+            opt_failure = "the exact walk needs a training context with no recurrent rollback slots (n_rs_seq = 0, it has "
+                + std::to_string(cparams.n_rs_seq) + "): create the training context with n_rs_seq = 0";
+            LLAMA_LOG_ERROR("%s: %s\n", __func__, opt_failure.c_str());
+            opt_alloc_failed.store(true);
+            opt_stop_requested.store(true);
+            return;
+        }
         // A recurrent model's state is carried the same way, one chunk at a time: the state
         // entering a chunk is a GRAD leaf, the state it leaves meets the next chunk's gradient
         // (build_rs, build_walk_state_exit). The reverse pass restores each chunk's entry state
@@ -3856,19 +3867,26 @@ void llama_context::opt_epoch_iter(
         // snapshotted at each chunk's start (empty before the first: it starts from zero)
         struct state_snapshot {
             bool                            empty = true;
-            std::vector<std::vector<float>> r, s; // per layer
+            llama_pos                       pos   = -1; // the cell's last position at the snapshot
+            std::vector<std::vector<float>> r, s;       // per layer
         };
         std::vector<state_snapshot> snaps(recr != nullptr ? chunks.size() : 0);
         auto state_rows = [&](bool write, state_snapshot & snap) {
             const int32_t cell = recr->cells[0].tail; // the cell holding sequence 0's state
             if (!write) {
                 snap.empty = cell < 0;
+                snap.pos   = cell < 0 ? -1 : recr->cells[cell].pos;
                 snap.r.assign(recr->r_l.size(), {});
                 snap.s.assign(recr->s_l.size(), {});
             }
             if (cell < 0) {
                 return;
             }
+            // the restored state is the state AFTER snap.pos, so the cell says so: left at the
+            // window's end, the next ubatch read as non-consecutive ("position 65640 after 66345")
+            if (write) {
+                recr->cells[cell].pos = snap.pos;
+            }
             for (size_t il = 0; il < recr->r_l.size(); ++il) {
                 for (int kind = 0; kind < 2; ++kind) {
                     ggml_tensor * t = kind == 0 ? recr->r_l[il] : recr->s_l[il];

tests/test-walk-exact.cpp

@@ -21,6 +21,7 @@
 #include "../src/llama-adapter.h"
 
 #include <cmath>
+#include <cstring>
 #include <map>
 #include <cstdlib>
 #include <cstdio>
@@ -41,6 +42,21 @@ static ggml_opt_optimizer_params sgd_pars(void *) {
 // q8_0, in a flash-attention context that stores V untransposed) measure the straight-through
 // approximation the walk makes there (Fable on #47: the gradient is taken at the cache's values)
 static ggml_type g_cache_type = GGML_TYPE_F32;
+// recurrent rollback slots in the context (serving keeps some for speculative decoding)
+static uint32_t g_n_rs_seq = 0;
+
+// counts the recurrent memory's "non-consecutive token position" warnings: a state restore that
+// leaves the cell's position at the window's end makes every next chunk read as non-consecutive
+// (the 5090 log, 2026-10-07: "position 65640 after 66345"), passing everything else through
+static int g_nonconsecutive = 0;
+static void count_log(ggml_log_level level, const char * text, void * ud) {
+    if (text && strstr(text, "non-consecutive token position")) {
+        ++g_nonconsecutive;
+    }
+    GGML_UNUSED(level);
+    GGML_UNUSED(ud);
+    fputs(text, stderr);
+}
 
 static llama_context * make_ctx(const common_params & params, llama_model * model, uint32_t chunk) {
     auto cparams = common_context_params_to_llama(params);
@@ -54,6 +70,7 @@ static llama_context * make_ctx(const common_params & params, llama_model * mode
     cparams.flash_attn_type = g_cache_type == GGML_TYPE_F32 ? LLAMA_FLASH_ATTN_TYPE_DISABLED : LLAMA_FLASH_ATTN_TYPE_ENABLED;
     cparams.type_k          = g_cache_type;
     cparams.type_v          = g_cache_type;
+    cparams.n_rs_seq        = g_n_rs_seq;
     return llama_init_from_model(model, cparams);
 }
 
@@ -282,6 +299,55 @@ int main(int argc, char ** argv) {
         llama_adapter_lora_free(adapter);
     }
 
+    if (llama_model_is_hybrid(model)) {
+        // 5. what this catches: a state restore that does not restore the recurrent cell's
+        // position. The reverse pass restores each chunk's entry state; with the cell left at the
+        // window's end, every chunk after the first restore reads as non-consecutive.
+        std::vector<uint8_t> all(WINDOW + 1, 1);
+        all[0] = 0;
+        g_nonconsecutive = 0;
+        llama_log_set(count_log, nullptr);
+        step_delta(params, model, init, tokens, all, WINDOW / 4, true);
+        llama_log_set(nullptr, nullptr); // the default logger again
+        printf("  state restore: %d non-consecutive position warnings in the reverse pass\n", g_nonconsecutive);
+        if (g_nonconsecutive != 0) {
+            fprintf(stderr, "FAILED: a restored recurrent state left its cell at another position\n");
+            ++failures;
+        }
+    }
+    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);
+        g_n_rs_seq = 4;
+        llama_context * ctx = make_ctx(params, model, WINDOW / 4);
+        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");
+            ++failures;
+        }
+        ggml_opt_dataset_free(dataset);
+        llama_free(ctx);
+        llama_adapter_lora_free(adapter);
+    }
+
     std::remove(init.c_str());
     llama_model_free(model);
     llama_backend_free();

tools/server/server-train.cpp

@@ -751,6 +751,11 @@ void server_trainer::run(json req, examples_data ex) {
 
     // The training context: the SAME model, its own graph; flash attention has no backward.
     llama_context_params cparams = common_context_params_to_llama(params_base);
+    // NO ROLLBACK SLOTS. Serving keeps recurrent-state snapshots for speculative decoding's
+    // rollback (n_rs_seq, from the MTP draft); a training context never rolls back, and the exact
+    // walk carries one recurrent state per sequence. Inherited, they crashed the 5090's serving
+    // process on Kimi's first exact run (2026-10-07 04:09Z: the delta-net invariant asserted).
+    cparams.n_rs_seq = 0;
     // The context holds the whole window; one training chunk is one batch is one ubatch, the
     // largest multiple of 256 that divides the window and is at most the caller's "chunk" (the
     // memory gate's S). Context she did not write is decoded; her replies train chunk by chunk,