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

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:

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,
+    // group s holds the conv window s tokens back.
-        conv_x->nb[1], conv_x->nb[2], n_seq_tokens * conv_x->nb[0]);
+    // [TAG_RECURRENT_ROLLBACK_SPLITS]: the last K_rs tokens must share one ubatch.
-    ggml_build_forward_expand(gf,
+    for (int64_t s = 0; s < K_rs; ++s) {
-        ggml_cpy(ctx0, last_conv_x,
+        const int64_t s_idx = std::max<int64_t>(0, n_seq_tokens - s);
-            ggml_view_3d(ctx0, conv_states_all, d_conv - 1, d_inner, n_seqs,
+        ggml_tensor * conv_x_s = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs,
-                (d_conv - 1)   * ggml_element_size(conv_states_all),
+            conv_x->nb[1], conv_x->nb[2], s_idx * conv_x->nb[0]);
-                n_embd_r_total * ggml_element_size(conv_states_all),
+        ggml_build_forward_expand(gf,
-                (kv_head * n_embd_r_total + qkv * conv_state_size) * ggml_element_size(conv_states_all))));
+            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);
+    const int64_t mem_size = mctx_cur->get_size();
-    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);
+    const int64_t K_rs     = (int64_t) cparams.n_rs_seq + 1;
-    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);
+
+    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 = build_recurrent_attn(inp_rs, ssm_states_all, Qcur, Kcur, Vcur, g1, beta, state, il);
-
+    output = ggml_cont(ctx0, output);
-    ggml_tensor * output    = ggml_cont(ctx0, attn_out.first);
     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

@@ -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

@@ -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) {
+static int test_rollback(const common_params & params, llama_model * model, uint8_t fill) {
-    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;
-    }
-
     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_src = make_ctx(params, model, fill);
-    llama_context * ctx_dst = make_ctx(params, model);
+    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;
+}