Commit 6afa9a313 for llama.cpp
commit 6afa9a3134dc4833ab58e625464aa67e5fc7eea5
Author: Aman Gupta <amangupta052@gmail.com>
Date: Fri Oct 16 08:44:36 2026 +0530
llama : share the k-pool input views between layers (#30613)
The scheduler uploads each view of a host input on its own, so the per-layer reshape of new_pool_idxs (and of sel_mask in glm5-next) was copied once per layer.
diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp
index a1b722950..90870bd96 100644
--- a/src/models/glm5-next.cpp
+++ b/src/models/glm5-next.cpp
@@ -271,8 +271,10 @@ public:
ggml_tensor * pool_mask = nullptr; // F32/F16 [n_pool, n_tokens]
ggml_tensor * tail_idxs = nullptr; // I32 [kpool - 1, n_tokens]
ggml_tensor * sel_mask = nullptr; // F32 [n_sel, 1, 1, n_tokens] 0 for live selection slots, -inf for dead ones
+ ggml_tensor * sel_mask_2d = nullptr; // F32 [n_sel, n_tokens] view of sel_mask, shared by the layers
// 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_flat = nullptr; // I32 [kpool*n_new] view of new_pool_idxs, shared by the layers
ggml_tensor * new_pool_rep = nullptr; // I64 [n_new] cell to write each new pooled key into
const llama_memory_hybrid_idx_context * mctx;
@@ -326,12 +328,15 @@ llama_model_glm5_next::llm_graph_input_kpool * llama_model_glm5_next::graph::bui
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);
+ // one view for all layers: each view of a host input is a separate upload
+ inp->sel_mask_2d = ggml_reshape_2d(ctx0, inp->sel_mask, n_sel, n_tokens);
}
inp->n_new = n_new;
inp->new_pool_idxs = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, kpool, n_new);
ggml_set_input(inp->new_pool_idxs);
+ inp->new_pool_flat = ggml_reshape_1d(ctx0, inp->new_pool_idxs, kpool*n_new);
// 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);
@@ -917,7 +922,7 @@ ggml_tensor * llama_model_glm5_next::graph::build_kpool_select(
const int64_t n_kv = mctx_lid->get_n_kv();
// Pool the entries completed by this ubatch. The last one is a dummy when the ubatch completes none.
- ggml_tensor * rows = kpool_cache.gather_key_gate(ggml_reshape_1d(ctx0, inp_kpool->new_pool_idxs, kpool*n_new));
+ ggml_tensor * rows = kpool_cache.gather_key_gate(inp_kpool->new_pool_flat);
rows = ggml_reshape_3d(ctx0, rows, 2*n_embd_indexer, kpool, n_new);
ggml_tensor * pk = ggml_view_3d(ctx0, rows, n_embd_indexer, kpool, n_new, rows->nb[1], rows->nb[2], 0);
@@ -1005,7 +1010,7 @@ 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->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 * live = ggml_exp(ctx0, inp_kpool->sel_mask_2d);
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);
diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp
index e055abdf5..77430e7af 100644
--- a/src/models/qwen4exp.cpp
+++ b/src/models/qwen4exp.cpp
@@ -682,6 +682,7 @@ public:
ggml_tensor * pool_mask = nullptr; // F16 [n_pool, n_tokens]
ggml_tensor * tail_idxs = nullptr; // I32 [kpool - 1, n_tokens]
ggml_tensor * new_pool_idxs = nullptr; // I32 [kpool, n_new] members of the blocks to re-pool this ubatch
+ ggml_tensor * new_pool_flat = nullptr; // I32 [kpool*n_new] view of new_pool_idxs, shared by the QSA layers
ggml_tensor * new_pool_rep = nullptr; // I64 [n_new] cell to write each new pooled key into
ggml_tensor * new_pool_pos = nullptr; // I32 [4*n_new] M-RoPE position of each new block's first member
@@ -724,6 +725,8 @@ 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);
+ // one view for all layers: each view of a host input is a separate upload
+ inp->new_pool_flat = ggml_reshape_1d(ctx0, inp->new_pool_idxs, kpool*inp->n_new);
// 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);
@@ -765,7 +768,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_sel(
auto kpool_cache = mctx_hyb->get_kpool_access(ctx0, il, idx_dim);
// pool only the blocks this ubatch completes or regroups
- ggml_tensor * rows = kpool_cache.gather_key_gate(ggml_reshape_1d(ctx0, inp_kpool->new_pool_idxs, kpool*n_new));
+ ggml_tensor * rows = kpool_cache.gather_key_gate(inp_kpool->new_pool_flat);
rows = ggml_reshape_3d(ctx0, rows, idx_dim, kpool, n_new);
// mean over the members; kpool is small, so summing slices beats a transpose plus sum_rows