Commit 9dcf84e5a for llama.cpp
commit 9dcf84e5ae2718947188b539aab8b9c2b15d3ba1
Author: Frank Dai <soulmachine@gmail.com>
Date: Mon Sep 7 20:31:48 2026 -0700
model : support Kimi-K3 recurrent-state rollback (#28466)
diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp
index 15f651919..b5efb7206 100644
--- a/src/llama-arch.cpp
+++ b/src/llama-arch.cpp
@@ -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:
diff --git a/src/models/kimi-k3.cpp b/src/models/kimi-k3.cpp
index b7604cbf2..112b09849 100644
--- a/src/models/kimi-k3.cpp
+++ b/src/models/kimi-k3.cpp
@@ -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);
diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
index c46377c76..5531c4ce3 100644
--- a/tests/CMakeLists.txt
+++ b/tests/CMakeLists.txt
@@ -238,6 +238,15 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS)
set_tests_properties(test-recurrent-state-rollback-dsv4 PROPERTIES
FIXTURES_REQUIRED generate-models
)
+ llama_test(
+ test-recurrent-state-rollback
+ NAME test-recurrent-state-rollback-kimi-k3
+ LABEL main
+ ARGS -m "${MODEL_DIR}/kimi-k3-moe.gguf"
+ )
+ set_tests_properties(test-recurrent-state-rollback-kimi-k3 PROPERTIES
+ FIXTURES_REQUIRED generate-models
+ )
# Test state save/load functionality across all architectures, using the generated dummy models
llama_test(
diff --git a/tests/test-recurrent-state-rollback.cpp b/tests/test-recurrent-state-rollback.cpp
index c6f599e58..ef05de67d 100644
--- a/tests/test-recurrent-state-rollback.cpp
+++ b/tests/test-recurrent-state-rollback.cpp
@@ -1,22 +1,19 @@
#include "arg.h"
#include "common.h"
+#include "ggml-backend.h"
#include "llama.h"
+#include "../src/llama-io.h"
+#include "../src/llama-memory.h"
+
#include <algorithm>
#include <clocale>
#include <cmath>
#include <cstdio>
+#include <limits>
+#include <set>
#include <vector>
-static llama_context * make_ctx(const common_params & params, llama_model * model) {
- auto cparams = common_context_params_to_llama(params);
- cparams.n_seq_max = 1;
- cparams.n_rs_seq = 8;
- cparams.n_batch = std::max(cparams.n_batch, (uint32_t) (cparams.n_rs_seq + 1));
- cparams.n_ubatch = std::max(cparams.n_ubatch, (uint32_t) (cparams.n_rs_seq + 1));
- return llama_init_from_model(model, cparams);
-}
-
static bool decode_tokens(llama_context * ctx, const std::vector<llama_token> & tokens, uint32_t count) {
llama_batch batch = llama_batch_init(count, 0, 1);
for (uint32_t pos = 0; pos < count; ++pos) {
@@ -35,12 +32,70 @@ static bool decode_one(llama_context * ctx, llama_token tok, llama_pos pos) {
return ok;
}
+struct cache_buffer_collector : llama_io_write_i {
+ std::set<ggml_backend_buffer_t> buffers;
+ size_t size = 0;
+
+ void write(const void *, size_t n) override {
+ size += n;
+ }
+
+ void write_tensor(ggml_tensor * tensor, size_t, size_t n) override {
+ buffers.insert(tensor->buffer);
+ size += n;
+ }
+
+ size_t n_bytes() override {
+ return size;
+ }
+};
+
+static llama_context * init_ctx(llama_model * model, llama_context_params cparams, uint8_t fill) {
+ llama_context * ctx = llama_init_from_model(model, cparams);
+ if (ctx == nullptr || fill == 0) {
+ return ctx;
+ }
+
+ // Use a full ubatch so buffer discovery preserves prefill allocation sizes.
+ const uint32_t n_tokens = llama_n_ubatch(ctx);
+ if (!decode_tokens(ctx, std::vector<llama_token>(n_tokens, 0), n_tokens)) {
+ llama_free(ctx);
+ return nullptr;
+ }
+ llama_synchronize(ctx);
+ cache_buffer_collector collector;
+ llama_get_memory(ctx)->state_write(collector);
+ llama_memory_clear(llama_get_memory(ctx), true);
+ if (collector.buffers.empty()) {
+ fprintf(stderr, "%s : no cache buffers found\n", __func__);
+ llama_free(ctx);
+ return nullptr;
+ }
+ for (auto * buffer : collector.buffers) {
+ ggml_backend_buffer_clear(buffer, fill);
+ }
+ return ctx;
+}
+
+static llama_context * make_ctx(const common_params & params, llama_model * model, uint8_t fill) {
+ auto cparams = common_context_params_to_llama(params);
+ cparams.n_seq_max = 1;
+ cparams.n_rs_seq = 8;
+ cparams.n_batch = std::max(cparams.n_batch, (uint32_t) (cparams.n_rs_seq + 1));
+ cparams.n_ubatch = std::max(cparams.n_ubatch, (uint32_t) (cparams.n_rs_seq + 1));
+ return init_ctx(model, cparams, fill);
+}
+
+static float logit_diff(float a, float b) {
+ return std::isfinite(a) && std::isfinite(b) ? std::fabs(a - b) : std::numeric_limits<float>::infinity();
+}
+
// Roll back multiple sequences, then replay them in a single batch whose
// per-seq token count exceeds n_ubatch: each seq's replay spans several
// ubatches while its rollback restore is still pending. Compared against a
// reference context that never advanced past the rollback point and decodes
// the identical replay batch.
-static bool test_multi_seq_split_replay(const common_params & params, llama_model * model, const int n_vocab) {
+static bool test_multi_seq_split_replay(const common_params & params, llama_model * model, const int n_vocab, uint8_t fill) {
constexpr uint32_t n_seqs = 2;
constexpr uint32_t n_ubatch = 16;
constexpr uint32_t n_prompt = 19;
@@ -56,7 +111,7 @@ static bool test_multi_seq_split_replay(const common_params & params, llama_mode
cparams.n_batch = 256;
cparams.n_ubatch = n_ubatch;
cparams.kv_unified = false;
- return llama_init_from_model(model, cparams);
+ return init_ctx(model, cparams, fill);
};
llama_context * ctx_roll = make_ctx_multi();
@@ -143,7 +198,7 @@ static bool test_multi_seq_split_replay(const common_params & params, llama_mode
return false;
}
for (int t = 0; t < n_vocab; ++t) {
- const float diff = std::fabs(l_roll[t] - l_ref[t]);
+ const float diff = logit_diff(l_roll[t], l_ref[t]);
if (diff > eps && pos_first < 0) {
seq_first = i/n_replay;
pos_first = p0 + (int32_t) (i%n_replay);
@@ -191,7 +246,7 @@ static bool test_multi_seq_split_replay(const common_params & params, llama_mode
const float * l_ref = llama_get_logits_ith(ctx_ref, 0);
ok = l_roll != nullptr && l_ref != nullptr;
for (int t = 0; ok && t < n_vocab; ++t) {
- diff_tail = std::max(diff_tail, std::fabs(l_roll[t] - l_ref[t]));
+ diff_tail = std::max(diff_tail, logit_diff(l_roll[t], l_ref[t]));
}
}
@@ -207,38 +262,12 @@ static bool test_multi_seq_split_replay(const common_params & params, llama_mode
return true;
}
-int main(int argc, char ** argv) {
- std::setlocale(LC_NUMERIC, "C");
-
- common_params params;
- params.sampling.seed = 1234;
- params.n_predict = 1;
-
- common_init();
-
- if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_COMMON)) {
- return 1;
- }
-
- ggml_backend_load_all();
-
- common_init_result_ptr llama_init = common_init_from_params(params);
- llama_model * model = llama_init->model();
- if (model == nullptr) {
- fprintf(stderr, "%s : failed to init model\n", __func__);
- return 1;
- }
-
- if (!llama_model_is_recurrent(model) && !llama_model_is_hybrid(model)) {
- fprintf(stderr, "%s : skipping for non-recurrent model\n", __func__);
- return 0;
- }
-
+static int test_rollback(const common_params & params, llama_model * model, uint8_t fill) {
const llama_vocab * vocab = llama_model_get_vocab(model);
const int n_vocab = llama_vocab_n_tokens(vocab);
- llama_context * ctx_src = make_ctx(params, model);
- llama_context * ctx_dst = make_ctx(params, model);
+ llama_context * ctx_src = make_ctx(params, model, fill);
+ llama_context * ctx_dst = make_ctx(params, model, fill);
if (ctx_src == nullptr || ctx_dst == nullptr) {
fprintf(stderr, "%s : failed to init contexts\n", __func__);
return 1;
@@ -311,7 +340,7 @@ int main(int argc, char ** argv) {
logits_src_replay[i].assign(logits_src, logits_src + n_vocab);
for (int token = 0; token < n_vocab; ++token) {
- if (std::fabs(logits_src[token] - logits_dst[token]) > eps) {
+ if (logit_diff(logits_src[token], logits_dst[token]) > eps) {
fprintf(stderr, "%s : %s logits mismatch at position %d, token %d (%g != %g)\n",
__func__, mode, pos, token, (double) logits_src[token], (double) logits_dst[token]);
return false;
@@ -342,7 +371,7 @@ int main(int argc, char ** argv) {
// Repeat the load into a context that already has its own rollback state:
// groups 1..n_rs_seq hold a different prompt's history, and rs_idx[0] is
// non-zero at load time. The restore must wipe that state and still match.
- llama_context * ctx_dirty = make_ctx(params, model);
+ llama_context * ctx_dirty = make_ctx(params, model, fill);
if (ctx_dirty == nullptr) {
fprintf(stderr, "%s : failed to init dirty ctx\n", __func__);
return 1;
@@ -380,7 +409,7 @@ int main(int argc, char ** argv) {
}
for (int token = 0; token < n_vocab; ++token) {
- if (std::fabs(logits_src_replay[i][token] - logits_dirty[token]) > eps) {
+ if (logit_diff(logits_src_replay[i][token], logits_dirty[token]) > eps) {
fprintf(stderr, "%s : dirty-ctx logits mismatch at position %d, token %d (%g != %g)\n",
__func__, pos, token, (double) logits_src_replay[i][token], (double) logits_dirty[token]);
return 1;
@@ -393,9 +422,46 @@ int main(int argc, char ** argv) {
llama_free(ctx_dst);
llama_free(ctx_dirty);
- if (!test_multi_seq_split_replay(params, model, n_vocab)) {
+ if (!test_multi_seq_split_replay(params, model, n_vocab, fill)) {
+ return 1;
+ }
+
+ return 0;
+}
+
+int main(int argc, char ** argv) {
+ std::setlocale(LC_NUMERIC, "C");
+
+ common_params params;
+ params.sampling.seed = 1234;
+ params.n_predict = 1;
+
+ common_init();
+
+ if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_COMMON)) {
return 1;
}
+ ggml_backend_load_all();
+
+ common_init_result_ptr llama_init = common_init_from_params(params);
+ llama_model * model = llama_init->model();
+ if (model == nullptr) {
+ fprintf(stderr, "%s : failed to init model\n", __func__);
+ return 1;
+ }
+
+ if (!llama_model_is_recurrent(model) && !llama_model_is_hybrid(model)) {
+ fprintf(stderr, "%s : skipping for non-recurrent model\n", __func__);
+ return 0;
+ }
+
+ for (uint8_t fill : { 0, 0x3e }) {
+ fprintf(stderr, "%s : testing with cache fill 0x%02x\n", __func__, fill);
+ if (test_rollback(params, model, fill) != 0) {
+ return 1;
+ }
+ }
+
return 0;
}