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
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;
}