From 93f17cd4ccf21100e9667ad62c39d16a751a4c7f Mon Sep 17 00:00:00 2001 From: rcfa <1831995+rcfa@users.noreply.github.com> Date: Wed, 30 Sep 2026 05:58:55 +0000 Subject: [PATCH] qsa : track a sequence's prefix while other sequences share the unified cache The incremental QSA state (qsa_prefix: one sequence's cells in position order, the pooled indexer-key cache, the compact metadata) could only be rebuilt when the tracked sequence was the only occupant of the cache: qsa_recover() gave up at the first cell of another sequence. A multi-slot server with a unified KV pool therefore never regained the incremental path once a second conversation had cells in the pool. - qsa_rebuild_prefix() (src/prefix.h, unit-tested in test-qsa-recover) skips other sequences' cells; holes, shared cells and duplicate positions are still refused. - qsa_apply() recovers the prefix of whichever single-sequence ubatch comes next, including right on a switch between sequences, and tries a sequence that cannot form a prefix only once until the state is invalidated again. Co-Authored-By: Claude Opus 5.5 Assisted-by: Claude Opus 5.5 --- src/llama-memory-hybrid-idx.cpp | 26 ++++++------ src/llama-memory-hybrid-idx.h | 1 + src/prefix.h | 20 ++++++++++ tests/CMakeLists.txt | 1 + tests/test-qsa-recover.cpp | 70 +++++++++++++++++++++++++++++++++ 5 files changed, 104 insertions(+), 14 deletions(-) create mode 100644 tests/test-qsa-recover.cpp diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index a976cfa1584f..32a2dd1af7f4 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -1053,19 +1053,7 @@ bool llama_memory_hybrid_idx::qsa_recover(llama_seq_id seq) { const auto & raw = mem_idx->get_cells(seq); if (raw.size() != qsa_prefix.positions.size()) { return false; } qsa_prefix_state rebuilt(qsa_prefix.positions.size()); - rebuilt.cells.resize(raw.get_used(), -1); - for (uint32_t cell=raw.used_min(); cell= rebuilt.cells.size() || rebuilt.cells[pos] >= 0 || - !raw.seq_has(cell, seq) || raw.seq_get_all(cell).count() != 1 || - !((ext.x == 0 && ext.y == 0) || (ext.x == pos && ext.y == pos))) { return false; } - rebuilt.cells[pos] = cell; rebuilt.positions[cell] = pos; - } - if (std::find(rebuilt.cells.begin(), rebuilt.cells.end(), -1) != rebuilt.cells.end()) { return false; } - rebuilt.sequence = seq; - for (size_t b=0; b= 0 && qsa_prefix.sequence != seq) { + qsa_invalidate(); // another sequence's turn: track this one from its first ubatch on + } + if (!qsa_prefix.valid) { + if (!qsa_recover_pending && seq == qsa_recover_failed) { return; } + qsa_recover_pending = false; + if (!qsa_recover(seq)) { qsa_recover_failed = seq; return; } + } qsa_recover_pending = false; + qsa_recover_failed = -1; if (!qsa_prefix.apply(seq, u.pos[0], slots.idxs[0])) { reject(); } } diff --git a/src/llama-memory-hybrid-idx.h b/src/llama-memory-hybrid-idx.h index f26209435dca..50d344e6e207 100644 --- a/src/llama-memory-hybrid-idx.h +++ b/src/llama-memory-hybrid-idx.h @@ -122,6 +122,7 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid { // the QSA layers run where a kernel reads the maskless block selection (HIP on RDNA3.5, Vulkan); otherwise the masked top-k bool selected_key_attn = false; bool qsa_recover_pending = false; + llama_seq_id qsa_recover_failed = -1; // the sequence the last recovery attempt failed for bool qsa_recover(llama_seq_id seq); qsa_prefix_state qsa_prefix; mutable std::vector qsa_ready; diff --git a/src/prefix.h b/src/prefix.h index cac9926197f3..ef7d2173d702 100644 --- a/src/prefix.h +++ b/src/prefix.h @@ -2,6 +2,7 @@ #include "llama-batch.h" #include "llama-kv-cells.h" +#include "qsa-prefix-state.h" #include #include #include @@ -54,6 +55,25 @@ static bool qsa_single_sequence_prefix(const llama_kv_cells & cells, uint32_t co return std::adjacent_find(positions.begin(),positions.end())==positions.end(); } +// rebuild the QSA prefix of seq from the cells (out: fresh, sized to the cache): each position in [0, n) must be in one cell that no other sequence shares +// a unified cache also holds other sequences (other slots of a server): skip their cells, seq does not attend to them +static bool qsa_rebuild_prefix(const llama_kv_cells & raw, llama_seq_id seq, qsa_prefix_state & out) { + out.cells.assign(raw.seq_pos_max(seq) + 1, -1); + for (uint32_t cell=raw.used_min(); cell= out.cells.size() || out.cells[pos] >= 0 || + raw.seq_get_all(cell).count() != 1 || + !((ext.x == 0 && ext.y == 0) || (ext.x == pos && ext.y == pos))) { return false; } + out.cells[pos] = cell; out.positions[cell] = pos; + } + if (std::find(out.cells.begin(), out.cells.end(), -1) != out.cells.end()) { return false; } + out.sequence = seq; + for (size_t b=0; b qsa_prefix_limits(const llama_pos * pos, int64_t tokens, int64_t strip, int64_t ratio, int64_t blocks, int64_t budget) { if (!pos || tokens<=0 || strip<=0 || ratio<=0 || budget<=0 || blocks + +static llama_kv_cells make_cells(uint32_t n) { llama_kv_cells c; c.resize(n); return c; } +static void put(llama_kv_cells & c, uint32_t cell, llama_seq_id seq, llama_pos pos) { c.pos_set(cell, pos); c.seq_add(cell, seq); } + +int main() { + // one sequence alone, in position order: what already worked + { + auto c = make_cells(64); + for (int p = 0; p < 10; ++p) { put(c, p, 0, p); } + qsa_prefix_state s(64); + GGML_ASSERT(qsa_rebuild_prefix(c, 0, s)); + GGML_ASSERT(s.sequence == 0 && s.cells.size() == 10 && s.block_positions.size() == 2); + for (int p = 0; p < 10; ++p) { GGML_ASSERT(s.cells[p] == p && s.positions[p] == p); } + } + // another conversation resident first (cells 0..14), ours after it, stored in reverse cell order + { + auto c = make_cells(64); + for (int p = 0; p < 15; ++p) { put(c, p, 1, p); } + for (int p = 0; p < 10; ++p) { put(c, 29 - p, 0, p); } + qsa_prefix_state s(64); + GGML_ASSERT(qsa_rebuild_prefix(c, 0, s)); + GGML_ASSERT(s.sequence == 0 && s.cells.size() == 10 && s.block_positions.size() == 2); + for (int p = 0; p < 10; ++p) { GGML_ASSERT(s.cells[p] == 29 - p && s.positions[29 - p] == p); } + for (int i = 0; i < 15; ++i) { GGML_ASSERT(s.positions[i] == -1); } + qsa_prefix_state t(64); + GGML_ASSERT(qsa_rebuild_prefix(c, 1, t) && t.cells.size() == 15 && t.positions[29] == -1); + } + // two conversations decoded together: their cells alternate + { + auto c = make_cells(64); + for (int p = 0; p < 20; ++p) { put(c, 2*p, 0, p); put(c, 2*p + 1, 1, p); } + qsa_prefix_state s(64); + GGML_ASSERT(qsa_rebuild_prefix(c, 1, s) && s.cells.size() == 20); + for (int p = 0; p < 20; ++p) { GGML_ASSERT(s.cells[p] == 2*p + 1); } + } + // a new conversation, nothing cached yet, others resident: an empty prefix its first ubatch extends + { + auto c = make_cells(64); + for (int p = 0; p < 15; ++p) { put(c, p, 1, p); } + qsa_prefix_state s(64); + GGML_ASSERT(qsa_rebuild_prefix(c, 0, s) && s.cells.empty() && s.sequence == 0); + GGML_ASSERT(s.apply(0, 0, {20, 21, 22})); + } + // still refused: a hole in the positions, a cell shared with another sequence, a duplicated position + { + auto c = make_cells(64); + for (int p = 0; p < 10; ++p) { if (p != 5) { put(c, p, 0, p); } } + qsa_prefix_state s(64); + GGML_ASSERT(!qsa_rebuild_prefix(c, 0, s)); + } + { + auto c = make_cells(64); + for (int p = 0; p < 10; ++p) { put(c, p, 0, p); } + c.seq_add(3, 1); + qsa_prefix_state s(64); + GGML_ASSERT(!qsa_rebuild_prefix(c, 0, s)); + } + { + auto c = make_cells(64); + for (int p = 0; p < 10; ++p) { put(c, p, 0, p); } + put(c, 40, 0, 3); + qsa_prefix_state s(64); + GGML_ASSERT(!qsa_rebuild_prefix(c, 0, s)); + } + std::cout << "PASS: prefixes rebuilt next to other sequences; holes, shared cells and duplicates refused\n"; +}