Commit 310991409 for llama.cpp
commit 3109914090564b4c5280f30896369d55b86bbdbf
Author: Pascal <admin@serveurperso.com>
Date: Tue Oct 6 18:31:13 2026 +0200
llama: remove the gather path of the glm5-next sparse attention (#30042)
The gather path attended over the selected latents with a plain
matmul and softmax. It only ran with n_ubatch <= 16, and the flash
attention backends now skip the masked rows through n_kv_max, so
the scatter path covers every case.
Drop the gather flag, the gathered attention branch and
gather_mla_rows. set_input_kpool always maps padding to the n_kv
sentinel, and the slot mask becomes sel_mask since only the
scatter reads it.
diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp
index 13f8f7c37..21ec0b2fe 100644
--- a/src/llama-memory-hybrid-idx.cpp
+++ b/src/llama-memory-hybrid-idx.cpp
@@ -609,16 +609,6 @@ llama_memory_hybrid_idx_context::kpool_access llama_memory_hybrid_idx_context::g
return kpool_access(ctx, mem->get_mem_idx()->get_k_storage(il), n_embd);
}
-ggml_tensor * llama_memory_hybrid_idx_context::gather_mla_rows(
- ggml_context * ctx, ggml_tensor * idxs, int64_t n_rows, int64_t n_embd, int32_t il) const {
- GGML_ASSERT(mem != nullptr);
- ggml_tensor * k = mem->get_mem_attn()->get_k_storage(il);
- GGML_ASSERT(k->ne[0] == n_embd);
-
- ggml_tensor * rows = ggml_view_2d(ctx, k, k->ne[0], k->ne[1]*k->ne[2], k->nb[1], 0);
- return ggml_get_rows(ctx, rows, ggml_reshape_1d(ctx, idxs, n_rows));
-}
-
// k-pool DSA indexer (glm5-next, qwen4exp QSA)
// Sizes only, used by the full cache context so get_n_kpool() works during graph reserve.
@@ -745,7 +735,7 @@ uint32_t llama_memory_hybrid_idx_context::get_n_kpool_new() const {
}
void llama_memory_hybrid_idx_context::set_input_kpool(ggml_tensor * pool_cells, ggml_tensor * pool_idxs, ggml_tensor * pool_mask, ggml_tensor * tail_idxs,
- ggml_tensor * gather_mask, bool gather, ggml_tensor * new_pool_idxs, ggml_tensor * new_pool_rep,
+ ggml_tensor * sel_mask, ggml_tensor * new_pool_idxs, ggml_tensor * new_pool_rep,
const llama_ubatch * ubatch, ggml_tensor * new_pool_pos) const {
GGML_ASSERT(mem != nullptr && mem->get_mem_idx() != nullptr);
GGML_ASSERT(ggml_backend_buffer_is_host(pool_cells->buffer));
@@ -797,7 +787,7 @@ void llama_memory_hybrid_idx_context::set_input_kpool(ggml_tensor * pool_cells,
}
}
- // Use the first ubatch cell for padded gathers.
+ // a cell of this ubatch, written before any read, so the padded pools read a finite K row
int64_t dummy_cell = 0;
{
const llama_seq_id s = ubatch->seq_id[0][0];
@@ -820,21 +810,21 @@ void llama_memory_hybrid_idx_context::set_input_kpool(ggml_tensor * pool_cells,
}
}
- // Gather maps padding to a real cell and masks it separately.
- const int32_t sentinel = gather ? (int32_t) dummy_cell : (int32_t) n_kv;
+ // padding and absent cells point at the n_kv sentinel row, one past the live cells
+ const int32_t sentinel = (int32_t) n_kv;
float * gm = nullptr;
uint32_t n_sel = 0;
uint32_t n_top = 0; // Pools per token in the selection.
- if (gather_mask != nullptr) {
- GGML_ASSERT(ggml_backend_buffer_is_host(gather_mask->buffer));
- GGML_ASSERT(gather_mask->type == GGML_TYPE_F32);
- GGML_ASSERT(gather_mask->ne[3] == (int64_t) n_tokens && gather_mask->ne[1] == 1 && gather_mask->ne[2] == 1);
- n_sel = (uint32_t) gather_mask->ne[0];
+ if (sel_mask != nullptr) {
+ GGML_ASSERT(ggml_backend_buffer_is_host(sel_mask->buffer));
+ GGML_ASSERT(sel_mask->type == GGML_TYPE_F32);
+ GGML_ASSERT(sel_mask->ne[3] == (int64_t) n_tokens && sel_mask->ne[1] == 1 && sel_mask->ne[2] == 1);
+ n_sel = (uint32_t) sel_mask->ne[0];
// The tail slots, when selected, are the n_sel % kpool != 0 remainder.
n_top = n_sel / kpool;
GGML_ASSERT(n_sel % kpool == 0 || n_sel % kpool == kpool - 1);
- gm = (float *) gather_mask->data;
+ gm = (float *) sel_mask->data;
}
// pools are laid out per sequence
@@ -857,7 +847,7 @@ void llama_memory_hybrid_idx_context::set_input_kpool(ggml_tensor * pool_cells,
const auto & sq = lay.seqs[s];
seq_pool_start[s] = (uint32_t) pool_end.size();
- const bool inert = !gather && n_stream_kv > 1 && !seq_in_ub[s];
+ const bool inert = n_stream_kv > 1 && !seq_in_ub[s];
for (size_t pi = 0; pi < sq.pools.size(); ++pi) {
const uint32_t j = sq.pools[pi];
@@ -869,8 +859,7 @@ void llama_memory_hybrid_idx_context::set_input_kpool(ggml_tensor * pool_cells,
pcell[ip] = (int32_t) gcell(sq, rep);
for (uint32_t k = 0; k < kpool; ++k) {
- pidx[(size_t) ip*kpool + k] = inert ? sentinel :
- (int32_t) (gather ? gcell(sq, sq.cells[j + k].second) : (int64_t) sq.cells[j + k].second);
+ pidx[(size_t) ip*kpool + k] = inert ? sentinel : (int32_t) sq.cells[j + k].second;
}
if (st.is_new[ip] == st.generation) {
@@ -972,13 +961,13 @@ void llama_memory_hybrid_idx_context::set_input_kpool(ggml_tensor * pool_cells,
bool real = false;
if (k < n_tail && by_order) {
const uint32_t c = sq.cells[rank[i] - k].second;
- cell = (int32_t) (gather ? gcell(sq, c) : (int64_t) c);
+ cell = (int32_t) c;
real = true;
} else if (k < n_tail) {
const llama_pos pt = p - (llama_pos) k;
auto it = std::lower_bound(sq.cells.begin(), sq.cells.end(), std::make_pair(pt, 0u));
if (it != sq.cells.end() && it->first == pt) {
- cell = (int32_t) (gather ? gcell(sq, it->second) : (int64_t) it->second);
+ cell = (int32_t) it->second;
real = true;
}
}
diff --git a/src/llama-memory-hybrid-idx.h b/src/llama-memory-hybrid-idx.h
index cbfb67943..0ac4265df 100644
--- a/src/llama-memory-hybrid-idx.h
+++ b/src/llama-memory-hybrid-idx.h
@@ -192,10 +192,10 @@ public:
uint32_t get_n_kpool () const; // Padded pool count, where the last pool is always unused.
uint32_t get_n_kpool_new() const; // Pools to re-pool this ubatch, padded to a stable bound, never below 1.
kpool_access get_kpool_access(ggml_context * ctx, int32_t il, int64_t n_embd) const;
- ggml_tensor * gather_mla_rows(ggml_context * ctx, ggml_tensor * idxs, int64_t n_rows, int64_t n_embd, int32_t il) const;
+ // sel_mask (F32 [n_sel, 1, 1, n_tokens], can be null): 0 for the live selection slots, -inf for the dead ones
// new_pool_pos (I32 [4*n_new]): M-RoPE position of each new pool's first member, for pooled keys rotated at pooling time
void set_input_kpool(ggml_tensor * pool_cells, ggml_tensor * pool_idxs, ggml_tensor * pool_mask, ggml_tensor * tail_idxs,
- ggml_tensor * gather_mask, bool gather, ggml_tensor * new_pool_idxs, ggml_tensor * new_pool_rep,
+ ggml_tensor * sel_mask, ggml_tensor * new_pool_idxs, ggml_tensor * new_pool_rep,
const llama_ubatch * ubatch, ggml_tensor * new_pool_pos = nullptr) const;
private:
diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp
index a062deb7a..955405c75 100644
--- a/src/models/glm5-next.cpp
+++ b/src/models/glm5-next.cpp
@@ -243,7 +243,7 @@ public:
void set_input(const llama_ubatch * ubatch) override {
mctx->get_idx()->set_input_k_idxs(k_idxs, ubatch);
- mctx->set_input_kpool(pool_cells, pool_idxs, pool_mask, tail_idxs, gather_mask, gather, new_pool_idxs, new_pool_rep, ubatch);
+ mctx->set_input_kpool(pool_cells, pool_idxs, pool_mask, tail_idxs, sel_mask, new_pool_idxs, new_pool_rep, ubatch);
}
bool can_reuse(const llm_graph_params & params) override {
@@ -272,7 +272,7 @@ public:
ggml_tensor * pool_idxs = nullptr; // I32 [kpool, n_pool] member cells per pool, n_kv sentinel for the padded pools
ggml_tensor * pool_mask = nullptr; // F32/F16 [n_pool, n_tokens]
ggml_tensor * tail_idxs = nullptr; // I32 [kpool - 1, n_tokens]
- ggml_tensor * gather_mask = nullptr; // F32 [n_sel, 1, 1, n_tokens] 0 for live selection slots, -inf for dead ones
+ ggml_tensor * sel_mask = nullptr; // F32 [n_sel, 1, 1, n_tokens] 0 for live selection slots, -inf for dead ones
// n_new is never below 1, see build_inp_kpool
ggml_tensor * new_pool_idxs = nullptr; // I32 [kpool, n_new] members of the pools completed this ubatch
ggml_tensor * new_pool_rep = nullptr; // I64 [n_new] cell to write each new pooled key into
@@ -281,7 +281,6 @@ public:
const uint32_t kpool;
uint32_t n_new = 0;
uint32_t n_sel = 0;
- bool gather = false;
uint32_t n_kv = 0;
};
@@ -318,23 +317,17 @@ llama_model_glm5_next::llm_graph_input_kpool * llama_model_glm5_next::graph::bui
inp->n_kv = n_kv;
- // Gather selected latents for small batches when the context exceeds the selection width.
+ // selection width: the top pools plus the optional tail, also the sparse attention bound
{
- constexpr int64_t max_ub = 16;
-
const int64_t n_top_pool = std::min<int64_t>(n_pool, hparams.indexer_top_k / kpool);
const int64_t n_sel = kpool*n_top_pool + (hparams.indexer_kpool_select_tail ? kpool - 1 : 0);
inp->n_sel = (uint32_t) n_sel;
- // both terms are context constants: n_ubatch bounds every ubatch and top_k + kpool - 1 bounds
- // n_sel, so the graph shape follows neither n_tokens nor n_kv, which the reserve cannot predict
- // TODO: remove "gather" logic and everything related. the backends now support sparse attension so this is obsolete
- inp->gather = (int64_t) cparams.n_ubatch <= max_ub && (int64_t) cparams.n_ctx > hparams.indexer_top_k + kpool - 1;
-
- // Both paths read the slot mask: gather adds it to the scores, scatter maps its dead slots to dump rows.
- inp->gather_mask = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_sel, 1, 1, n_tokens);
- ggml_set_input(inp->gather_mask);
- // Keep the mask allocated even when no op reads it, because set_input_kpool always fills it.
- ggml_build_forward_expand(gf, inp->gather_mask);
+
+ // the scatter maps the dead slots of this mask to dump rows
+ inp->sel_mask = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_sel, 1, 1, n_tokens);
+ ggml_set_input(inp->sel_mask);
+ // set_input_kpool always fills the mask, so it stays allocated in every graph
+ ggml_build_forward_expand(gf, inp->sel_mask);
}
inp->n_new = n_new;
@@ -845,7 +838,7 @@ ggml_tensor * llama_model_glm5_next::graph::build_kpool_select(
const int64_t n_top_pool = std::min<int64_t>(n_pool, hparams.indexer_top_k / kpool);
ggml_tensor * top_k = ggml_top_k(ctx0, score, n_top_pool); // [n_top_pool, n_tokens], UNORDERED
- // The gather mask marks the first min(nv, n_top_pool) slots as the visible pools, so order the set by descending score.
+ // The selection mask marks the first min(nv, n_top_pool) slots as the visible pools, so order the set by descending score.
ggml_tensor * sel_score = ggml_get_rows(ctx0,
ggml_reshape_3d(ctx0, score, 1, n_pool, n_tokens), top_k); // [1, n_top_pool, n_tokens]
ggml_tensor * sel_order = ggml_argsort(ctx0,
@@ -866,13 +859,6 @@ ggml_tensor * llama_model_glm5_next::graph::build_kpool_select(
}
const int64_t n_sel = sel_idx->ne[0];
- // Gather returns selected cell indices and masks padding separately.
- if (inp_kpool->gather) {
- GGML_ASSERT(inp_kpool->gather_mask->ne[0] == n_sel && inp_kpool->gather_mask->ne[3] == n_tokens);
- cb(sel_idx, "indexer_sel_idx", il);
- return sel_idx;
- }
-
ggml_build_forward_expand(gf, sel_idx);
ggml_tensor * mask_all = ggml_new_tensor_4d(ctx0, kq_mask->type, n_kv + n_sel, 1, 1, 1);
@@ -887,8 +873,8 @@ ggml_tensor * llama_model_glm5_next::graph::build_kpool_select(
// Live slots (visible pools, real tail cells) address disjoint cells. Each dead slot writes its own dump row
// n_kv + slot, so the scatter indices of a token are unique: idx = dump + live*(idx - dump), live = exp(mask).
- GGML_ASSERT(inp_kpool->gather_mask->ne[0] == n_sel && inp_kpool->gather_mask->ne[3] == n_tokens);
- ggml_tensor * live = ggml_exp(ctx0, ggml_reshape_2d(ctx0, inp_kpool->gather_mask, n_sel, n_tokens));
+ GGML_ASSERT(inp_kpool->sel_mask->ne[0] == n_sel && inp_kpool->sel_mask->ne[3] == n_tokens);
+ ggml_tensor * live = ggml_exp(ctx0, ggml_reshape_2d(ctx0, inp_kpool->sel_mask, n_sel, n_tokens));
ggml_tensor * dump = ggml_arange(ctx0, (float) n_kv, (float) (n_kv + n_sel), 1.0f);
ggml_tensor * idx_f = ggml_cast(ctx0, sel_idx, GGML_TYPE_F32);
idx_f = ggml_add(ctx0, ggml_mul(ctx0, ggml_sub(ctx0, idx_f, dump), live), dump);
@@ -954,43 +940,14 @@ ggml_tensor * llama_model_glm5_next::graph::build_dsa_layer(
ggml_build_forward_expand(gf, kv_cmpr);
ggml_build_forward_expand(gf, mctx_mla->cpy_k(ctx0, kv_cmpr, inp_attn->get_k_idxs(), il));
- ggml_tensor * out = nullptr;
- if (inp_kpool->gather) {
- // Attend over gathered latents with the token dimension in ne[3].
-
- ggml_build_forward_expand(gf, kq_mask);
-
- ggml_tensor * sel_idx = sel; // I32 [n_sel, n_tokens]
- const int64_t n_sel = sel_idx->ne[0];
+ // The scatter selection already includes the causal mask.
+ ggml_tensor * mask = ggml_reshape_4d(ctx0, sel, kq_mask->ne[0], kq_mask->ne[1], kq_mask->ne[2], kq_mask->ne[3]);
+ cb(mask, "kq_mask_dsa", il);
- ggml_tensor * k_g = mctx_hyb->gather_mla_rows(ctx0, sel_idx, n_sel*n_tokens, kv_lora_rank, il);
- k_g = ggml_reshape_4d(ctx0, k_g, kv_lora_rank, n_sel, 1, n_tokens); // F32 [kv_lora_rank, n_sel, 1, n_tokens]
- cb(k_g, "kv_gathered", il);
+ ggml_tensor * k = mctx_mla->get_k(ctx0, il);
+ ggml_tensor * v = ggml_view_4d(ctx0, k, kv_lora_rank, k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0);
- ggml_tensor * q_g = ggml_permute(ctx0, q_absorbed, 0, 2, 3, 1); // [kv_lora_rank, 1, n_head, n_tokens]
-
- ggml_tensor * kq = ggml_mul_mat(ctx0, k_g, q_g); // [n_sel, 1, n_head, n_tokens]
- ggml_prec_set_acc(kq, GGML_PREC_F32);
- kq = ggml_soft_max_ext(ctx0, kq, inp_kpool->gather_mask, kq_scale, 0.0f);
- cb(kq, "kq_soft_max_gathered", il);
-
- ggml_tensor * v_t = ggml_cont(ctx0, ggml_transpose(ctx0, k_g)); // [n_sel, kv_lora_rank, 1, n_tokens]
- ggml_tensor * kqv = ggml_mul_mat(ctx0, v_t, kq); // [kv_lora_rank, 1, n_head, n_tokens]
- kqv = ggml_mul_mat(ctx0, layer.wv_b, kqv); // [n_embd_head_v, 1, n_head, n_tokens]
- cb(kqv, "kqv_gathered", il);
-
- out = ggml_cont(ctx0, ggml_permute(ctx0, kqv, 0, 2, 1, 3)); // [n_embd_head_v, n_head, 1, n_tokens]
- out = ggml_reshape_2d(ctx0, out, kqv->ne[0]*n_head, n_tokens);
- } else {
- // The scatter selection already includes the causal mask.
- ggml_tensor * mask = ggml_reshape_4d(ctx0, sel, kq_mask->ne[0], kq_mask->ne[1], kq_mask->ne[2], kq_mask->ne[3]);
- cb(mask, "kq_mask_dsa", il);
-
- ggml_tensor * k = mctx_mla->get_k(ctx0, il);
- ggml_tensor * v = ggml_view_4d(ctx0, k, kv_lora_rank, k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0);
-
- out = build_attn_mha(q_absorbed, k, v, nullptr, mask, nullptr, layer.wv_b, inp_kpool->n_sel, kq_scale, il);
- }
+ ggml_tensor * out = build_attn_mha(q_absorbed, k, v, nullptr, mask, nullptr, layer.wv_b, inp_kpool->n_sel, kq_scale, il);
cb(out, "kqv_out", il);
out = ggml_mul_mat(ctx0, layer.wo, out);
diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp
index aca8f6065..d05f72fc3 100644
--- a/src/models/qwen4exp.cpp
+++ b/src/models/qwen4exp.cpp
@@ -651,7 +651,7 @@ public:
void set_input(const llama_ubatch * ubatch) override {
mctx->get_idx()->set_input_k_idxs(k_idxs, ubatch);
- mctx->set_input_kpool(pool_cells, pool_idxs, pool_mask, tail_idxs, nullptr, false, new_pool_idxs, new_pool_rep,
+ mctx->set_input_kpool(pool_cells, pool_idxs, pool_mask, tail_idxs, nullptr, new_pool_idxs, new_pool_rep,
ubatch, new_pool_pos);
}