Commit 43fe9c642 for llama.cpp
commit 43fe9c64281ef735046adc025e9e7559a1f659a5
Author: Pascal <admin@serveurperso.com>
Date: Tue Oct 6 07:41:57 2026 +0200
llama: fix k-pool scatter data race on shared sequences (#29994)
* llama: re-pool each shared k-pool rep once
With shared cells every pool is re-pooled, and since the pooled keys
are always scattered, the pools a seq_cp shares between sequences
wrote the same rep row from several scatter entries, a data race on
the CPU backend. Mark each rep once: the sharing sequences read the
same row through pool_cells.
* llama: assert whole-sequence seq_cp in the hybrid idx memory
The recurrent state is always copied whole whatever the range, and a
k-pool cell shared by a partial copy could carry two pool groupings
with a single pooled row. Every caller copies whole sequences, so
reject partial ranges instead of supporting them.
* llama: drop the k-pool cache_safe mode
With whole-sequence seq_cp, sequences sharing cells share their pools
too, so the pooled row of a shared rep is valid for all of them. Mark
each rep once in every ubatch instead of re-pooling everything while
cells are shared, which removes the sharing scan and the stale-all
workarounds in seq_rm, state_read and state_drop. seq_cp now only
stales the destination.
diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp
index 65ce4f6fb..13f8f7c37 100644
--- a/src/llama-memory-hybrid-idx.cpp
+++ b/src/llama-memory-hybrid-idx.cpp
@@ -192,24 +192,21 @@ bool llama_memory_hybrid_idx::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_po
const llama_pos stale = mem_idx_stale_pos(seq_id, p0);
mem_idx->seq_rm(seq_id, p0, p1);
mem_idx_stale_set(seq_id, stale);
-
- // removing a sequence can free cells another sequence shared, but only this one is marked stale, so the
- // survivor would keep shared = true and pin cache_safe off forever; stale every sequence to re-derive it
- if (kpool_layout_shared()) {
- mem_idx_stale_set(-1, 0);
- }
}
return get_mem_attn()->seq_rm(seq_id, p0, p1);
}
void llama_memory_hybrid_idx::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {
+ // only whole sequences are copied: the recurrent state ignores the range, and a shared cell holds a single pool grouping
+ GGML_ASSERT(p0 <= 0 && p1 < 0 && "partial seq_cp is not supported");
+
llama_memory_hybrid::seq_cp(seq_id_src, seq_id_dst, p0, p1);
if (mem_idx) {
mem_idx->seq_cp(seq_id_src, seq_id_dst, p0, p1);
- // the copy shares cells, which cannot hold two groupings, so both sides drop their cached keys
- mem_idx_stale_set(seq_id_src, 0);
+ // a whole sequence copy gives the destination the source's pools, rep rows included: the source keeps its
+ // pooled keys, the destination rebuilds its layout and re-pools into the same rows
mem_idx_stale_set(seq_id_dst, 0);
}
}
@@ -219,7 +216,7 @@ void llama_memory_hybrid_idx::seq_keep(llama_seq_id seq_id) {
if (mem_idx) {
mem_idx->seq_keep(seq_id);
- // cells shared with the dropped sequences become exclusive again, their keys were never cached
+ // every other sequence loses its cells, so their layouts must rebuild
mem_idx_stale_set(-1, 0);
}
}
@@ -291,10 +288,6 @@ void llama_memory_hybrid_idx::state_read(llama_io_read_i & io, llama_seq_id seq_
mem_idx->state_read_sinfo(io, seq_id, flags, nullptr, &sinfos_attn);
// the restore rewrites the cells behind the pool layout's back
mem_idx_stale_set(seq_id, 0);
- // it can also change which cells are shared; re-derive sharing for every sequence, as seq_rm does
- if (kpool_layout_shared()) {
- mem_idx_stale_set(-1, 0);
- }
}
}
@@ -321,10 +314,6 @@ void llama_memory_hybrid_idx::state_drop(llama_seq_id seq_id) {
if (mem_idx) {
mem_idx->state_clear(seq_id);
mem_idx_stale_set(seq_id, 0);
- // clearing this sequence can end a sharing the survivor would otherwise keep flagged (see seq_rm)
- if (kpool_layout_shared()) {
- mem_idx_stale_set(-1, 0);
- }
}
}
@@ -358,26 +347,22 @@ struct llama_memory_hybrid_idx::kpool_layout {
// Where the pool scan stopped, so an append resumes instead of starting over.
size_t j_next = 0;
-
- // Whether any cell also carries another sequence, which rules out caching this sequence's pooled keys.
- bool shared = false;
};
std::array<seq, LLAMA_MAX_SEQ> seqs;
uint32_t n_pool_real = 0;
- bool cache_safe = true;
};
// Which pools of the layout the current ubatch must re-pool, in the layout's pool order.
struct llama_memory_hybrid_idx_context::kpool_state {
std::vector<uint32_t> is_new;
+ std::vector<uint32_t> rep_gen; // per global cell, the generation that last marked a pool with that rep
uint32_t generation = 0;
uint32_t n_pool_real = 0;
uint32_t n_new = 0;
uint32_t n_new_g = 1; // graph size of the new pool list, stable across decode steps
- bool cache_safe = true;
};
namespace {
@@ -404,10 +389,6 @@ const llama_memory_hybrid_idx::kpool_layout & llama_memory_hybrid_idx::kpool_lay
return *kpool_lay;
}
-bool llama_memory_hybrid_idx::kpool_layout_shared() const {
- return kpool_lay && !kpool_lay->cache_safe;
-}
-
// Pools are fixed by the positions relative to the sequence's first one, so the layout survives a plain
// append. A sequence edit can regroup them, and mem_idx_stale tells us it happened.
const llama_memory_hybrid_idx::kpool_layout & llama_memory_hybrid_idx::kpool_layout_update() {
@@ -424,7 +405,6 @@ const llama_memory_hybrid_idx::kpool_layout & llama_memory_hybrid_idx::kpool_lay
const bool unified = n_stream_kv == 1;
lay.n_pool_real = 0;
- lay.cache_safe = true;
for (llama_seq_id s = 0; s < LLAMA_MAX_SEQ; ++s) {
auto & sq = lay.seqs[s];
@@ -440,10 +420,8 @@ const llama_memory_hybrid_idx::kpool_layout & llama_memory_hybrid_idx::kpool_lay
sq.strm = unified ? 0 : mem_idx->get_stream(s);
- size_t n_kept = 0;
if (mem_idx_stale[s] == POS_CLEAN && !sq.cells.empty() && !sp.empty() &&
sq.pos_min == sp.begin()->first) {
- n_kept = sq.cells.size();
for (auto it = sp.upper_bound(sq.cells.back()); it != sp.end(); ++it) {
sq.cells.push_back(*it);
}
@@ -455,21 +433,7 @@ const llama_memory_hybrid_idx::kpool_layout & llama_memory_hybrid_idx::kpool_lay
sq.cells.assign(sp.begin(), sp.end());
sq.pools.clear();
sq.j_next = 0;
- sq.shared = false;
sq.pos_min = sp.empty() ? 0 : sp.begin()->first;
- n_kept = 0;
- }
-
- // sharing starts with a seq_cp; it ends with an edit, or a seq_rm/state_drop/state_read that frees the
- // shared cells - each stales every sequence so the rebuild above re-derives it, so once set it holds
- // until then and the rescan can be skipped
- if (unified && !sq.shared) {
- for (size_t j = n_kept; j < sq.cells.size(); ++j) {
- if (cells.seq_count(sq.cells[j].second) > 1) {
- sq.shared = true;
- break;
- }
- }
}
// Pools start at the first valid token
@@ -504,7 +468,6 @@ const llama_memory_hybrid_idx::kpool_layout & llama_memory_hybrid_idx::kpool_lay
sq.j_next = j;
lay.n_pool_real += (uint32_t) sq.pools.size();
- lay.cache_safe = lay.cache_safe && !sq.shared;
}
return lay;
@@ -681,34 +644,40 @@ void llama_memory_hybrid_idx_context::kpool_build_state(const llama_ubatch & uba
const auto & lay = mem->kpool_layout_get();
auto & st = *kpool_st;
+ const auto * idx = mem->get_mem_idx();
+ const uint32_t kv_size = idx->get_size();
+ const uint32_t kpool = mem->get_kpool();
+
st.n_pool_real = lay.n_pool_real;
- st.cache_safe = lay.cache_safe;
st.n_new = 0;
if (++st.generation == 0) {
- std::fill(st.is_new.begin(), st.is_new.end(), 0);
+ std::fill(st.is_new.begin(), st.is_new.end(), 0);
+ std::fill(st.rep_gen.begin(), st.rep_gen.end(), 0);
st.generation = 1;
}
st.is_new.resize(lay.n_pool_real, 0);
+ st.rep_gen.resize((size_t) kv_size*idx->get_n_stream(), 0);
- auto mark = [&](uint32_t ip) {
- if (st.is_new[ip] != st.generation) {
- st.is_new[ip] = st.generation;
+ std::array<uint32_t, LLAMA_MAX_SEQ> pool_start;
+
+ // a pool is marked once per rep: sequences sharing cells (a seq_cp, or tokens decoded for several sequences)
+ // share their pools, whose single pooled row they all read through pool_cells, so the scatter rows stay unique
+ auto mark = [&](llama_seq_id s, size_t k) {
+ const auto & sq = lay.seqs[s];
+ const size_t rep = (size_t) sq.strm*kv_size + sq.cells[sq.pools[k] + kpool - 1].second;
+ if (st.rep_gen[rep] != st.generation) {
+ st.rep_gen[rep] = st.generation;
+ st.is_new[pool_start[s] + k] = st.generation;
++st.n_new;
}
};
- const uint32_t kpool = mem->get_kpool();
- std::array<uint32_t, LLAMA_MAX_SEQ> pool_start;
uint32_t ip = 0;
for (llama_seq_id s = 0; s < LLAMA_MAX_SEQ; ++s) {
const auto & sq = lay.seqs[s];
pool_start[s] = ip;
ip += (uint32_t) sq.pools.size();
- if (!st.cache_safe) {
- continue;
- }
-
// A sequence edit invalidates only pools ending after the edited position.
const llama_pos stale_from = i_cur == 0 ?
mem_idx_stale_batch[s] : llama_memory_hybrid_idx::POS_CLEAN;
@@ -719,18 +688,11 @@ void llama_memory_hybrid_idx_context::kpool_build_state(const llama_ubatch & uba
auto first = std::lower_bound(sq.pools.begin(), sq.pools.end(), stale_from,
[&](uint32_t j, llama_pos p) { return sq.cells[j + kpool - 1].first < p; });
for (auto it = first; it != sq.pools.end(); ++it) {
- mark(pool_start[s] + (uint32_t) (it - sq.pools.begin()));
+ mark(s, it - sq.pools.begin());
}
}
GGML_ASSERT(ip == st.is_new.size());
- if (!st.cache_safe) {
- std::fill(st.is_new.begin(), st.is_new.end(), st.generation);
- st.n_new = st.n_pool_real;
- st.n_new_g = std::max(st.n_new, 1u);
- return;
- }
-
// in order mode a token's cell gives its rank, and the rank its pool: positions cannot, as an image shares one
const bool by_order = mem->get_kpool_by_order();
const auto * sinfo = by_order ? &sinfos_kpool[i_cur] : nullptr;
@@ -745,7 +707,7 @@ void llama_memory_hybrid_idx_context::kpool_build_state(const llama_ubatch & uba
const int64_t r = kpool_rank(sq.cells, p, sinfo->idxs[i / n_tps][i % n_tps]);
GGML_ASSERT(r >= 0);
if ((size_t) r / kpool < sq.pools.size()) {
- mark(pool_start[s] + (uint32_t) (r / kpool));
+ mark(s, (size_t) r / kpool);
}
continue;
}
@@ -756,15 +718,14 @@ void llama_memory_hybrid_idx_context::kpool_build_state(const llama_ubatch & uba
}
--it;
if (p <= sq.cells[*it + kpool - 1].first) {
- mark(pool_start[s] + (uint32_t) (it - sq.pools.begin()));
+ mark(s, it - sq.pools.begin());
}
}
}
// a ubatch touches at most t_s/kpool + 1 pools per sequence, pad to that bound so the graph keeps its shape
// as the count moves; reserve sizes the list for every pool the cache can hold, so never pad past n_pool_max
- const auto * idx = mem->get_mem_idx();
- const uint32_t n_pool_max = idx->get_size() / kpool * idx->get_n_seq_max();
+ const uint32_t n_pool_max = kv_size / kpool * idx->get_n_seq_max();
const uint32_t bound = ubatch.n_tokens/kpool + ubatch.n_seqs_unq;
st.n_new_g = std::max({st.n_new, 1u, std::min({bound, kpool_pad(st.n_pool_real) - 1, n_pool_max})});
}
diff --git a/src/llama-memory-hybrid-idx.h b/src/llama-memory-hybrid-idx.h
index 18e3413b1..cbfb67943 100644
--- a/src/llama-memory-hybrid-idx.h
+++ b/src/llama-memory-hybrid-idx.h
@@ -121,9 +121,6 @@ private:
// unique_ptr because kpool_layout is incomplete here
std::unique_ptr<kpool_layout> kpool_lay;
- // whether the current layout has cells shared between sequences (kpool_layout is incomplete here, so out of line)
- bool kpool_layout_shared() const;
-
// seq_id < 0 stales every sequence, p0 < 0 stales the sequence from its first position
void mem_idx_stale_set(llama_seq_id seq_id, llama_pos p0);
diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp
index 839f223d2..48f4c1af8 100644
--- a/src/models/glm5-next.cpp
+++ b/src/models/glm5-next.cpp
@@ -341,8 +341,7 @@ llama_model_glm5_next::llm_graph_input_kpool * llama_model_glm5_next::graph::bui
inp->new_pool_idxs = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, kpool, n_new);
ggml_set_input(inp->new_pool_idxs);
- // the scatter target is part of the graph shape: llama_context reserves the full-context graph,
- // so this must not depend on cache_safe, which only the decode-time graph can know
+ // one scatter row per new pool, each a distinct rep row (see kpool_build_state)
inp->new_pool_rep = ggml_new_tensor_1d(ctx0, GGML_TYPE_I64, n_new);
ggml_set_input(inp->new_pool_rep);
@@ -813,9 +812,8 @@ ggml_tensor * llama_model_glm5_next::graph::build_kpool_select(
pooled_new = ggml_reshape_2d(ctx0, pooled_new, n_embd_indexer, n_new);
cb(pooled_new, "indexer_pool_k_new", il);
- // scatter the fresh pooled keys, then gather all n_pool of them by cell, in both cache modes:
- // the reserved graph cannot branch on cache_safe, and without sharing every pool is re-pooled
- // anyway (n_new == n_pool_real, layout order), so the gather returns exactly pooled_new
+ // scatter the fresh pooled keys into their rep rows, then gather all n_pool of them by cell:
+ // the older pools come from the rows earlier ubatches wrote
ggml_build_forward_expand(gf, kpool_cache.scatter_pooled(pooled_new, inp_kpool->new_pool_rep));
ggml_tensor * pooled = kpool_cache.gather_pooled(inp_kpool->pool_cells);
pooled = ggml_reshape_3d(ctx0, pooled, n_embd_indexer, 1, n_pool);
diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp
index 250eb74da..71ce2d503 100644
--- a/src/models/qwen4exp.cpp
+++ b/src/models/qwen4exp.cpp
@@ -726,8 +726,7 @@ llama_model_qwen4exp::llm_graph_input_kpool * llama_model_qwen4exp::graph::build
inp->new_pool_idxs = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, kpool, inp->n_new);
ggml_set_input(inp->new_pool_idxs);
- // the scatter target is part of the graph shape: llama_context reserves the full-context graph,
- // so this must not depend on cache_safe, which only the decode-time graph can know
+ // one scatter row per new pool, each a distinct rep row (see kpool_build_state)
inp->new_pool_rep = ggml_new_tensor_1d(ctx0, GGML_TYPE_I64, inp->n_new);
ggml_set_input(inp->new_pool_rep);
inp->new_pool_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, 4*inp->n_new);
@@ -787,9 +786,8 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_sel(
pooled_new = ggml_reshape_2d(ctx0, pooled_new, idx_dim, n_new);
cb(pooled_new, "indexer_pool_k_new", il);
- // scatter the fresh pooled keys, then gather all n_pool of them by cell, in both cache modes:
- // the reserved graph cannot branch on cache_safe, and without sharing every pool is re-pooled
- // anyway (n_new == n_pool_real, layout order), so the gather returns exactly pooled_new
+ // scatter the fresh pooled keys into their rep rows, then gather all n_pool of them by cell:
+ // the older pools come from the rows earlier ubatches wrote
ggml_build_forward_expand(gf, kpool_cache.scatter_pooled(pooled_new, inp_kpool->new_pool_rep));
ggml_tensor * pooled = kpool_cache.gather_pooled(inp_kpool->pool_cells);
pooled = ggml_reshape_3d(ctx0, pooled, idx_dim, 1, n_pool);