Commit ab3768cdc for llama.cpp
commit ab3768cdc24644bdc285bc07b575ca367b1e1dd9
Author: Pascal <admin@serveurperso.com>
Date: Fri Oct 30 09:11:58 2026 +0100
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.
diff --git a/src/llama-kv-cache-dsv4.cpp b/src/llama-kv-cache-dsv4.cpp
index 35eb57f76..027f3c61a 100644
--- a/src/llama-kv-cache-dsv4.cpp
+++ b/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;
diff --git a/src/llama-kv-cache-dsv4.h b/src/llama-kv-cache-dsv4.h
index 44bf95955..5a9dcfc55 100644
--- a/src/llama-kv-cache-dsv4.h
+++ b/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;
diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp
index 5fd644645..9213ac1dc 100644
--- a/src/llama-memory-recurrent.cpp
+++ b/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];
diff --git a/src/llama-memory-recurrent.h b/src/llama-memory-recurrent.h
index d0fc14779..6442fc1c1 100644
--- a/src/llama-memory-recurrent.h
+++ b/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 {
diff --git a/tests/llama-context/cases-kv.cpp b/tests/llama-context/cases-kv.cpp
index d5d5e1f09..b970ea940 100644
--- a/tests/llama-context/cases-kv.cpp
+++ b/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 {
diff --git a/tests/llama-context/cases-rs.cpp b/tests/llama-context/cases-rs.cpp
index 33347cef0..a26322b30 100644
--- a/tests/llama-context/cases-rs.cpp
+++ b/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>());
diff --git a/tests/llama-context/test-util.cpp b/tests/llama-context/test-util.cpp
index d66f90334..b67e2a03b 100644
--- a/tests/llama-context/test-util.cpp
+++ b/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;
}
//