darthrubik/llama-cpp
1
0
Fork 0
Code Issues Pull Requests 1 Packages Projects Releases Wiki Activity

glm5 : optimize inputs (#30932)
Copilot Setup Steps / copilot-setup-steps (push) Has been cancelled
Check Pre-Tokenizer Hashes / pre-tokenizer-hashes (push) Has been cancelled
Python check requirements.txt / check-requirements (push) Has been cancelled
Python Type-Check / python type-check (push) Has been cancelled
Update Operations Documentation / update-ops-docs (push) Has been cancelled

Browse Source 
* glm5 : avoid sel_mask_2d input

* cont : avoid new_pool_flat input
This commit is contained in:
Georgi Gerganov
2026-10-22 16:59:29 +03:00
committed by GitHub
parent d9ee007a06
commit 5c19a5a756
2 changed files with 5 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/llama-memory-hybrid-idx.cpp
+1 -1
View File
@@ -758,7 +758,7 @@ void llama_memory_hybrid_idx_context::set_input_kpool(ggml_tensor * pool_cells,
GGML_ASSERT(tail_idxs->ne[0] == (int64_t) kpool - 1 && tail_idxs->ne[1] == (int64_t) n_tokens);
GGML_ASSERT(pool_idxs->ne[0] == (int64_t) kpool && pool_idxs->ne[1] == (int64_t) n_pool);
GGML_ASSERT(ggml_backend_buffer_is_host(new_pool_idxs->buffer));
GGML_ASSERT(new_pool_idxs->ne[0] == (int64_t) kpool && new_pool_idxs->ne[1] == (int64_t) n_new_g);
GGML_ASSERT(ggml_nelements(new_pool_idxs) == (int64_t) kpool*n_new_g);
// the graph always scatters the fresh pooled keys back into the cache, see build_qsa_sel
GGML_ASSERT(new_pool_rep != nullptr && ggml_backend_buffer_is_host(new_pool_rep->buffer));
GGML_ASSERT(new_pool_rep->ne[0] == (int64_t) n_new_g);
src/models/glm5-next.cpp
+4 -9
View File
@@ -271,10 +271,8 @@ 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_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
const llama_memory_hybrid_idx_context * mctx;
@@ -329,16 +327,13 @@ 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);
inp->new_pool_idxs = ggml_new_tensor_1d(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);
@@ -924,7 +919,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(inp_kpool->new_pool_flat);
ggml_tensor * rows = kpool_cache.gather_key_gate(inp_kpool->new_pool_idxs);
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);
@@ -1023,7 +1018,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, inp_kpool->sel_mask_2d);
ggml_tensor * live = ggml_reshape_2d(ctx0, ggml_exp(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);