From 27c299a6ec3db3304193584580625d07085410c6 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Sat, 15 Aug 2026 22:40:41 +0200 Subject: [PATCH] wip --- common/debug.cpp | 10 ++++++++-- src/models/kimi-k3.cpp | 27 ++++++++++++++++++++------- src/models/models.h | 3 ++- 3 files changed, 30 insertions(+), 10 deletions(-) diff --git a/common/debug.cpp b/common/debug.cpp index 102c6924dc..62ac61f78e 100644 --- a/common/debug.cpp +++ b/common/debug.cpp @@ -162,14 +162,20 @@ bool common_debug_cb_eval(struct ggml_tensor * t, bool ask, void * user_data) { } } + // some ops have no sources, for example ggml_arange + char src0_str[128] = { 0 }; + if (src0) { + snprintf(src0_str, sizeof(src0_str), "%s{%s}", src0->name, common_ggml_ne_string(src0).c_str()); + } + char src1_str[128] = { 0 }; if (src1) { snprintf(src1_str, sizeof(src1_str), "%s{%s}", src1->name, common_ggml_ne_string(src1).c_str()); } if (matches_filter) { - LOG("%s: %24s = (%s) %10s(%s{%s}, %s}) = {%s}\n", __func__, t->name, ggml_type_name(t->type), - ggml_op_desc(t), src0->name, common_ggml_ne_string(src0).c_str(), src1 ? src1_str : "", + LOG("%s: %24s = (%s) %10s(%s, %s}) = {%s}\n", __func__, t->name, ggml_type_name(t->type), + ggml_op_desc(t), src0 ? src0_str : "", src1 ? src1_str : "", common_ggml_ne_string(t).c_str()); } diff --git a/src/models/kimi-k3.cpp b/src/models/kimi-k3.cpp index d952d72cdf..fb5abfff54 100644 --- a/src/models/kimi-k3.cpp +++ b/src/models/kimi-k3.cpp @@ -185,24 +185,30 @@ static ggml_tensor * kimi_k3_situ(ggml_context * ctx0, ggml_tensor * gate, ggml_ // cross-layer residual attention // -// layout is [n_embd, n_ckpt, n_tokens]: rms_norm reduces over ne0, dsv4_hc_pre over ne1 -// append the new checkpoint, do not re-fold the whole chain +// layout is [n_embd, n_ckpt_max, n_tokens]: rms_norm reduces over ne0, dsv4_hc_pre over ne1 +// the bank is allocated once, each push only writes the new checkpoint into its slot void llama_model_kimi_k3::graph::res_push(ggml_tensor * cur, int64_t n_embd, int64_t n_tokens) { ggml_tensor * ckpt = ggml_reshape_3d(ctx0, cur, n_embd, 1, n_tokens); - resi_stack = resi_stack ? ggml_concat(ctx0, resi_stack, ckpt, 1) : ckpt; + // keep the result view so that later reads depend on this write + resi_stack = ggml_set_inplace(ctx0, resi_stack, ckpt, + resi_stack->nb[1], resi_stack->nb[2], resi_stack->nb[3], + n_ckpt*resi_stack->nb[1]); + + n_ckpt++; } ggml_tensor * llama_model_kimi_k3::graph::res_mix(ggml_tensor * cur, ggml_tensor * score_w, int64_t n_tokens, int il) { - if (!resi_stack) { + if (n_ckpt == 0) { return cur; // layer 0: nothing banked yet } - const int n_ckpt = (int) resi_stack->ne[1]; - const float eps = hparams.f_norm_rms_eps; + const float eps = hparams.f_norm_rms_eps; - ggml_tensor * src = resi_stack; // [n_embd, n_ckpt, n_tokens] + // the written slots of the bank, strided over ne2 while the bank is not full + ggml_tensor * src = ggml_view_3d(ctx0, resi_stack, resi_stack->ne[0], n_ckpt, n_tokens, + resi_stack->nb[1], resi_stack->nb[2], 0); // one rms_norm scores all checkpoints at once // note: the scores use the normalized values, but the sum below uses the raw ones @@ -272,6 +278,13 @@ llama_model_kimi_k3::graph::graph(const llama_model & model, const llm_graph_par const bool use_attn_res = res_bs > 0; const int64_t n_embd_latent = hparams.n_expert_latent > 0 ? hparams.n_expert_latent : n_embd; + if (use_attn_res) { + // one checkpoint per res_bs layers, see res_push for the bank layout + const int64_t n_ckpt_max = (n_layer + res_bs - 1) / res_bs; + + resi_stack = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, n_embd, n_ckpt_max, n_tokens); + } + for (int il = 0; il < n_layer; ++il) { const auto & layer = model.layers[il]; diff --git a/src/models/models.h b/src/models/models.h index a181370474..c94abf580a 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2296,7 +2296,8 @@ struct llama_model_kimi_k3 : public llama_model_base { const llama_model & model; // Cross-layer residual attention (K3's `_apply_attn_res`). - ggml_tensor * resi_stack = nullptr; + ggml_tensor * resi_stack = nullptr; // checkpoint bank [n_embd, n_ckpt_max, n_tokens] + int n_ckpt = 0; // bank slots written so far void res_push(ggml_tensor * cur, int64_t n_embd, int64_t n_tokens); ggml_tensor * res_mix(ggml_tensor * cur, ggml_tensor * score_w,