From 910cc2a38e2d0daddebfdfc8a836bd0f8213a8f2 Mon Sep 17 00:00:00 2001 From: Aman Gupta Date: Sun, 25 Oct 2026 22:53:42 +0800 Subject: [PATCH] add GLM5.3 --- ggml/src/ggml-backend-meta.cpp | 6 ++++- ggml/src/ggml-rpc/ggml-rpc.cpp | 8 ++++++ src/models/glm5-next.cpp | 46 +++++++++++++++++++++++----------- 3 files changed, 45 insertions(+), 15 deletions(-) diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index 2d3356d695..ae2f59eeca 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -704,8 +704,12 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( if (ggml_is_permuted(tensor) && !ggml_is_permuted(tensor->src[0]) && axis >= 0 && axis < GGML_MAX_DIMS-1) { const size_t nb_axis = tensor->src[0]->nb[axis]; const size_t nb_next = tensor->src[0]->nb[axis+1]; + // the span can also be one repetition of a repeated segment, e.g. the q, k or v part of a fused conv state + const size_t nb_rep = src_ss[0].n_segments == 1 && nb_next % src_ss[0].nr[0] == 0 ? nb_next/src_ss[0].nr[0] : nb_next; for (int dim = 0; dim < GGML_MAX_DIMS; dim++) { - if (tensor->ne[dim] == 1 || tensor->nb[dim] < nb_axis || tensor->nb[dim]*tensor->ne[dim] != nb_next) { + const size_t span = tensor->nb[dim]*tensor->ne[dim]; + if (tensor->ne[dim] == 1 || tensor->nb[dim] < nb_axis || + (span != nb_next && (span != nb_rep || tensor->view_offs % nb_next % nb_rep != 0))) { continue; } bool others_ok = true; diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp index 49f11073e9..f6c227594c 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp @@ -509,6 +509,7 @@ public: void busy_spin_acquire(); void busy_spin_release(); void graph_compute(uint32_t device, const ggml_cgraph * cgraph); + void graph_cache_clear(); void start(const std::string & endpoint); void work(); @@ -729,6 +730,7 @@ static void ggml_backend_rpc_buffer_free_buffer(ggml_backend_buffer_t buffer) { request->remote_ptr = ctx->remote_ptr; LOG_DBG2("[%s] remote_ptr: 0x%" PRIx64 "\n", __func__, request->remote_ptr); ctx->dispatcher->send(RPC_CMD_FREE_BUFFER, request, sizeof(*request)); + ctx->dispatcher->graph_cache_clear(); delete ctx; } @@ -1267,6 +1269,12 @@ void rpc_dispatcher::graph_compute(uint32_t device, const ggml_cgraph * cgraph) } } +// the server drops all of its cached graphs when a buffer is freed, see rpc_server::free_buffer +void rpc_dispatcher::graph_cache_clear() { + std::lock_guard lock(graph_mutex); + graph_uids.clear(); +} + static enum ggml_status ggml_backend_rpc_graph_compute(ggml_backend_t backend, ggml_cgraph * cgraph) { ggml_backend_rpc_context * rpc_ctx = (ggml_backend_rpc_context *)backend->context; rpc_ctx->dispatcher->graph_compute(rpc_ctx->device, cgraph); diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp index 31c019e785..4a790cf11d 100644 --- a/src/models/glm5-next.cpp +++ b/src/models/glm5-next.cpp @@ -196,7 +196,7 @@ static ggml_tensor * glm5_conv1d(ggml_cgraph * gf, ggml_context * ctx0, int64_t qkv, ggml_tensor * x, ggml_tensor * proj_w, ggml_tensor * conv_w, int64_t d_conv, int64_t head_dim, int64_t n_head, int64_t n_seq_tokens, int64_t n_seqs, int64_t n_tokens, int64_t kv_head, - int64_t mem_size, int64_t K_rs) { + int64_t mem_size, int64_t K_rs, ggml_tensor * rs_slot_rev) { const int64_t d_inner = head_dim * n_head; const int64_t conv_state_size = (d_conv - 1) * d_inner; const int64_t n_embd_r_total = 3 * conv_state_size; @@ -210,18 +210,29 @@ static ggml_tensor * glm5_conv1d(ggml_cgraph * gf, ggml_context * ctx0, ggml_tensor * x_3d = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs); ggml_tensor * conv_x = ggml_concat(ctx0, conv_state_x, ggml_transpose(ctx0, x_3d), 0); - // group s holds the conv window s tokens back. - // [TAG_RECURRENT_ROLLBACK_SPLITS]: the last K_rs tokens must share one ubatch. - for (int64_t s = 0; s < K_rs; ++s) { - const int64_t s_idx = std::max(0, n_seq_tokens - s); - ggml_tensor * conv_x_s = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs, - conv_x->nb[1], conv_x->nb[2], s_idx * conv_x->nb[0]); + const size_t row_size = n_embd_r_total * ggml_element_size(conv_states_all); + const size_t offs = (kv_head * n_embd_r_total + qkv * conv_state_size) * ggml_element_size(conv_states_all); + + if (K_rs == 1) { + ggml_tensor * conv_x_last = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs, + conv_x->nb[1], conv_x->nb[2], n_seq_tokens * conv_x->nb[0]); ggml_build_forward_expand(gf, - ggml_cpy(ctx0, conv_x_s, + ggml_cpy(ctx0, conv_x_last, ggml_view_3d(ctx0, conv_states_all, d_conv - 1, d_inner, n_seqs, - (d_conv - 1) * ggml_element_size(conv_states_all), - n_embd_r_total * ggml_element_size(conv_states_all), - ((s * mem_size + kv_head) * n_embd_r_total + qkv * conv_state_size) * ggml_element_size(conv_states_all)))); + (d_conv - 1) * ggml_element_size(conv_states_all), row_size, offs))); + } else { + // [TAG_RECURRENT_ROLLBACK_SPLITS]: the last K_rs tokens must share one ubatch. + // window w starts at column n_seq_tokens - (Kw - 1) + w and is rollback slot Kw - 1 - w, see build_conv_state + const int64_t Kw = rs_slot_rev->ne[0]; + ggml_tensor * windows = ggml_view_4d(ctx0, conv_x, d_conv - 1, 1, d_inner, n_seqs, + conv_x->nb[0], conv_x->nb[1], conv_x->nb[2], (n_seq_tokens - (Kw - 1)) * conv_x->nb[0]); + windows->ne[1] = Kw; + + // [d_conv - 1, slot, channel, seq] view of this part of the cells of this ubatch + ggml_tensor * slots = ggml_view_4d(ctx0, conv_states_all, d_conv - 1, K_rs, d_inner, n_seqs, + mem_size * row_size, (d_conv - 1) * ggml_element_size(conv_states_all), row_size, offs); + + ggml_build_forward_expand(gf, ggml_set_rows(ctx0, slots, windows, rs_slot_rev)); } ggml_tensor * conv_weight = ggml_reshape_2d(ctx0, conv_w, d_conv, d_inner); @@ -819,9 +830,16 @@ ggml_tensor * llama_model_glm5_next::graph::build_kda_layer( const int64_t mem_size = mctx_cur->get_size(); const int64_t K_rs = (int64_t) cparams.n_rs_seq + 1; - ggml_tensor * Qcur = glm5_conv1d(gf, ctx0, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head, mem_size, K_rs); - ggml_tensor * Kcur = glm5_conv1d(gf, ctx0, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head, mem_size, K_rs); - ggml_tensor * Vcur = glm5_conv1d(gf, ctx0, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head, mem_size, K_rs); + // {Kw-1, ..., 0} for the rollback windows that exist in this ubatch, shared by every layer + if (K_rs > 1 && rs_slot_rev == nullptr) { + const int64_t Kw = std::min(n_seq_tokens + 1, K_rs); + rs_slot_rev = ggml_argsort(ctx0, ggml_arange(ctx0, 0.0f, (float) Kw, 1.0f), GGML_SORT_ORDER_DESC); + cb(rs_slot_rev, "rs_slot_rev", -1); + } + + ggml_tensor * Qcur = glm5_conv1d(gf, ctx0, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head, mem_size, K_rs, rs_slot_rev); + ggml_tensor * Kcur = glm5_conv1d(gf, ctx0, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head, mem_size, K_rs, rs_slot_rev); + ggml_tensor * Vcur = glm5_conv1d(gf, ctx0, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head, mem_size, K_rs, rs_slot_rev); cb(Qcur, "kda_q_conv", il); cb(Kcur, "kda_k_conv", il); cb(Vcur, "kda_v_conv", il);