From 5c19a5a7560024bb19745b4ae04a75613d4b9aa3 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Thu, 22 Oct 2026 16:59:29 +0300 Subject: [PATCH] glm5 : optimize inputs (#30932) * glm5 : avoid sel_mask_2d input * cont : avoid new_pool_flat input --- src/llama-memory-hybrid-idx.cpp | 2 +- src/models/glm5-next.cpp | 13 ++++--------- 2 files changed, 5 insertions(+), 10 deletions(-) diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index 2c5a57680..8113f5c21 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -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); diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp index 5056c901b..eb2e5bd59 100644 --- a/src/models/glm5-next.cpp +++ b/src/models/glm5-next.cpp @@ -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);