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 #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(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 #include #include #include +#include +#include #include -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 & 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 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(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::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; +}