qsa : track a sequence's prefix while other sequences share the unified cache

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<raw.used_max_p1(); ++cell) {
-        if (raw.is_empty(cell)) { continue; }
-        const int32_t pos = raw.pos_get(cell);
-        const auto & ext = raw.ext_get(cell);
-        if (pos < 0 || size_t(pos) >= 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<rebuilt.cells.size()/4; ++b) { rebuilt.block_positions.push_back(b*4); }
+    if (!qsa_rebuild_prefix(raw, seq, rebuilt)) { return false; }
     qsa_prefix = std::move(rebuilt);
     return true;
 }
@@ -1086,8 +1074,18 @@ void llama_memory_hybrid_idx::qsa_apply(const llama_ubatch & u, const llama_kv_c
             if (u.pos[i+axis*u.n_tokens] != u.pos[i]) { reject(); return; }
         }
     }
-    if (!qsa_prefix.valid && (!qsa_recover_pending || !qsa_recover(seq))) { qsa_recover_pending = false; return; }
+    // a unified cache serves its sequences in turn: when another sequence comes next, rebuild the prefix from its cells
+    // a sequence whose cells do not form a prefix is tried once, until the state changes again
+    if (qsa_prefix.valid && qsa_prefix.sequence >= 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(); }
 }
 

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<int64_t> qsa_ready;

src/prefix.h

@@ -2,6 +2,7 @@
 
 #include "llama-batch.h"
 #include "llama-kv-cells.h"
+#include "qsa-prefix-state.h"
 #include <algorithm>
 #include <cstdint>
 #include <unordered_map>
@@ -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<raw.used_max_p1(); ++cell) {
+        if (raw.is_empty(cell) || !raw.seq_has(cell, seq)) { continue; }
+        const int32_t pos = raw.pos_get(cell);
+        const auto & ext = raw.ext_get(cell);
+        if (pos < 0 || size_t(pos) >= 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<out.cells.size()/4; ++b) { out.block_positions.push_back(b*4); }
+    return true;
+}
+
 static std::vector<int64_t> 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<budget) { return {}; }

tests/CMakeLists.txt

@@ -150,6 +150,7 @@ endif ()
 
 llama_build_and_test(test-qsa-prefix.cpp)
 llama_build_and_test(test-qsa-fast.cpp)
+llama_build_and_test(test-qsa-recover.cpp)
 llama_build_and_test(test-ple-disk.cpp)
 llama_build(test-recurrent-state-rollback.cpp)
 llama_build(test-mtp-draft-vocab.cpp) # needs a qwen35/qwen35moe model with an MTP head: -m MODEL

tests/test-qsa-recover.cpp

@@ -0,0 +1,70 @@
+// qsa_rebuild_prefix: rebuild the QSA prefix of one sequence from a unified KV cache that also holds other sequences
+#include "../src/prefix.h"
+#include "ggml.h"
+#include <iostream>
+
+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";
+}