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);
tests/CMakeLists.txt
+9
View File
@@ -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(
tests/test-recurrent-state-rollback.cpp
+112 -46
View File
@@ -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;
}