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 & llama_kv_cache_dsv4::get_rs_idx() const { return rs_idx; } -void llama_kv_cache_dsv4::reset_rs_idx_for_ubatches(const std::vector & ubatches) { +void llama_kv_cache_dsv4::update_rs_for_ubatches(const std::vector & ubatches) { if (n_rs_seq == 0) { return; } + std::vector 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 & get_rs_idx() const; - void reset_rs_idx_for_ubatches(const std::vector & ubatches); + void update_rs_for_ubatches(const std::vector & ubatches); private: llama_hparams hparams_raw; @@ -168,6 +168,10 @@ private: std::vector 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 rs_max; + // [n_layer] storage layer -> layer whose rows it holds (itself, or the source of a mirror) std::vector 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 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 #include -#include #include #include #include @@ -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 #include #include #include @@ -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 make_rs_cases() { std::vector cases; cases.push_back(std::make_unique()); + cases.push_back(std::make_unique()); cases.push_back(std::make_unique()); cases.push_back(std::make_unique()); 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; } //