memory: refuse rollbacks past the snapshots of the last ubatch (#31368)

* memory: refuse rollbacks past the snapshots of the last ubatch

The recurrent memory and the DSV4 cache accepted any rollback up to
n_rs_seq, whether or not the last ubatch wrote that snapshot. Each now
records the deepest rollback its last ubatch guarantees and seq_rm
refuses anything past it:

- recurrent: n - 1 for a ubatch of n <= n_rs_seq + 1 tokens, 0 for a
  longer one, which only guarantees the final state
- DSV4: min(n, n_rs_seq), the snapshots are rebuilt from the state
  before the ubatch

The dead rollback block of the rollback test becomes the assertion it
asked for, and rollback_bound checks the exact bound of each memory
after ubatches of 1, 2, n_rs_seq + 1 and n_rs_seq + 2 tokens.

* tests: address review from ggerganov

Detect architectures with llama_arch_from_string in comp_padding and
arch_reserves_single_seq, as in rollback_bound.

src/llama-kv-cache-dsv4.cpp

@@ -1261,7 +1261,8 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4(
     hparams_lid(model.hparams),
     n_seq_max(n_seq_max),
     n_rs_seq(n_rs_seq),
-    rs_idx(n_seq_max, 0) {
+    rs_idx(n_seq_max, 0),
+    rs_max(n_seq_max, 0) {
 
     const layer_filter_cb filter_raw = [&](int32_t il) {
         if (filter && !filter(il)) {
@@ -1618,12 +1619,8 @@ bool llama_kv_cache_dsv4::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1
             return res;
         }
 
-        if (n_rs_seq == 0) {
-            return false;
-        }
-
         const llama_pos rollback = pos_max - (p0 - 1);
-        if (rollback < 1 || rollback > (llama_pos) n_rs_seq) {
+        if (rollback < 1 || rollback > (llama_pos) rs_max[seq_id]) {
             return false;
         }
 
@@ -1663,6 +1660,7 @@ void llama_kv_cache_dsv4::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_ds
 
     if (seq_id_src != seq_id_dst) {
         rs_idx[seq_id_dst] = 0;
+        rs_max[seq_id_dst] = 0;
     }
 }
 
@@ -1807,8 +1805,10 @@ void llama_kv_cache_dsv4::state_read(llama_io_read_i & io, llama_seq_id seq_id,
     if (seq_id >= 0) {
         GGML_ASSERT((uint32_t) seq_id < n_seq_max);
         rs_idx[seq_id] = 0;
+        rs_max[seq_id] = 0;
     } else {
         std::fill(rs_idx.begin(), rs_idx.end(), 0);
+        std::fill(rs_max.begin(), rs_max.end(), 0);
     }
 }
 
@@ -1852,20 +1852,32 @@ const std::vector<uint32_t> & llama_kv_cache_dsv4::get_rs_idx() const {
     return rs_idx;
 }
 
-void llama_kv_cache_dsv4::reset_rs_idx_for_ubatches(const std::vector<llama_ubatch> & ubatches) {
+void llama_kv_cache_dsv4::update_rs_for_ubatches(const std::vector<llama_ubatch> & ubatches) {
     if (n_rs_seq == 0) {
         return;
     }
 
+    std::vector<uint32_t> n_seq_tokens(n_seq_max);
+
     for (const llama_ubatch & ubatch : ubatches) {
+        std::fill(n_seq_tokens.begin(), n_seq_tokens.end(), 0);
+
         for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
             for (int32_t s = 0; s < ubatch.n_seq_id[i]; ++s) {
                 const llama_seq_id seq_id = ubatch.seq_id[i][s];
                 if (seq_id >= 0 && (uint32_t) seq_id < n_seq_max) {
-                    rs_idx[seq_id] = 0;
+                    n_seq_tokens[seq_id]++;
                 }
             }
         }
+
+        // the ubatch consumes any pending rollback and leaves snapshots for its own tokens only
+        for (uint32_t seq_id = 0; seq_id < n_seq_max; ++seq_id) {
+            if (n_seq_tokens[seq_id] > 0) {
+                rs_idx[seq_id] = 0;
+                rs_max[seq_id] = std::min(n_seq_tokens[seq_id], n_rs_seq);
+            }
+        }
     }
 }
 
@@ -1899,8 +1911,10 @@ void llama_kv_cache_dsv4::clear_compressed(llama_seq_id seq_id, bool data) {
 
     if (seq_id >= 0) {
         rs_idx[seq_id] = 0;
+        rs_max[seq_id] = 0;
     } else {
         std::fill(rs_idx.begin(), rs_idx.end(), 0);
+        std::fill(rs_max.begin(), rs_max.end(), 0);
     }
 }
 
@@ -2235,7 +2249,7 @@ llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context(
     lid_state(kv->get_lid_state()),
     status(ctx_raw->get_status()),
     kv(kv) {
-    kv->reset_rs_idx_for_ubatches(this->ubatches);
+    kv->update_rs_for_ubatches(this->ubatches);
 }
 
 llama_kv_cache_dsv4_context::~llama_kv_cache_dsv4_context() = default;

src/llama-kv-cache-dsv4.h

@@ -155,7 +155,7 @@ public:
 
     uint32_t get_n_rs_seq() const;
     const std::vector<uint32_t> & get_rs_idx() const;
-    void reset_rs_idx_for_ubatches(const std::vector<llama_ubatch> & ubatches);
+    void update_rs_for_ubatches(const std::vector<llama_ubatch> & ubatches);
 
 private:
     llama_hparams hparams_raw;
@@ -168,6 +168,10 @@ private:
 
     std::vector<uint32_t> rs_idx;
 
+    // [n_seq_max] largest rollback whose snapshot the last ubatch of the seq wrote: snapshot d is
+    // exact for d <= n tokens of the seq in that ubatch (d == n is the state before the ubatch)
+    std::vector<uint32_t> rs_max;
+
     // [n_layer] storage layer -> layer whose rows it holds (itself, or the source of a mirror)
     std::vector<int32_t> comp_il_src;
 

src/llama-memory-recurrent.cpp

@@ -148,6 +148,7 @@ void llama_memory_recurrent::clear(bool data) {
         cells[i].seq_id.clear();
         cells[i].src = -1;
         cells[i].tail = -1;
+        cells[i].rs_max = 0;
     }
 
     head = 0;
@@ -194,7 +195,7 @@ 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 snapshots of the last ubatch)
             if (0 < p0 && p0 <= cell.pos && p1 > cell.pos) {
                 // a cell shared by several sequences cannot move back for only one of them
                 if (cell.seq_id.size() > 1) {
@@ -208,7 +209,7 @@ bool llama_memory_recurrent::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_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) cell.rs_max) {
                     set_rs_idx(seq_id, (uint32_t) rollback);
                     cell.pos = p0 - 1;
                     return true;
@@ -672,6 +673,7 @@ bool llama_memory_recurrent::find_slot(const llama_ubatch & ubatch) {
                 __func__, last_pos, cell.pos, ubatch.seq_id[i][0], n_seq_tokens);
         }
         cell.pos = last_pos;
+        cell.rs_max = n_seq_tokens <= n_rs_seq + 1 ? n_seq_tokens - 1 : 0;
         cell.seq_id.clear();
         for (int32_t j = 0; j < ubatch.n_seq_id[i]; ++j) {
             const llama_seq_id seq_id = ubatch.seq_id[i][j];

src/llama-memory-recurrent.h

@@ -92,6 +92,10 @@ public:
         int32_t   src0 = -1; // like src, but only used when setting the inputs (allowing to copy once)
         int32_t   tail = -1;
 
+        // largest rollback whose snapshot the last ubatch guarantees: a ubatch of n <= n_rs_seq + 1 tokens
+        // writes snapshots 0..n-1, a longer one only guarantees the final state (ggml_gated_delta_net)
+        uint32_t rs_max = 0;
+
         std::set<llama_seq_id> seq_id;
 
         bool has_seq_id(const llama_seq_id & id) const {

tests/llama-context/cases-kv.cpp

@@ -7,11 +7,10 @@
 
 #include "common.h"
 
-#include "../../src/llama-ext.h" // llama_model_set_attn_rot
+#include "../../src/llama-ext.h" // llama_model_set_attn_rot, llama_arch_from_string
 
 #include <algorithm>
 #include <cstdio>
-#include <cstring>
 #include <fstream>
 #include <functional>
 #include <iterator>
@@ -770,9 +769,10 @@ struct case_comp_padding : ctx_test {
     }
 
     std::string skip_reason(model_run & mr) const override {
-        char arch[64] = {};
-        llama_model_meta_val_str(mr.model, "general.architecture", arch, sizeof(arch));
-        return strcmp(arch, "deepseek4") == 0 || strcmp(arch, "deepseek41") == 0 ? "" : "requires DeepSeek V4 or V4.1";
+        char arch_name[64] = {};
+        llama_model_meta_val_str(mr.model, "general.architecture", arch_name, sizeof(arch_name));
+        const llm_arch arch = llama_arch_from_string(arch_name);
+        return arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_DEEPSEEK41 ? "" : "requires DeepSeek V4 or V4.1";
     }
 
     case_result run(model_run & mr) const override {

tests/llama-context/cases-rs.cpp

@@ -7,6 +7,9 @@
 
 #include "common.h"
 
+#include "../../src/llama-ext.h" // llama_arch_from_string
+
+#include <algorithm>
 #include <cmath>
 #include <string>
 #include <vector>
@@ -278,18 +281,10 @@ struct case_rollback : ctx_test {
                     "replay logits mismatch at position " + std::to_string(pos) + ", nmse " + std::to_string(diff.nmse()));
         }
 
-        // TODO: this test is invalid because RS rollback is only correct once after a ubatch with more than n_rs_seq tokens
-        //       this is not the case here. add asserts and guardrails to prevent such attempts
-        //if (!llama_memory_seq_rm(llama_get_memory(ctx_src.get()), 0, rollback_pos, -1) ||
-        //    !llama_memory_seq_rm(llama_get_memory(ctx_dst.get()), 0, rollback_pos, -1)) {
-        //    TLOG_ERR("%s: partial rollback failed\n", __func__);
-        //    return case_result::fail();
-        //}
-
-        //constexpr llama_state_seq_flags partial_flags = LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY;
-        //common_prompt_checkpoint ckpt_partial;
-        //ckpt_partial.update_tgt(ctx_src.get(), 0, partial_flags);
-        //ckpt_partial.load_tgt(ctx_dst.get(), 0, partial_flags);
+        // the replay ran in single-token ubatches, which hold no snapshot that deep
+        CTX_REQUIRE(!llama_memory_seq_rm(llama_get_memory(ctx_src.get()), 0, rollback_pos, -1) &&
+                    !llama_memory_seq_rm(llama_get_memory(ctx_dst.get()), 0, rollback_pos, -1),
+                "a rollback past the snapshots of single-token ubatches was accepted");
 
         // Repeat the load into a context that already has its own rollback state:
         // groups 1..n_rs_seq hold a different prompt's history, and rs_idx[0] is
@@ -341,6 +336,80 @@ struct case_rollback : ctx_test {
     }
 };
 
+// after a ubatch of n tokens, probe rollbacks from deepest to shallowest: every rollback past the
+// snapshots that ubatch guarantees is refused and the deepest guaranteed one is accepted
+struct case_rollback_bound : ctx_test {
+    test_group group() const override { return TEST_GROUP_RS; }
+
+    const char * name() const override { return "rollback_bound"; }
+
+    const char * description() const override {
+        return "seq_rm accepts exactly the rollbacks whose snapshots the last ubatch wrote";
+    }
+
+    case_result run(model_run & mr) const override {
+        const int n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(mr.model));
+
+        constexpr llama_pos n_prefix = 16;
+
+        char arch_name[64] = {};
+        llama_model_meta_val_str(mr.model, "general.architecture", arch_name, sizeof(arch_name));
+        const llm_arch arch = llama_arch_from_string(arch_name);
+
+        // the DSV4 cache rebuilds snapshot d from the state before the ubatch, so d <= n is exact;
+        // the recurrent memory gets snapshots 0..n-1 from ggml_gated_delta_net, only for n <= n_rs_seq + 1
+        const bool from_ubatch_start = arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_DEEPSEEK41;
+
+        const auto make_ctx = [&]() {
+            ctx_opts opts;
+            opts.n_batch   = 64;
+            opts.n_ubatch  = 64;
+            opts.n_seq_max = 1;
+            opts.n_rs_seq  = 8;
+            return mr.ctx(opts);
+        };
+
+        llama_context_ptr ctx = make_ctx();
+        if (!ctx) {
+            return case_result::fail();
+        }
+
+        const uint32_t n_rs_seq = llama_n_rs_seq(ctx.get());
+        if (n_rs_seq == 0) {
+            return case_result::skip("n_rs_seq is disabled");
+        }
+
+        for (const uint32_t n : { 1u, 2u, n_rs_seq + 1, n_rs_seq + 2 }) {
+            const uint32_t r_max = from_ubatch_start ? std::min(n, n_rs_seq) : (n <= n_rs_seq + 1 ? n - 1 : 0);
+
+            ctx = make_ctx();
+            if (!ctx) {
+                return case_result::fail();
+            }
+
+            const llama_pos pos_end = n_prefix + (llama_pos) n;
+
+            CTX_REQUIRE(decode_gen(ctx.get(), n_vocab, 0, 0, n_prefix) &&
+                        decode_gen(ctx.get(), n_vocab, 0, n_prefix, pos_end),
+                    "the prefix / ubatch decode failed");
+
+            for (uint32_t r = n_rs_seq + 1; r > r_max; --r) {
+                CTX_REQUIRE(!llama_memory_seq_rm(llama_get_memory(ctx.get()), 0, pos_end - (llama_pos) r, -1),
+                        "rollback " + std::to_string(r) + " after a " + std::to_string(n) + "-token ubatch was accepted");
+            }
+
+            if (r_max > 0) {
+                CTX_REQUIRE(llama_memory_seq_rm(llama_get_memory(ctx.get()), 0, pos_end - (llama_pos) r_max, -1),
+                        "rollback " + std::to_string(r_max) + " after a " + std::to_string(n) + "-token ubatch was refused");
+            }
+        }
+
+        TLOG_INF("%s: rollback bounds hold for ubatches of 1, 2, %u and %u tokens\n", __func__, n_rs_seq + 1, n_rs_seq + 2);
+
+        return case_result::pass();
+    }
+};
+
 // decode a prompt into seq 0, share its cells with a second sequence, then keep decoding both
 struct case_shared_seq_reserve : ctx_test {
     test_group group() const override { return TEST_GROUP_RS; }
@@ -422,6 +491,7 @@ std::vector<ctx_test_ptr> make_rs_cases() {
     std::vector<ctx_test_ptr> cases;
 
     cases.push_back(std::make_unique<case_rollback>());
+    cases.push_back(std::make_unique<case_rollback_bound>());
     cases.push_back(std::make_unique<case_multi_seq_split_replay>());
     cases.push_back(std::make_unique<case_shared_seq_reserve>());
 

tests/llama-context/test-util.cpp

@@ -2,7 +2,7 @@
 
 #include "test-log.h"
 
-#include "../../src/llama-ext.h" // llama_memory_get_buffers
+#include "../../src/llama-ext.h" // llama_memory_get_buffers, llama_arch_from_string
 
 #include "common.h"
 
@@ -368,10 +368,11 @@ llama_sampler_ptr make_dist_sampler(uint32_t seed) {
 }
 
 bool arch_reserves_single_seq(const llama_model * model) {
-    char arch_str[64] = {};
-    llama_model_meta_val_str(model, "general.architecture", arch_str, sizeof(arch_str));
+    char arch_name[64] = {};
+    llama_model_meta_val_str(model, "general.architecture", arch_name, sizeof(arch_name));
+    const llm_arch arch = llama_arch_from_string(arch_name);
 
-    return strcmp(arch_str, "kimi-linear") == 0 || strcmp(arch_str, "minimax-01") == 0;
+    return arch == LLM_ARCH_KIMI_LINEAR || arch == LLM_ARCH_MINIMAX_01;
 }
 
 //