Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

model : support Kimi-K3 recurrent-state rollback (#28466)

Browse Source
This commit is contained in:
Frank Dai
2026-09-08 11:31:48 +08:00
committed by GitHub
parent 050dde50c9
commit 9dcf84e5ae
4 changed files with 147 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/llama-arch.cpp
+1
View File
@@ -1103,6 +1103,7 @@ bool llm_arch_is_diffusion(const llm_arch & arch) {
bool llm_arch_supports_rs_rollback(const llm_arch & arch) {
switch (arch) {
case LLM_ARCH_KIMI_K3:
case LLM_ARCH_QWEN35:
case LLM_ARCH_QWEN35MOE:
case LLM_ARCH_QWEN4EXP:
src/models/kimi-k3.cpp
+25 -21
View File
@@ -1,4 +1,6 @@
#include "models.h"
#include <algorithm>
#include "llama-memory-recurrent.h"
//
@@ -357,7 +359,8 @@ static ggml_tensor * kimi_k3_conv1d(ggml_cgraph * gf, ggml_context * ctx0,
ggml_tensor * conv_states_all, ggml_tensor * conv_state_all,
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 n_seq_tokens, int64_t n_seqs, int64_t n_tokens, int64_t kv_head,
int64_t mem_size, int64_t K_rs) {
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;
@@ -371,14 +374,19 @@ static ggml_tensor * kimi_k3_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);
ggml_tensor * last_conv_x = 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, last_conv_x,
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),
(kv_head * n_embd_r_total + qkv * conv_state_size) * ggml_element_size(conv_states_all))));
// 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<int64_t>(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]);
ggml_build_forward_expand(gf,
ggml_cpy(ctx0, conv_x_s,
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))));
}
ggml_tensor * conv_weight = ggml_reshape_2d(ctx0, conv_w, d_conv, d_inner);
ggml_tensor * Xcur = ggml_ssm_conv(ctx0, conv_x, conv_weight);
@@ -399,9 +407,12 @@ ggml_tensor * llama_model_kimi_k3::graph::build_kda_layer(
ggml_tensor * conv_states_all = mctx_cur->get_r_l(il);
ggml_tensor * conv_state_all = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs);
ggml_tensor * Qcur = kimi_k3_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);
ggml_tensor * Kcur = kimi_k3_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);
ggml_tensor * Vcur = kimi_k3_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);
const int64_t mem_size = mctx_cur->get_size();
const int64_t K_rs = (int64_t) cparams.n_rs_seq + 1;
ggml_tensor * Qcur = kimi_k3_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 = kimi_k3_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 = kimi_k3_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);
cb(Qcur, "kda_q_conv", il);
cb(Kcur, "kda_k_conv", il);
cb(Vcur, "kda_v_conv", il);
@@ -445,16 +456,9 @@ ggml_tensor * llama_model_kimi_k3::graph::build_kda_layer(
Qcur = build_gdn_l2_norm(ctx0, Qcur, eps_norm);
Kcur = build_gdn_l2_norm(ctx0, Kcur, eps_norm);
auto attn_out = build_delta_net(Qcur, Kcur, Vcur, g1, beta, state, il);
ggml_tensor * output = ggml_cont(ctx0, attn_out.first);
ggml_tensor * output = build_recurrent_attn(inp_rs, ssm_states_all, Qcur, Kcur, Vcur, g1, beta, state, il);
output = ggml_cont(ctx0, output);
cb(output, "kda_scan_out", il);
ggml_tensor * new_state = attn_out.second;
ggml_build_forward_expand(gf,
ggml_cpy(ctx0, new_state,
ggml_view_1d(ctx0, ssm_states_all, hparams.n_embd_s() * n_seqs,
kv_head * hparams.n_embd_s() * ggml_element_size(ssm_states_all))));
// K3: single full-rank gate (kimi-linear factors this as g_b(g_a(x)))
ggml_tensor * cur_2d = ggml_reshape_2d(ctx0, cur_3d, cur_3d->ne[0], n_seq_tokens * n_seqs);