Commit a17e94d22 for llama.cpp
commit a17e94d229fc7fed1567a799baebd11b771fe5af
Author: Georgi Gerganov <ggerganov@gmail.com>
Date: Tue Oct 13 16:21:55 2026 +0300
tests : consolidate context state tests into test-llama-context (#30023)
* tests : consolidate the context state tests into test-llama-context
test-save-load-state, test-recurrent-state-rollback and
test-state-restore-fragmented check the same object from three places:
the KV cache and the recurrent state of a llama_context, saved, loaded,
copied, rolled back. Each had its own main, its own loop over the
models, its own copy of the token and the comparison helpers, and a run
of all three loaded the same models several times.
They are one binary now, and the case is the unit of the test:
tests/llama-context/test-case.h the interface a case implements
tests/llama-context/cases-kv.cpp the cases of the KV cache
tests/llama-context/cases-rs.cpp the cases of the recurrent state
tests/llama-context/test-util.cpp tokens, generation, comparison
tests/llama-context/test-runner.cpp the models of a case, on threads
tests/llama-context/test-log.cpp the log file of a model
A run is a table: a row per model, a column per case, a cell holds the
verdict of that pair. The models load on several threads (-j, --threads,
min(cores, 4) by default), the table is printed when it is complete and
the log of a model that failed is printed after it, in the order of the
rows. A case returns a verdict and the reason for it, and the runner is
the only one who reports - the test does not use tests/testing.h, a run
over 130 models is a matrix of verdicts and not a tree of subtests.
Staging no longer travels through the environment: llama_model_set_attn_rot
replaces the LLAMA_ATTN_ROT_DISABLE variable, which a model loaded in one
thread used to change for the models the other threads were loading, and
llama_memory_get_buffers is added for the cases that copy the buffers of
a memory. common_log_set_stdo keeps the console quiet while the table is
being built.
ggml_aligned_malloc allocates with posix_memalign when the thread
sanitizer is on. On macOS it allocates with the Mach VM, which the
sanitizer does not intercept, so it never sees a page range change hands
and reports the first write into recycled memory as a race - 30 of them
in this test, and every one of them false.
Assisted-by: pi:llama.cpp/Qwen3.8-Flash-Next
* cont : run single-threaded by default
* cont : try to build the tests on all platforms
* cont : trim down some body fat
* cont : fix
* cont : fix the fix
diff --git a/.github/workflows/build-apple.yml b/.github/workflows/build-apple.yml
index f792256f6..70aefd5c8 100644
--- a/.github/workflows/build-apple.yml
+++ b/.github/workflows/build-apple.yml
@@ -105,7 +105,7 @@ jobs:
cd build
# Metal Paravirtual devices are difficult to support -> disable
# ref: https://github.com/ggml-org/llama.cpp/pull/19802#issuecomment-4013704023
- ctest -L main -E "test-llama-archs|test-save-load-state|test-recurrent-state-rollback" --verbose --timeout 900
+ ctest -L main -E "test-llama-archs|test-llama-context" --verbose --timeout 900
macos-latest-x64:
runs-on: macos-15-intel
diff --git a/.github/workflows/build-openvino.yml b/.github/workflows/build-openvino.yml
index 25e10c46d..f4c13b611 100644
--- a/.github/workflows/build-openvino.yml
+++ b/.github/workflows/build-openvino.yml
@@ -37,7 +37,7 @@ env:
LLAMA_ARG_LOG_PREFIX: 1
LLAMA_ARG_LOG_TIMESTAMPS: 1
# TODO: fix failing tests on OpenVINO backend
- CTEST_EXCLUDE: "test-llama-archs|^test-recurrent-state-|test-save-load-state"
+ CTEST_EXCLUDE: "test-llama-archs|test-llama-context"
jobs:
ubuntu-24-openvino:
diff --git a/ci/run.sh b/ci/run.sh
index e96b6ba49..8bebb8b9a 100755
--- a/ci/run.sh
+++ b/ci/run.sh
@@ -180,7 +180,7 @@ if [ ! -z ${GG_BUILD_OPENVINO} ]; then
CMAKE_EXTRA="${CMAKE_EXTRA} -DGGML_OPENVINO=ON"
# TODO: fix failing tests on OpenVINO backend
- CTEST_EXTRA="-E test-llama-archs|^test-recurrent-state-|test-save-load-state"
+ CTEST_EXTRA="-E test-llama-archs|test-llama-context"
fi
## helpers
@@ -447,10 +447,10 @@ function gg_run_qwen3_0_6b {
(time ./bin/llama-imatrix --model ${model_f16} -f ${wiki_test} -ngl 99 -c 1024 -b 512 --chunks 2 ) 2>&1 | tee -a $OUT/${ci}-imatrix.log
- (time ./bin/test-save-load-state --model ${model_q4_0} -ngl 10 -c 1024 -fa off --no-op-offload) 2>&1 | tee -a $OUT/${ci}-save-load-state.log
- (time ./bin/test-save-load-state --model ${model_q4_0} -ngl 10 -c 1024 -fa on --no-op-offload) 2>&1 | tee -a $OUT/${ci}-save-load-state.log
- (time ./bin/test-save-load-state --model ${model_q4_0} -ngl 99 -c 1024 -fa off ) 2>&1 | tee -a $OUT/${ci}-save-load-state.log
- (time ./bin/test-save-load-state --model ${model_q4_0} -ngl 99 -c 1024 -fa on ) 2>&1 | tee -a $OUT/${ci}-save-load-state.log
+ (time ./bin/test-llama-context --model ${model_q4_0} -ngl 10 -c 1024 -fa off --no-op-offload) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
+ (time ./bin/test-llama-context --model ${model_q4_0} -ngl 10 -c 1024 -fa on --no-op-offload) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
+ (time ./bin/test-llama-context --model ${model_q4_0} -ngl 99 -c 1024 -fa off ) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
+ (time ./bin/test-llama-context --model ${model_q4_0} -ngl 99 -c 1024 -fa on ) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
function check_ppl {
qnt="$1"
diff --git a/common/log.cpp b/common/log.cpp
index 86c98e8e6..25f12684e 100644
--- a/common/log.cpp
+++ b/common/log.cpp
@@ -162,6 +162,7 @@ struct common_log {
// default capacity
common_log(size_t capacity = 512) {
file = nullptr;
+ stdo = true;
colors = false;
prefix = false;
timestamps = false;
@@ -190,6 +191,7 @@ private:
FILE * file;
+ bool stdo;
bool colors;
bool prefix;
bool timestamps;
@@ -205,7 +207,9 @@ private:
bool print_entry(const common_log_entry & e) const {
if (e.is_end) return true;
- e.print();
+ if (stdo) {
+ e.print();
+ }
if (file) {
e.print(file);
}
@@ -400,6 +404,14 @@ public:
resume();
}
+ void set_stdo(bool enabled) {
+ pause();
+
+ stdo = enabled;
+
+ resume();
+ }
+
bool get_colors() const {
return colors;
}
@@ -497,6 +509,10 @@ void common_log_set_file(struct common_log * log, const char * file) {
log->set_file(file);
}
+void common_log_set_stdo(struct common_log * log, bool enabled) {
+ log->set_stdo(enabled);
+}
+
void common_log_set_colors(struct common_log * log, log_colors colors) {
if (colors == LOG_COLORS_AUTO) {
log->set_colors(tty_can_use_colors());
diff --git a/common/log.h b/common/log.h
index e9e1f1761..93ea58e76 100644
--- a/common/log.h
+++ b/common/log.h
@@ -92,6 +92,7 @@ void common_log_add(struct common_log * log, enum ggml_log_level level, const ch
//
void common_log_set_file (struct common_log * log, const char * file); // not thread-safe
+void common_log_set_stdo (struct common_log * log, bool enabled); // whether to output to stdout/stderr
void common_log_set_colors (struct common_log * log, log_colors colors); // not thread-safe
bool common_log_get_colors (struct common_log * log); // whether colors are enabled
void common_log_set_prefix (struct common_log * log, bool prefix); // whether to output prefix to each log
diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c
index 882cdf4a5..ef0134ba0 100644
--- a/ggml/src/ggml.c
+++ b/ggml/src/ggml.c
@@ -329,6 +329,17 @@ void ggml_log_callback_default(enum ggml_log_level level, const char * text, voi
#endif
+// the thread sanitizer does not intercept vm_allocate and vm_deallocate, so it produces false positives
+#if defined(__has_feature)
+#if __has_feature(thread_sanitizer)
+#define GGML_TSAN_ENABLED 1
+#endif
+#else
+#if defined(__SANITIZE_THREAD__)
+#define GGML_TSAN_ENABLED 1
+#endif
+#endif
+
void * ggml_aligned_malloc(size_t size) {
#if defined(__s390x__)
const int alignment = 256;
@@ -346,7 +357,7 @@ void * ggml_aligned_malloc(size_t size) {
void * aligned_memory = NULL;
#ifdef GGML_USE_CPU_HBM
int result = hbw_posix_memalign(&aligned_memory, alignment, size);
- #elif TARGET_OS_OSX
+ #elif TARGET_OS_OSX && !defined(GGML_TSAN_ENABLED)
GGML_UNUSED(alignment);
kern_return_t alloc_status = vm_allocate((vm_map_t) mach_task_self(), (vm_address_t *) &aligned_memory, size, VM_FLAGS_ANYWHERE);
int result = EFAULT;
@@ -393,7 +404,7 @@ void ggml_aligned_free(void * ptr, size_t size) {
if (ptr != NULL) {
hbw_free(ptr);
}
-#elif TARGET_OS_OSX
+#elif TARGET_OS_OSX && !defined(GGML_TSAN_ENABLED)
if (ptr != NULL) {
vm_deallocate((vm_map_t)mach_task_self(), (vm_address_t)ptr, size);
}
diff --git a/src/llama-ext.h b/src/llama-ext.h
index faf4beb69..b8c5e8033 100644
--- a/src/llama-ext.h
+++ b/src/llama-ext.h
@@ -8,6 +8,7 @@
#include <cstdint>
#include <map>
+#include <vector>
// Reserve a new compute graph. It is valid until the next call to llama_graph_reserve.
LLAMA_API struct ggml_cgraph * llama_graph_reserve(
@@ -90,6 +91,13 @@ LLAMA_API ggml_backend_dev_t llama_model_get_device(const struct llama_model * m
LLAMA_API llama_memory_breakdown llama_get_memory_breakdown(const struct llama_context * ctx);
+// the backend buffers that a state write of this memory would touch - empty caches report none
+LLAMA_API std::vector<ggml_backend_buffer_t> llama_memory_get_buffers(const llama_memory_t mem);
+
+// Turn the KV cache attention rotation on or off for the contexts created from this model.
+// Useful for testing purposes.
+LLAMA_API void llama_model_set_attn_rot(struct llama_model * model, bool value);
+
// Set whether the context outputs nextn embeddings or not
// If masked == true, output the embeddings only for the tokens with batch.logits != 0
// If masked == false, output the embeddings for all tokens in the batch regardless of batch.logits
diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp
index 804c90532..d5208d034 100644
--- a/src/llama-kv-cache.cpp
+++ b/src/llama-kv-cache.cpp
@@ -312,10 +312,14 @@ llama_kv_cache::llama_kv_cache(
n_rot_k = other->n_rot_k;
n_rot_v = other->n_rot_v;
} else {
- const char * LLAMA_ATTN_ROT_DISABLE = getenv("LLAMA_ATTN_ROT_DISABLE");
- const bool attn_rot_disable = LLAMA_ATTN_ROT_DISABLE ? atoi(LLAMA_ATTN_ROT_DISABLE) : false;
+ // two knobs turn the same thing off: the environment variable, for the whole process, and
+ // the model setting, for the contexts of one model - see llama_model_set_attn_rot
+ const char * env_rot_disable = getenv("LLAMA_ATTN_ROT_DISABLE");
+
+ const bool attn_rot_disable = !model.attn_rot || (env_rot_disable != nullptr && atoi(env_rot_disable) != 0);
if (attn_rot_disable) {
- LLAMA_LOG_WARN("%s: attention rotation force disabled (LLAMA_ATTN_ROT_DISABLE)\n", __func__);
+ LLAMA_LOG_WARN("%s: attention rotation force disabled (%s)\n", __func__,
+ model.attn_rot ? "LLAMA_ATTN_ROT_DISABLE" : "llama_model_set_attn_rot");
}
bool attn_rot_k =
diff --git a/src/llama-memory.cpp b/src/llama-memory.cpp
index ca6844c32..961dd74ef 100644
--- a/src/llama-memory.cpp
+++ b/src/llama-memory.cpp
@@ -1,5 +1,11 @@
#include "llama-memory.h"
+#include "llama-ext.h"
+#include "llama-io.h"
+
+#include <set>
+#include <vector>
+
llama_memory_status llama_memory_status_combine(llama_memory_status s0, llama_memory_status s1) {
bool has_update = false;
@@ -57,3 +63,30 @@ bool llama_memory_status_is_fail(llama_memory_status status) {
return false;
}
+
+// collects the buffers of the tensors that a state write visits, without keeping the data
+struct llama_memory_buffers_writer : public 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;
+ }
+};
+
+std::vector<ggml_backend_buffer_t> llama_memory_get_buffers(const llama_memory_t mem) {
+ llama_memory_buffers_writer writer;
+
+ mem->state_write(writer);
+
+ return std::vector<ggml_backend_buffer_t>(writer.buffers.begin(), writer.buffers.end());
+}
diff --git a/src/llama-model.cpp b/src/llama-model.cpp
index 6872589fa..e57ede5fd 100644
--- a/src/llama-model.cpp
+++ b/src/llama-model.cpp
@@ -3781,6 +3781,10 @@ ggml_backend_dev_t llama_model_get_device(const struct llama_model * model, int
return model->devices[i].dev;
}
+void llama_model_set_attn_rot(struct llama_model * model, bool value) {
+ model->attn_rot = value;
+}
+
//
// llama_model_base
//
diff --git a/src/llama-model.h b/src/llama-model.h
index dff84c574..bc640e0dd 100644
--- a/src/llama-model.h
+++ b/src/llama-model.h
@@ -690,6 +690,9 @@ struct llama_model {
llama_hadamard hdmd;
+ // see llama_model_set_attn_rot
+ bool attn_rot = true;
+
// for classifier models
std::vector<std::string> classifier_labels;
diff --git a/tests/.gitignore b/tests/.gitignore
index 04095c9dd..6a6373101 100644
--- a/tests/.gitignore
+++ b/tests/.gitignore
@@ -2,6 +2,7 @@
!*.*
!snapshots/
!fusion/
+!llama-context/
*.o
ggml-common.h
**/*.swp
diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
index daccea063..1a9bc9157 100644
--- a/tests/CMakeLists.txt
+++ b/tests/CMakeLists.txt
@@ -166,8 +166,17 @@ if (LLAMA_LLGUIDANCE)
llama_build_and_test(test-grammar-llguidance.cpp ARGS ${PROJECT_SOURCE_DIR}/models/ggml-vocab-llama-bpe.gguf)
endif ()
-llama_build(test-recurrent-state-rollback.cpp)
-llama_build(test-save-load-state.cpp)
+# llama_context state tests: KV and recurrent state save, load, copy, rollback
+set(LLAMA_CONTEXT_TEST_SRC
+ test-llama-context.cpp
+ llama-context/cases-kv.cpp
+ llama-context/cases-rs.cpp
+ llama-context/test-case.cpp
+ llama-context/test-log.cpp
+ llama-context/test-runner.cpp
+ llama-context/test-util.cpp
+)
+llama_build(${LLAMA_CONTEXT_TEST_SRC})
if (NOT WIN32 OR NOT BUILD_SHARED_LIBS)
# these tests are disabled on Windows because they use internal functions not exported with LLAMA_API (when building with shared libraries)
@@ -213,6 +222,8 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS)
# llama_build_and_test(test-double-float.cpp) # SLOW
+ # TODO: after `test-llama-archs` is refactored to not link internal symbols, we can move all related tests out of
+ # this branch and run them on all platforms
llama_build(test-llama-archs.cpp)
set(MODEL_DIR "${CMAKE_CURRENT_BINARY_DIR}/test-models/")
@@ -228,26 +239,18 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS)
FIXTURES_SETUP generate-models
)
- # Test recurrent-state rollback across all architectures, using the generated dummy models
+ # Test context state handling across all architectures, using the generated dummy models
llama_test(
- test-recurrent-state-rollback
+ test-llama-context
LABEL main
ARGS --models "${MODEL_DIR}"
)
# abort on graph reallocations that the initial reserve should have covered
- set_tests_properties(test-recurrent-state-rollback PROPERTIES
+ set_tests_properties(test-llama-context PROPERTIES
FIXTURES_REQUIRED generate-models
ENVIRONMENT "GGML_SCHED_DEBUG_REALLOC=1"
)
- # Test state save/load functionality across all architectures, using the generated dummy models
- llama_test(
- test-save-load-state
- LABEL main
- ARGS --models "${MODEL_DIR}"
- )
- set_tests_properties(test-save-load-state PROPERTIES FIXTURES_REQUIRED generate-models)
-
llama_build(test-fusion.cpp)
endif()
@@ -288,6 +291,10 @@ add_test(NAME test-download-model COMMAND ${CMAKE_COMMAND}
)
set_tests_properties(test-download-model PROPERTIES FIXTURES_SETUP test-download-model)
+# the context state cases against a real model, the sweep above uses the generated dummy ones
+llama_test(test-llama-context NAME test-llama-context-model LABEL "model" ARGS -m "${MODEL_DEST}")
+set_tests_properties(test-llama-context-model PROPERTIES FIXTURES_REQUIRED test-download-model)
+
llama_build_and_test(test-thread-safety.cpp ARGS -m "${MODEL_DEST}" -ngl 99 -p "The meaning of life is" -n 128 -c 256 -ub 32 -np 4 -t 2)
set_tests_properties(test-thread-safety PROPERTIES FIXTURES_REQUIRED test-download-model)
@@ -313,12 +320,6 @@ llama_build_and_test(test-model-load-cancel.cpp LABEL "model")
llama_build_and_test(test-autorelease.cpp LABEL "model")
llama_build_and_test(test-backend-sampler.cpp LABEL "model")
-# Test for state restore with fragmented KV cache
-# Requires a model, uses same args pattern as test-thread-safety
-# TODO: run on all dummy models
-llama_build_and_test(test-state-restore-fragmented.cpp LABEL "model" ARGS -m "${MODEL_DEST}")
-set_tests_properties(test-state-restore-fragmented PROPERTIES FIXTURES_REQUIRED test-download-model)
-
if (APPLE)
llama_build(test-rset-release.cpp)
endif()
diff --git a/tests/llama-context/cases-kv.cpp b/tests/llama-context/cases-kv.cpp
new file mode 100644
index 000000000..152977468
--- /dev/null
+++ b/tests/llama-context/cases-kv.cpp
@@ -0,0 +1,694 @@
+// KV cache state cases
+//
+#include "test-case.h"
+
+#include "test-log.h"
+#include "test-util.h"
+
+#include "common.h"
+
+#include "../../src/llama-ext.h" // llama_model_set_attn_rot
+
+#include <algorithm>
+#include <cstdio>
+#include <cstring>
+#include <fstream>
+#include <functional>
+#include <iterator>
+#include <string>
+#include <utility>
+#include <vector>
+
+//
+// reference run
+//
+
+// - decode all but the last token, saving the state to disk
+// - decode the last token
+// - generate n_predict tokens
+// the logits it produces are the reference of every case with needs_baseline()
+struct case_baseline : ctx_test {
+ const char * name() const override { return "baseline"; }
+
+ const char * description() const override {
+ return "reference generation over the whole prompt, saves the state before the last token";
+ }
+
+ case_result run(model_run & mr) const override {
+ ctx_opts opts;
+ opts.n_seq_max = 2;
+
+ auto ctx = mr.ctx(opts);
+ auto smpl = mr.sampler();
+ if (!ctx) {
+ return case_result::fail();
+ }
+
+ auto n_past = 0;
+ if (!common_prompt_batch_decode(ctx.get(), mr.tokens, (int) mr.tokens.size(), n_past, mr.params.n_batch, mr.params.out_file, true)) {
+ TLOG_ERR("%s: failed to decode the prompt\n", __func__);
+ return case_result::fail();
+ }
+
+ generation_result result;
+ if (!generate_tokens(ctx.get(), smpl.get(), n_past, mr.params.n_predict, 0, result)) {
+ return case_result::fail();
+ }
+ if (result.empty()) {
+ return case_result::fail();
+ }
+
+ mr.baseline = std::move(result);
+ mr.have_baseline = true;
+
+ return case_result::pass();
+ }
+};
+
+//
+// sequences
+//
+
+// - decode the same prefix into two sequences
+// - remove sequence 0
+// - verify that sequence 1 remains unchanged
+struct case_seq_rm_isolated : ctx_test {
+ const char * name() const override { return "seq_rm"; }
+
+ const char * description() const override {
+ return "removing one sequence leaves the state of the other sequences untouched";
+ }
+
+ case_result run(model_run & mr) const override {
+ ctx_opts opts;
+ opts.n_ctx = 256;
+ opts.n_seq_max = 2;
+ opts.kv_unified = 1;
+
+ auto ctx = mr.ctx(opts);
+ if (!ctx) {
+ return case_result::fail();
+ }
+
+ const size_t n_tokens = std::min<size_t>(mr.tokens.size(), 128);
+ for (llama_seq_id seq_id = 0; seq_id < 2; ++seq_id) {
+ if (!decode_seq(ctx.get(), mr.tokens, 0, (llama_pos) n_tokens, seq_id, true)) {
+ TLOG_ERR("%s: failed to decode the prompt for sequence %d\n", __func__, seq_id);
+ return case_result::fail();
+ }
+ }
+
+ std::vector<uint8_t> before;
+ if (!get_seq_state(ctx.get(), 1, LLAMA_STATE_SEQ_FLAGS_NONE, before)) {
+ return case_result::fail();
+ }
+
+ if (!seq_rm(ctx.get(), 0)) {
+ return case_result::fail();
+ }
+
+ std::vector<uint8_t> after;
+ if (!get_seq_state(ctx.get(), 1, LLAMA_STATE_SEQ_FLAGS_NONE, after)) {
+ return case_result::fail();
+ }
+
+ CTX_REQUIRE(blobs_equal(before, after), "removing seq 0 changed the state of seq 1");
+
+ return case_result::pass();
+ }
+};
+
+// - load the state saved by the baseline case
+// - replay its last token
+// - generate again and compare against the baseline logits
+struct case_state_load : ctx_test {
+ const char * name() const override { return "state_load"; }
+
+ const char * description() const override {
+ return "a saved session replays to the same logits as the run that saved it";
+ }
+
+ bool needs_baseline() const override { return true; }
+
+ case_result run(model_run & mr) const override {
+ ctx_opts opts;
+ opts.n_seq_max = 2;
+
+ auto ctx = mr.ctx(opts);
+ auto smpl = mr.sampler();
+ if (!ctx) {
+ return case_result::fail();
+ }
+
+ int n_past = 0;
+ if (!load_state_replay_last(ctx.get(), mr.params.out_file, mr.tokens, n_past)) {
+ return case_result::fail();
+ }
+
+ if (!generate_tokens_compare(ctx.get(), smpl.get(), n_past, mr.params.n_predict, 0, mr.baseline, NMSE_THRESHOLD)) {
+ return case_result::fail();
+ }
+
+ return case_result::pass();
+ }
+};
+
+// - load the state saved by the baseline case and replay its last token
+// - copy sequence 0 to sequence 1 over the given io path, clearing the cache in between
+// - continue on sequence 1 and compare against the baseline logits
+struct case_seq_cp : ctx_test {
+ explicit case_seq_cp(bool on_device_) : on_device(on_device_) {}
+
+ const char * name() const override { return on_device ? "cp_d" : "cp_h"; }
+
+ const char * description() const override {
+ return on_device ? "copying a sequence over the device io path keeps its state"
+ : "copying a sequence over the host io path keeps its state";
+ }
+
+ bool needs_baseline() const override { return true; }
+
+ case_result run(model_run & mr) const override {
+ ctx_opts opts;
+ opts.n_seq_max = 2;
+
+ auto ctx = mr.ctx(opts);
+ auto smpl = mr.sampler();
+ if (!ctx) {
+ return case_result::fail();
+ }
+
+ const uint32_t flags = on_device ? LLAMA_STATE_SEQ_FLAGS_ON_DEVICE : LLAMA_STATE_SEQ_FLAGS_NONE;
+ TLOG_INF("io path: %s\n", on_device ? "device" : "host");
+
+ int n_past = 0;
+ if (!load_state_replay_last(ctx.get(), mr.params.out_file, mr.tokens, n_past)) {
+ return case_result::fail();
+ }
+
+ {
+ std::vector<uint8_t> seq_store;
+ if (!get_seq_state(ctx.get(), 0, flags, seq_store)) {
+ return case_result::fail();
+ }
+ TLOG_TRC("%s: seq 0 copied, %zu bytes\n", __func__, seq_store.size());
+
+ llama_memory_clear(llama_get_memory(ctx.get()), true);
+ TLOG_TRC("%s: kv cache cleared\n", __func__);
+
+ if (!set_seq_state(ctx.get(), 1, flags, seq_store)) {
+ return case_result::fail();
+ }
+ TLOG_TRC("%s: seq 1 restored, %zu bytes\n", __func__, seq_store.size());
+ }
+
+ if (!generate_tokens_compare(ctx.get(), smpl.get(), n_past, mr.params.n_predict, 1, mr.baseline, NMSE_THRESHOLD)) {
+ return case_result::fail();
+ }
+
+ return case_result::pass();
+ }
+
+private:
+ bool on_device;
+};
+
+// - decode the same prefix on two sequences, interleaving seq 0 cells between the seq 1 cells
+// - save the seq 1 state, free the interleaved seq 0 cells, and restore via the given io path
+// - the restore destination is non-contiguous: scatter reads are batched per contiguous run
+// - save again on the host and compare the two blobs byte for byte
+struct case_seq_cp_scatter : ctx_test {
+ explicit case_seq_cp_scatter(bool on_device_) : on_device(on_device_) {}
+
+ const char * name() const override { return on_device ? "cp_d_s" : "cp_h_s"; }
+
+ const char * description() const override {
+ return "saving and restoring a sequence with a fragmented (scattered) cell range";
+ }
+
+ case_result run(model_run & mr) const override {
+ ctx_opts opts;
+ opts.n_ctx = 256;
+ opts.n_seq_max = 2;
+ opts.kv_unified = 1;
+
+ auto ctx = mr.ctx(opts);
+ if (!ctx) {
+ return case_result::fail();
+ }
+
+ const uint32_t flags = on_device ? LLAMA_STATE_SEQ_FLAGS_ON_DEVICE : LLAMA_STATE_SEQ_FLAGS_NONE;
+ TLOG_INF("io path: %s\n", on_device ? "device" : "host");
+
+ // seq 0 cells 0,1,4 interleave the seq 1 cells 2,3,5
+ if (!decode_one(ctx.get(), mr.tokens[0], 0, 0) ||
+ !decode_one(ctx.get(), mr.tokens[1], 1, 0) ||
+ !decode_one(ctx.get(), mr.tokens[0], 0, 1) ||
+ !decode_one(ctx.get(), mr.tokens[1], 1, 1) ||
+ !decode_one(ctx.get(), mr.tokens[2], 2, 0) ||
+ !decode_one(ctx.get(), mr.tokens[2], 2, 1)) {
+ TLOG_ERR("%s: failed to build the interleaved state\n", __func__);
+ return case_result::fail();
+ }
+
+ // host blob: contains the KV data, used for the byte-for-byte comparison
+ std::vector<uint8_t> state_before;
+ if (!get_seq_state(ctx.get(), 1, LLAMA_STATE_SEQ_FLAGS_NONE, state_before)) {
+ return case_result::fail();
+ }
+
+ // save via the io path under test
+ std::vector<uint8_t> state_save;
+ if (!get_seq_state(ctx.get(), 1, flags, state_save)) {
+ return case_result::fail();
+ }
+ TLOG_TRC("%s: seq 1 saved via %s, %zu bytes\n", __func__, on_device ? "device" : "host", state_save.size());
+
+ // free seq 0's cells so the ring is fragmented: the restore destination (seq 1's interleaved cells) stays non-contiguous
+ if (!seq_rm(ctx.get(), 0)) {
+ return case_result::fail();
+ }
+
+ // restore via the io path under test
+ if (!set_seq_state(ctx.get(), 1, flags, state_save)) {
+ return case_result::fail();
+ }
+ TLOG_TRC("%s: seq 1 restored via %s, %zu bytes\n", __func__, on_device ? "device" : "host", state_save.size());
+
+ std::vector<uint8_t> state_after;
+ if (!get_seq_state(ctx.get(), 1, LLAMA_STATE_SEQ_FLAGS_NONE, state_after)) {
+ return case_result::fail();
+ }
+
+ // the blob is serialized in sequence cell order, so identical bytes iff the restore wrote the same KV
+ CTX_REQUIRE(blobs_equal(state_before, state_after), "the restored KV state is not byte-identical to the saved state");
+
+ return case_result::pass();
+ }
+
+private:
+ bool on_device;
+};
+
+//
+// sessions
+//
+
+// - save a sequence, erase it, restore it, save it again
+// - the two blobs must be identical
+// compares blobs rather than generated text: a partially restored cell still decodes to plausible tokens
+struct case_state_roundtrip : ctx_test {
+ const char * name() const override { return "rt"; }
+
+ const char * description() const override {
+ return "save, erase, restore, save again - the two blobs must be identical";
+ }
+
+ case_result run(model_run & mr) const override {
+ auto ctx = mr.ctx();
+ if (!ctx) {
+ return case_result::fail();
+ }
+
+ common_batch batch = common_batch_get_one(ctx.get(), mr.tokens);
+ if (llama_process(ctx.get(), LLAMA_PROCESS_TYPE_DECODE, batch.get())) {
+ TLOG_ERR("%s: failed to decode the prompt\n", __func__);
+ return case_result::fail();
+ }
+
+ std::vector<uint8_t> blob_a;
+ if (!get_seq_state(ctx.get(), 0, LLAMA_STATE_SEQ_FLAGS_NONE, blob_a)) {
+ return case_result::fail();
+ }
+
+ if (!seq_rm(ctx.get(), 0)) {
+ return case_result::fail();
+ }
+
+ if (!set_seq_state(ctx.get(), 0, LLAMA_STATE_SEQ_FLAGS_NONE, blob_a)) {
+ TLOG_ERR("%s: failed to restore seq 0\n", __func__);
+ return case_result::fail();
+ }
+
+ std::vector<uint8_t> blob_b;
+ if (!get_seq_state(ctx.get(), 0, LLAMA_STATE_SEQ_FLAGS_NONE, blob_b)) {
+ return case_result::fail();
+ }
+
+ CTX_REQUIRE(blobs_equal(blob_a, blob_b), "the state changed across a save / erase / restore cycle");
+
+ return case_result::pass();
+ }
+};
+
+// a failed restore must leave the sequence empty and must not change the logits of other sequences
+struct case_state_restore_failure : ctx_test {
+ const char * name() const override { return "rf"; }
+
+ const char * description() const override {
+ return "restoring a corrupted state fails and leaves the other sequences untouched";
+ }
+
+ case_result run(model_run & mr) const override {
+ ctx_opts opts;
+ opts.n_ctx = 256;
+ opts.n_seq_max = 4;
+ opts.kv_unified = 1;
+
+ // without flash attention, corrupted data left behind by the restore shows up as NaN logits on the other sequences
+ opts.fa = LLAMA_FLASH_ATTN_TYPE_DISABLED;
+
+ auto ctx = mr.ctx(opts);
+ if (!ctx) {
+ return case_result::fail();
+ }
+
+ llama_memory_t mem = llama_get_memory(ctx.get());
+ if (mem == nullptr) {
+ TLOG_INF("no memory to test\n");
+ return case_result::pass();
+ }
+
+ const auto decode = [&](const llama_tokens & inp, llama_seq_id seq_id, std::vector<float> * logits_out) {
+ if (!decode_tokens(ctx.get(), inp, seq_id)) {
+ TLOG_ERR("%s: failed to decode on sequence %d\n", __func__, seq_id);
+ return false;
+ }
+
+ if (logits_out && !get_current_logits(ctx.get(), *logits_out)) {
+ TLOG_ERR("%s: failed to get logits\n", __func__);
+ return false;
+ }
+
+ return true;
+ };
+
+ const llama_tokens tokens_save (mr.tokens.begin(), mr.tokens.begin() + std::min<size_t>(24, mr.tokens.size()));
+ const llama_tokens tokens_verify(mr.tokens.end() - std::min<size_t>(8, mr.tokens.size()), mr.tokens.end());
+
+ // the cases share a working directory, so the state file is named after the model
+ const std::string path = "state-restore-failure." + mr.name + ".tmp.bin";
+
+ llama_memory_clear(mem, true);
+
+ std::vector<float> baseline;
+ if (!decode(tokens_verify, 1, &baseline)) {
+ return case_result::fail();
+ }
+
+ const std::vector<std::pair<const char *, std::function<bool()>>> sub_cases = {
+ { "buffer", [&]() {
+ std::vector<uint8_t> state(llama_state_seq_get_size(ctx.get(), 0));
+ GGML_ASSERT(llama_state_seq_get_data(ctx.get(), state.data(), state.size(), 0) == state.size());
+ llama_memory_seq_rm(mem, 0, -1, -1);
+
+ if (!corrupt_state(state)) {
+ return false;
+ }
+
+ return llama_state_seq_set_data(ctx.get(), state.data(), state.size(), 0) == 0;
+ }},
+ { "file", [&]() {
+ GGML_ASSERT(llama_state_seq_save_file(ctx.get(), path.c_str(), 0, tokens_save.data(), tokens_save.size()) > 0);
+ llama_memory_seq_rm(mem, 0, -1, -1);
+
+ std::vector<uint8_t> data;
+ {
+ std::ifstream f(path, std::ios::binary);
+ data.assign(std::istreambuf_iterator<char>(f), std::istreambuf_iterator<char>());
+ }
+
+ if (!corrupt_state(data)) {
+ std::remove(path.c_str());
+ return false;
+ }
+
+ {
+ std::ofstream f(path, std::ios::binary);
+ f.write((const char *) data.data(), data.size());
+ }
+
+ llama_tokens tokens_out(tokens_save.size());
+ size_t n_token_count = 0;
+ const size_t nread = llama_state_seq_load_file(ctx.get(), path.c_str(), 0, tokens_out.data(), tokens_out.size(), &n_token_count);
+ std::remove(path.c_str());
+
+ return nread == 0;
+ }},
+ };
+
+ for (const auto & [name, restore_failed] : sub_cases) {
+ llama_memory_clear(mem, true);
+
+ if (!decode(tokens_save, 0, nullptr)) {
+ return case_result::fail();
+ }
+
+ CTX_REQUIRE(restore_failed(), std::string(name) + ": restoring a corrupted state did not fail");
+
+ CTX_REQUIRE(llama_memory_seq_pos_max(mem, 0) == -1,
+ std::string(name) + ": the sequence is not empty after a failed restore");
+
+ std::vector<float> logits;
+ if (!decode(tokens_verify, 1, &logits)) {
+ return case_result::fail();
+ }
+
+ const logits_diff diff = compare_logits(baseline.data(), logits.data(), logits.size());
+
+ CTX_REQUIRE(diff.n_nan == 0 && diff.diff_max <= 1e-6f,
+ std::string(name) + ": the logits changed after a failed restore (max diff = " +
+ std::to_string(diff.diff_max) + ", nan = " + std::to_string(diff.n_nan) + ")");
+
+ TLOG_TRC("%s: %s: logits match (max diff = %g)\n", __func__, name, (double) diff.diff_max);
+ }
+
+ return case_result::pass();
+ }
+};
+
+// a KV state saved with attention rotation enabled must restore only into a context with the
+// same setting
+// note: rotation is only active for quantized KV caches with a head size that is a multiple of
+// 64, for other models the restore into the rotation-disabled context is valid and the
+// test passes vacuously
+struct case_state_rotation : ctx_test {
+ const char * name() const override { return "rot"; }
+
+ const char * description() const override {
+ return "a state saved with attention rotation restores only into a matching context";
+ }
+
+ case_result run(model_run & mr) const override {
+ const auto make_context = [&](ggml_type type_k, ggml_type type_v, bool disable_rotation) {
+ // a setting of the model, not of the process - a sweep that runs several models at a
+ // time must not let this leak into the contexts of another model
+ llama_model_set_attn_rot(mr.model, !disable_rotation);
+
+ ctx_opts opts;
+ opts.n_ctx = 32;
+ opts.n_batch = 1;
+ opts.n_ubatch = 1;
+ opts.type_k = type_k;
+ opts.type_v = type_v;
+ opts.fa = LLAMA_FLASH_ATTN_TYPE_ENABLED;
+
+ // probing whether this KV type is supported is expected to fail for some models
+ opts.quiet = true;
+
+ return mr.ctx(opts);
+ };
+
+ std::vector<std::pair<ggml_type, ggml_type>> type_pairs;
+ for (const auto & types : { std::pair{GGML_TYPE_Q8_0, GGML_TYPE_Q8_0} }) {
+ if (make_context(types.first, types.second, false)) {
+ type_pairs.push_back(types);
+ }
+ }
+ if (type_pairs.empty()) {
+ return case_result::skip("no supported quantized KV cache type combination");
+ }
+
+ bool success = true;
+ for (const auto & types : type_pairs) {
+ auto src = make_context(types.first, types.second, false);
+ if (!src) {
+ TLOG_ERR("%s: failed to create the source context\n", __func__);
+ success = false;
+ break;
+ }
+
+ llama_token token = 0;
+ if (llama_decode(src.get(), llama_batch_get_one(&token, 1))) {
+ TLOG_ERR("%s: failed to decode a token\n", __func__);
+ success = false;
+ break;
+ }
+
+ std::vector<uint8_t> state;
+ if (!get_seq_state(src.get(), 0, LLAMA_STATE_SEQ_FLAGS_NONE, state)) {
+ if (state.empty()) {
+ continue; // no KV state to test
+ }
+ success = false;
+ break;
+ }
+
+ auto matching = make_context(types.first, types.second, false);
+ if (!matching || !set_seq_state(matching.get(), 0, LLAMA_STATE_SEQ_FLAGS_NONE, state)) {
+ TLOG_ERR("%s: failed to restore a matching rotation\n", __func__);
+ success = false;
+ break;
+ }
+
+ auto mismatched = make_context(types.first, types.second, true);
+ if (!mismatched) {
+ TLOG_ERR("%s: failed to create a mismatched rotation context\n", __func__);
+ success = false;
+ break;
+ }
+ // a model that does not use attention rotation accepts the mismatched restore
+ if (set_seq_state(mismatched.get(), 0, LLAMA_STATE_SEQ_FLAGS_NONE, state)) {
+ TLOG_TRC("%s: state restored into rotation-disabled context, model does not use attention rotation\n", __func__);
+ }
+ }
+
+ llama_model_set_attn_rot(mr.model, true); // the default, the cases that follow rely on it
+
+ if (!success) {
+ return case_result::fail();
+ }
+
+ return case_result::pass();
+ }
+};
+
+// restore a sequence into a fragmented cache: the slots freed by the removed sequence are
+// scattered between live cells of the other sequences, so the restore cannot claim one
+// contiguous block
+// the restored state must come back byte for byte, and the sequence must continue to behave
+// exactly as the same sequence in a context where it was never removed
+// ref: https://github.com/ggml-org/llama.cpp/issues/17527
+struct case_state_restore_fragmented : ctx_test {
+ const char * name() const override { return "frag"; }
+
+ const char * description() const override {
+ return "restore a sequence into a cache fragmented by the removal of another sequence";
+ }
+
+ std::string skip_reason(model_run & mr) const override {
+ if (arch_reserves_single_seq(mr.model)) {
+ return "the interleaved batch is a multi-seq graph";
+ }
+
+ return ctx_test::skip_reason(mr);
+ }
+
+ case_result run(model_run & mr) const override {
+ const int n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(mr.model));
+
+ ctx_opts opts;
+ opts.n_ctx = 256;
+ opts.n_seq_max = 3;
+ opts.kv_unified = 1;
+ opts.n_batch = 3*70;
+ opts.n_ubatch = 3*70; // the interleaved prompt is a single multi-seq batch
+
+ // ctx removes and restores seq 1, ctx_keep keeps it untouched as the reference
+ llama_context_ptr ctx = mr.ctx(opts);
+ llama_context_ptr ctx_keep = mr.ctx(opts);
+ if (!ctx || !ctx_keep) {
+ return case_result::fail();
+ }
+
+ const llama_pos n_tokens = (llama_pos) std::min<size_t>(70, mr.tokens.size());
+
+ // interleave the 3 sequences: 01201230123... - seq 1 owns every third cell
+ const auto decode_interleaved = [&](llama_context * cur) {
+ common_batch batch(cur);
+ for (llama_pos i = 0; i < n_tokens; i++) {
+ for (llama_seq_id s = 0; s < 3; ++s) {
+ batch.add(mr.tokens[i], i, s, false);
+ }
+ }
+ batch.set_output(batch.size() - 1, true);
+
+ return llama_process(cur, LLAMA_PROCESS_TYPE_DECODE, batch.get()) == 0;
+ };
+
+ if (!decode_interleaved(ctx.get()) || !decode_interleaved(ctx_keep.get())) {
+ TLOG_ERR("%s: failed to decode the interleaved prompt\n", __func__);
+ return case_result::fail();
+ }
+
+ TLOG_INF("%s: processed the prompt on seq 0, 1, 2 (%d tokens each)\n", __func__, n_tokens);
+
+ std::vector<uint8_t> state_before;
+ if (!get_seq_state(ctx.get(), 1, LLAMA_STATE_SEQ_FLAGS_NONE, state_before)) {
+ return case_result::fail();
+ }
+
+ // clearing seq 1 leaves holes where its cells were - no contiguous block large enough
+ // for the seq 1 state is left in the cache
+ llama_memory_seq_rm(llama_get_memory(ctx.get()), 1, -1, -1);
+
+ if (!set_seq_state(ctx.get(), 1, LLAMA_STATE_SEQ_FLAGS_NONE, state_before)) {
+ TLOG_ERR("%s: failed to restore the seq state into a fragmented cache\n", __func__);
+ return case_result::fail();
+ }
+
+ TLOG_INF("%s: restored seq 1 into a fragmented cache, %zu bytes\n", __func__, state_before.size());
+
+ // the restore has to be lossless - the cells the sequence ends up in are irrelevant, its
+ // state has to come back exactly as it was saved
+ std::vector<uint8_t> state_after;
+ if (!get_seq_state(ctx.get(), 1, LLAMA_STATE_SEQ_FLAGS_NONE, state_after)) {
+ return case_result::fail();
+ }
+
+ CTX_REQUIRE(blobs_equal(state_before, state_after), "restoring into a fragmented cache changed the state of seq 1");
+
+ // the restored sequence must continue exactly as the untouched one - decode the same token
+ // on seq 1 in both contexts and compare the logits they produce
+ const llama_token token = gen_token(n_vocab, 1, n_tokens);
+
+ if (!decode_one(ctx.get(), token, n_tokens, 1) || !decode_one(ctx_keep.get(), token, n_tokens, 1)) {
+ TLOG_ERR("%s: failed to decode with the restored state\n", __func__);
+ return case_result::fail();
+ }
+
+ const float * logits = llama_get_logits(ctx.get());
+ const float * logits_keep = llama_get_logits(ctx_keep.get());
+ if (!logits || !logits_keep) {
+ TLOG_ERR("%s: missing logits after the restore\n", __func__);
+ return case_result::fail();
+ }
+
+ const logits_diff diff = compare_logits(logits, logits_keep, n_vocab);
+
+ CTX_REQUIRE(diff.nmse() <= NMSE_THRESHOLD,
+ "the logits of the restored seq 1 differ from the reference: nmse " + std::to_string(diff.nmse()));
+
+ TLOG_INF("%s: restored seq 1 continued identically to the reference, nmse %g\n", __func__, diff.nmse());
+
+ return case_result::pass();
+ }
+};
+
+std::vector<ctx_test_ptr> make_kv_cases() {
+ std::vector<ctx_test_ptr> cases;
+
+ cases.push_back(std::make_unique<case_baseline>());
+ cases.push_back(std::make_unique<case_seq_rm_isolated>());
+ cases.push_back(std::make_unique<case_state_load>());
+ cases.push_back(std::make_unique<case_seq_cp>(false));
+ cases.push_back(std::make_unique<case_seq_cp>(true));
+ cases.push_back(std::make_unique<case_seq_cp_scatter>(false));
+ cases.push_back(std::make_unique<case_seq_cp_scatter>(true));
+ cases.push_back(std::make_unique<case_state_roundtrip>());
+ cases.push_back(std::make_unique<case_state_restore_failure>());
+ cases.push_back(std::make_unique<case_state_rotation>());
+ cases.push_back(std::make_unique<case_state_restore_fragmented>());
+
+ return cases;
+}
diff --git a/tests/llama-context/cases-rs.cpp b/tests/llama-context/cases-rs.cpp
new file mode 100644
index 000000000..b549bab4f
--- /dev/null
+++ b/tests/llama-context/cases-rs.cpp
@@ -0,0 +1,404 @@
+// recurrent state cases
+
+#include "test-case.h"
+
+#include "test-log.h"
+#include "test-util.h"
+
+#include "common.h"
+
+#include <algorithm>
+#include <cmath>
+#include <cstdio>
+#include <string>
+#include <vector>
+
+// 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.
+struct case_multi_seq_split_replay : ctx_test {
+ test_group group() const override { return TEST_GROUP_RS; }
+
+ const char * name() const override { return "replay"; }
+
+ const char * description() const override {
+ return "replay rolled-back sequences in a batch wider than n_ubatch, against a reference context";
+ }
+
+ case_result run(model_run & mr) const override {
+ const int n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(mr.model));
+
+ constexpr uint32_t n_seqs = 2;
+ constexpr uint32_t n_ubatch = 16;
+ constexpr uint32_t n_prompt = 20;
+ constexpr uint32_t n_rollback = 3;
+ constexpr uint32_t n_replay = 40; // > n_ubatch so each seq spans multiple ubatches
+ constexpr llama_pos p0 = n_prompt - n_rollback;
+
+ // both contexts have the same layout, only what they decode differs
+ const auto make_ctx_multi = [&]() {
+ ctx_opts opts;
+ opts.n_ctx = 256;
+ opts.n_batch = 256;
+ opts.n_ubatch = n_ubatch;
+ opts.n_seq_max = n_seqs;
+ opts.n_rs_seq = 8;
+ opts.kv_unified = 0;
+ return mr.ctx(opts);
+ };
+
+ llama_context_ptr ctx_roll = make_ctx_multi();
+ llama_context_ptr ctx_ref = make_ctx_multi();
+ if (!ctx_roll || !ctx_ref) {
+ return case_result::fail();
+ }
+
+ if (llama_n_rs_seq(ctx_roll.get()) < n_rollback) {
+ return case_result::skip("n_rs_seq is too small");
+ }
+
+ const auto tok = [&](uint32_t seq, llama_pos pos) {
+ return gen_token(n_vocab, (llama_seq_id) seq, pos);
+ };
+
+ bool ok = true;
+
+ // keep the anchor at p0 - 1 when rolling back the tail
+ for (uint32_t s = 0; s < n_seqs && ok; ++s) {
+ ok = ok && decode_gen(ctx_roll.get(), n_vocab, (llama_seq_id) s, 0, (llama_pos) p0 - 1);
+ ok = ok && decode_gen(ctx_ref.get(), n_vocab, (llama_seq_id) s, 0, (llama_pos) p0 - 1);
+
+ ok = ok && decode_gen(ctx_roll.get(), n_vocab, (llama_seq_id) s, (llama_pos) p0 - 1, (llama_pos) n_prompt);
+ ok = ok && decode_gen(ctx_ref.get(), n_vocab, (llama_seq_id) s, (llama_pos) p0 - 1, (llama_pos) p0);
+
+ ok = ok && llama_memory_seq_rm(llama_get_memory(ctx_roll.get()), (llama_seq_id) s, p0, -1);
+
+ // a second partial removal while one is pending must be refused
+ ok = ok && !llama_memory_seq_rm(llama_get_memory(ctx_roll.get()), (llama_seq_id) s, p0 - 1, -1);
+ }
+ CTX_REQUIRE(ok, "the multi-seq prefill / rollback failed");
+
+ // all seqs replay in a single batch
+ const auto decode_replay = [&](llama_context * ctx) {
+ common_batch batch(ctx);
+ for (uint32_t s = 0; s < n_seqs; ++s) {
+ for (uint32_t i = 0; i < n_replay; ++i) {
+ const llama_pos pos = p0 + (llama_pos) i;
+ batch.add(tok(s, pos), pos, (llama_seq_id) s, true);
+ }
+ }
+ return llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get()) == 0;
+ };
+
+ ok = decode_replay(ctx_roll.get()) && decode_replay(ctx_ref.get());
+ CTX_REQUIRE(ok, "the multi-seq replay decode failed");
+
+ // both contexts decode identical batches, so the logits should match;
+ // random dummy models can still drift up to ~1.7e-5, so the bound is 1e-4
+ constexpr float nmse_eps = 1e-4f;
+
+ logits_diff diff;
+ for (uint32_t i = 0; i < n_seqs*n_replay; ++i) {
+ const float * l_roll = llama_get_logits_ith(ctx_roll.get(), i);
+ const float * l_ref = llama_get_logits_ith(ctx_ref.get(), i);
+ if (l_roll == nullptr || l_ref == nullptr) {
+ TLOG_ERR("%s: missing multi-seq logits at index %u\n", __func__, i);
+ return case_result::fail();
+ }
+
+ diff.add(l_roll, l_ref, n_vocab);
+ }
+
+ CTX_REQUIRE(diff.nmse() <= nmse_eps,
+ "multi-seq split replay logits mismatch (max diff " + std::to_string(diff.diff_max) +
+ ", nmse " + std::to_string(diff.nmse()) + ", first difference at index " + std::to_string(diff.i_first) + ")");
+
+ TLOG_INF("%s: multi-seq split replay matched (max diff %g, nmse %g)\n", __func__, (double) diff.diff_max, diff.nmse());
+
+ // seq-1-only decodes must be independent of seq 0's content: diverge seq 0
+ // in ctx_ref only, then compare identical seq-1-only continuations
+ constexpr uint32_t n_tail = 4;
+
+ {
+ common_batch batch_tail(ctx_ref.get());
+ for (uint32_t i = 0; i < n_tail; ++i) {
+ const llama_pos pos = p0 + (llama_pos) (n_replay + i);
+ batch_tail.add(tok(0, pos + 7), pos, 0, false);
+ }
+ ok = llama_process(ctx_ref.get(), LLAMA_PROCESS_TYPE_DECODE, batch_tail.get()) == 0;
+ }
+
+ logits_diff diff_tail;
+ for (uint32_t i = 0; i < n_tail && ok; ++i) {
+ const llama_pos pos = p0 + (llama_pos) (n_replay + i);
+ ok = decode_one(ctx_roll.get(), tok(1, pos), pos, 1) &&
+ decode_one(ctx_ref.get(), tok(1, pos), pos, 1);
+ if (!ok) {
+ break;
+ }
+
+ const float * l_roll = llama_get_logits_ith(ctx_roll.get(), 0);
+ const float * l_ref = llama_get_logits_ith(ctx_ref.get(), 0);
+ if (l_roll == nullptr || l_ref == nullptr) {
+ TLOG_ERR("%s: missing tail logits at step %u\n", __func__, i);
+ return case_result::fail();
+ }
+
+ diff_tail.add(l_roll, l_ref, n_vocab);
+ }
+
+ CTX_REQUIRE(ok && diff_tail.nmse() <= nmse_eps,
+ "a seq-1-only decode leaked the state of seq 0 (nmse " + std::to_string(diff_tail.nmse()) + ")");
+
+ TLOG_INF("%s: seq-1-only decode is independent of seq 0 (max diff %g, nmse %g)\n",
+ __func__, (double) diff_tail.diff_max, diff_tail.nmse());
+
+ return case_result::pass();
+ }
+};
+
+// save a rolled-back single-seq state, restore it into a fresh and into a dirty context, and
+// verify exact logit matches on replay
+struct case_rollback : ctx_test {
+ test_group group() const override { return TEST_GROUP_RS; }
+
+ const char * name() const override { return "rollback"; }
+
+ const char * description() const override {
+ return "a rolled-back recurrent state restores into fresh and dirty contexts with identical replay";
+ }
+
+ case_result run(model_run & mr) const override {
+ const llama_vocab * vocab = llama_model_get_vocab(mr.model);
+ const int n_vocab = llama_vocab_n_tokens(vocab);
+
+ const auto make_rs_ctx = [&]() {
+ ctx_opts opts;
+ opts.n_seq_max = 1;
+ opts.n_rs_seq = 8;
+ return mr.ctx(opts);
+ };
+
+ llama_context_ptr ctx_src = make_rs_ctx();
+ llama_context_ptr ctx_dst = make_rs_ctx();
+ if (!ctx_src || !ctx_dst) {
+ return case_result::fail();
+ }
+
+ const uint32_t n_rs_seq = llama_n_rs_seq(ctx_src.get());
+ if (n_rs_seq == 0) {
+ return case_result::skip("n_rs_seq is disabled");
+ }
+
+ constexpr uint32_t n_rollback = 3;
+ if (n_rs_seq < n_rollback) {
+ return case_result::skip("n_rs_seq is too small");
+ }
+
+ llama_tokens tokens;
+ if (!mr.has_vocab()) {
+ tokens = { 1, 2, 3, 4, 5, 6, 7, 8, 9 };
+ } else {
+ tokens = common_tokenize(ctx_src.get(), "The quick brown fox jumps over the lazy dog", true);
+ }
+ if (tokens.empty()) {
+ TLOG_ERR("%s: not enough prompt tokens\n", __func__);
+ return case_result::fail();
+ }
+ tokens.resize(n_rs_seq + 1, tokens.back());
+
+ const llama_pos rollback_pos = (llama_pos) tokens.size() - n_rollback;
+
+ // Decode the full prompt on the source, then roll back three positions.
+ // Replaying them crosses DSV4's ratio-4 compressor boundary.
+ // Rollback leaves the recurrent memory in a snapshot state (rs_idx != 0).
+ if (!decode_tokens(ctx_src.get(), tokens)) {
+ TLOG_ERR("%s: failed to decode the prompt\n", __func__);
+ return case_result::fail();
+ }
+ if (!llama_memory_seq_rm(llama_get_memory(ctx_src.get()), 0, rollback_pos, -1)) {
+ TLOG_ERR("%s: rollback failed\n", __func__);
+ return case_result::fail();
+ }
+
+ // Save the rolled-back state and restore it into a fresh context.
+ common_prompt_checkpoint ckpt;
+ ckpt.update_tgt(ctx_src.get(), 0, 0);
+ ckpt.load_tgt(ctx_dst.get(), 0, 0);
+
+ // exact match - the same context configuration replaying the same tokens
+ constexpr float nmse_eps = 0.0;
+
+ std::vector<std::vector<float>> logits_src_replay(n_rollback);
+ for (uint32_t i = 0; i < n_rollback; ++i) {
+ const llama_pos pos = rollback_pos + i;
+ if (!decode_one(ctx_src.get(), tokens[pos], pos) ||
+ !decode_one(ctx_dst.get(), tokens[pos], pos)) {
+ TLOG_ERR("%s: replay failed at position %d\n", __func__, pos);
+ return case_result::fail();
+ }
+
+ const float * logits_src = llama_get_logits_ith(ctx_src.get(), 0);
+ const float * logits_dst = llama_get_logits_ith(ctx_dst.get(), 0);
+ if (logits_src == nullptr || logits_dst == nullptr) {
+ TLOG_ERR("%s: missing replay logits at position %d\n", __func__, pos);
+ return case_result::fail();
+ }
+
+ logits_src_replay[i].assign(logits_src, logits_src + n_vocab);
+
+ const logits_diff diff = compare_logits(logits_src, logits_dst, n_vocab);
+ CTX_REQUIRE(diff.nmse() <= nmse_eps,
+ "replay logits mismatch at position " + std::to_string(pos) + ", nmse " + std::to_string(diff.nmse()));
+ }
+
+ // TODO: this test is invalid because RS rollback is only correct once after a ubatch with more than n_rs_seq tokens
+ // this is not the case here. add asserts and guardrails to prevent such attempts
+ //if (!llama_memory_seq_rm(llama_get_memory(ctx_src.get()), 0, rollback_pos, -1) ||
+ // !llama_memory_seq_rm(llama_get_memory(ctx_dst.get()), 0, rollback_pos, -1)) {
+ // TLOG_ERR("%s: partial rollback failed\n", __func__);
+ // return case_result::fail();
+ //}
+
+ //constexpr llama_state_seq_flags partial_flags = LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY;
+ //common_prompt_checkpoint ckpt_partial;
+ //ckpt_partial.update_tgt(ctx_src.get(), 0, partial_flags);
+ //ckpt_partial.load_tgt(ctx_dst.get(), 0, partial_flags);
+
+ // 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_ptr ctx_dirty = make_rs_ctx();
+ if (!ctx_dirty) {
+ return case_result::fail();
+ }
+
+ llama_tokens noise = tokens;
+ for (auto & tok : noise) {
+ tok = (tok + 1) % n_vocab;
+ if (tok < 0) {
+ tok = 0;
+ }
+ }
+ if (!decode_tokens(ctx_dirty.get(), noise)) {
+ TLOG_ERR("%s: failed to decode the dirty prompt\n", __func__);
+ return case_result::fail();
+ }
+ if (!llama_memory_seq_rm(llama_get_memory(ctx_dirty.get()), 0, rollback_pos, -1)) {
+ TLOG_ERR("%s: dirty rollback failed\n", __func__);
+ return case_result::fail();
+ }
+
+ ckpt.load_tgt(ctx_dirty.get(), 0, 0);
+
+ for (uint32_t i = 0; i < n_rollback; ++i) {
+ const llama_pos pos = rollback_pos + i;
+ if (!decode_one(ctx_dirty.get(), tokens[pos], pos)) {
+ TLOG_ERR("%s: dirty replay failed at position %d\n", __func__, pos);
+ return case_result::fail();
+ }
+
+ const float * logits_dirty = llama_get_logits_ith(ctx_dirty.get(), 0);
+ if (logits_dirty == nullptr) {
+ TLOG_ERR("%s: missing dirty logits at position %d\n", __func__, pos);
+ return case_result::fail();
+ }
+
+ const logits_diff diff = compare_logits(logits_src_replay[i].data(), logits_dirty, n_vocab);
+ CTX_REQUIRE(diff.nmse() <= nmse_eps,
+ "dirty-ctx logits mismatch at position " + std::to_string(pos) + ", nmse " + std::to_string(diff.nmse()));
+ }
+
+ TLOG_INF("%s: recurrent rollback checkpoint restored successfully\n", __func__);
+
+ return case_result::pass();
+ }
+};
+
+// decode a prompt into seq 0, share its cells with a second sequence, then keep decoding both
+struct case_shared_seq_reserve : ctx_test {
+ test_group group() const override { return TEST_GROUP_RS; }
+
+ const char * name() const override { return "shared"; }
+
+ const char * description() const override {
+ return "sequences sharing cells keep decoding correctly after a seq_cp";
+ }
+
+ std::string skip_reason(model_run & mr) const override {
+ if (arch_reserves_single_seq(mr.model)) {
+ return "the shared-seq batch is a multi-seq graph";
+ }
+
+ return ctx_test::skip_reason(mr);
+ }
+
+ case_result run(model_run & mr) const override {
+ const int n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(mr.model));
+
+ constexpr uint32_t n_seqs = 2;
+ constexpr uint32_t n_prompt = 128;
+ constexpr uint32_t n_continue = 32;
+
+ ctx_opts opts;
+ opts.n_ctx = 512;
+ opts.n_batch = 256;
+ opts.n_ubatch = 64;
+ opts.n_seq_max = n_seqs;
+ opts.kv_unified = 1; // only a unified cache shares cells on seq_cp
+
+ auto ctx = mr.ctx(opts);
+ if (!ctx) {
+ return case_result::fail();
+ }
+
+ if (!decode_gen(ctx.get(), n_vocab, 0, 0, (llama_pos) n_prompt)) {
+ TLOG_ERR("%s: prompt decode failed\n", __func__);
+ return case_result::fail();
+ }
+
+ // this is what llama-batched-bench does for -pps
+ llama_memory_seq_cp(llama_get_memory(ctx.get()), 0, 1, -1, -1);
+
+ for (uint32_t i = 0; i < n_continue; ++i) {
+ const llama_pos pos = (llama_pos) (n_prompt + i);
+
+ common_batch batch(ctx.get());
+ for (uint32_t s = 0; s < n_seqs; ++s) {
+ batch.add(gen_token(n_vocab, (llama_seq_id) s, pos), pos, (llama_seq_id) s, true);
+ }
+ CTX_REQUIRE(llama_process(ctx.get(), LLAMA_PROCESS_TYPE_DECODE, batch.get()) == 0,
+ "shared-seq decode failed at step " + std::to_string(i));
+
+ for (uint32_t s = 0; s < n_seqs; ++s) {
+ const float * logits = llama_get_logits_ith(ctx.get(), (int) s);
+ if (logits == nullptr) {
+ TLOG_ERR("%s: missing shared-seq logits at index %u\n", __func__, s);
+ return case_result::fail();
+ }
+
+ bool finite = true;
+ for (int t_ = 0; t_ < n_vocab; ++t_) {
+ finite = finite && std::isfinite(logits[t_]);
+ }
+
+ CTX_REQUIRE(finite, "non-finite shared-seq logit at step " + std::to_string(i) + ", seq " + std::to_string(s));
+ }
+ }
+
+ TLOG_INF("%s: shared-seq decode succeeded (%u tokens after seq_cp)\n", __func__, n_continue*n_seqs);
+
+ return case_result::pass();
+ }
+};
+
+std::vector<ctx_test_ptr> make_rs_cases() {
+ std::vector<ctx_test_ptr> cases;
+
+ cases.push_back(std::make_unique<case_rollback>());
+ cases.push_back(std::make_unique<case_multi_seq_split_replay>());
+ cases.push_back(std::make_unique<case_shared_seq_reserve>());
+
+ return cases;
+}
diff --git a/tests/llama-context/test-case.cpp b/tests/llama-context/test-case.cpp
new file mode 100644
index 000000000..86bc4d313
--- /dev/null
+++ b/tests/llama-context/test-case.cpp
@@ -0,0 +1,69 @@
+#include "test-case.h"
+
+#include <iterator>
+
+const char * test_status_str(test_status status) {
+ switch (status) {
+ case test_status::PASS: return "\033[1;32mOK\033[0m";
+ case test_status::FAIL: return "\033[1;31mFAIL\033[0m";
+ case test_status::SKIP: return "\033[1;33mSKIP\033[0m";
+ }
+ return "";
+}
+
+// the printable width of a verdict, the color escapes of test_status_str are not part of it
+size_t test_status_width(test_status status) {
+ switch (status) {
+ case test_status::PASS: return 2; // "OK"
+ case test_status::FAIL: return 4; // "FAIL"
+ case test_status::SKIP: return 4; // "SKIP"
+ }
+
+ return 0;
+}
+
+test_status merge_status(test_status a, test_status b) {
+ if (a == test_status::FAIL || b == test_status::FAIL) {
+ return test_status::FAIL;
+ }
+ if (a == test_status::PASS || b == test_status::PASS) {
+ return test_status::PASS;
+ }
+ return test_status::SKIP;
+}
+
+case_result case_result::pass() {
+ return { test_status::PASS, {} };
+}
+
+case_result case_result::fail(const std::string & why) {
+ return { test_status::FAIL, why };
+}
+
+case_result case_result::skip(const std::string & why) {
+ return { test_status::SKIP, why };
+}
+
+std::string ctx_test::skip_reason(model_run & mr) const {
+ // the recurrent cases need a recurrent cache, the KV cases work with any model
+ if (group() == TEST_GROUP_RS && !mr.is_recurrent()) {
+ return "not a recurrent model";
+ }
+
+ return {};
+}
+
+const std::vector<ctx_test_ptr> & test_cases() {
+ static const std::vector<ctx_test_ptr> cases = [] {
+ std::vector<ctx_test_ptr> res;
+
+ for (const auto & maker : { make_kv_cases, make_rs_cases }) {
+ auto part = maker();
+ res.insert(res.end(), std::make_move_iterator(part.begin()), std::make_move_iterator(part.end()));
+ }
+
+ return res;
+ }();
+
+ return cases;
+}
diff --git a/tests/llama-context/test-case.h b/tests/llama-context/test-case.h
new file mode 100644
index 000000000..847fd1759
--- /dev/null
+++ b/tests/llama-context/test-case.h
@@ -0,0 +1,83 @@
+#pragma once
+
+#include "test-util.h"
+
+#include <memory>
+#include <string>
+#include <vector>
+
+//
+// the interface of a single context state case
+//
+
+enum class test_status {
+ PASS,
+ FAIL,
+ SKIP,
+};
+
+// the verdict as printed in the result table, colored
+const char * test_status_str(test_status status);
+
+// how many columns the verdict takes in the table - the escapes of the color are not printed
+size_t test_status_width(test_status status);
+
+// a skipped case does not fail the run, a failed case always does
+test_status merge_status(test_status a, test_status b);
+
+enum test_group {
+ TEST_GROUP_STATE, // KV state save / load / copy
+ TEST_GROUP_RS, // recurrent state, also run against a dirty cache
+};
+
+struct case_result {
+ test_status status = test_status::PASS;
+ std::string reason; // why the case was skipped or failed, empty when it passed
+
+ static case_result pass();
+ static case_result fail(const std::string & why = {});
+ static case_result skip(const std::string & why);
+};
+
+struct ctx_test {
+ virtual ~ctx_test() = default;
+
+ // the column header of the --models table, keep it short
+ virtual const char * name() const = 0;
+
+ // one line of what the case verifies, printed by --list
+ virtual const char * description() const { return ""; }
+
+ virtual test_group group() const { return TEST_GROUP_STATE; }
+
+ // the case compares against the logits produced by the baseline case
+ virtual bool needs_baseline() const { return false; }
+
+ // the case is also run against a cache pre-filled with a known pattern, so that a restore
+ // reading cells it never wrote cannot pass
+ virtual bool uses_fill() const { return group() == TEST_GROUP_RS; }
+
+ // non-empty means the case does not apply to this model, the text is the reported reason
+ virtual std::string skip_reason(model_run & mr) const;
+
+ // the scenario, the checks are written with CTX_REQUIRE
+ virtual case_result run(model_run & mr) const = 0;
+};
+
+// record an expectation that did not hold and give up on the case - the runner reports the
+// reason, the body does not print anything
+#define CTX_REQUIRE(cond, msg) \
+ do { \
+ if (!(cond)) { \
+ return case_result::fail(msg); \
+ } \
+ } while (0)
+
+using ctx_test_ptr = std::unique_ptr<const ctx_test>;
+
+// the cases of one source file, in table order
+std::vector<ctx_test_ptr> make_kv_cases(); // cases-kv.cpp
+std::vector<ctx_test_ptr> make_rs_cases(); // cases-rs.cpp
+
+// all cases, in table order
+const std::vector<ctx_test_ptr> & test_cases();
diff --git a/tests/llama-context/test-log.cpp b/tests/llama-context/test-log.cpp
new file mode 100644
index 000000000..17568f6b5
--- /dev/null
+++ b/tests/llama-context/test-log.cpp
@@ -0,0 +1,119 @@
+#include "test-log.h"
+
+#include <algorithm>
+#include <cstdarg>
+#include <cstdio>
+#include <filesystem>
+#include <fstream>
+#include <iterator>
+
+namespace {
+
+// a log is printed after the row of its model, what does not fit stays in the file
+const size_t LOG_DUMP_MAX_BYTES = 64 * 1024;
+
+// one worker thread runs one model and the log callback fires on the thread that logged the
+// message, so the file of a model is held by the thread that runs it
+thread_local FILE * t_file;
+
+char level_char(enum ggml_log_level level) {
+ switch (level) {
+ case GGML_LOG_LEVEL_DEBUG: return 'D';
+ case GGML_LOG_LEVEL_INFO: return 'I';
+ case GGML_LOG_LEVEL_WARN: return 'W';
+ case GGML_LOG_LEVEL_ERROR: return 'E';
+ case GGML_LOG_LEVEL_CONT: return 'C';
+ default: return '?';
+ }
+}
+
+// every line of a model log carries its level, so that a failure can be looked up with a grep
+void test_log_level(enum ggml_log_level level) {
+ if (level != GGML_LOG_LEVEL_NONE) {
+ fprintf(t_file, "%c ", level_char(level));
+ }
+}
+
+} // namespace
+
+void test_log_capture(const std::string & path) {
+ if (t_file) {
+ fclose(t_file);
+ t_file = nullptr;
+ }
+
+ if (path.empty()) {
+ return;
+ }
+
+ t_file = fopen(path.c_str(), "w");
+
+ if (!t_file) {
+ fprintf(stderr, "%s: failed to open '%s'\n", __func__, path.c_str());
+ }
+}
+
+void test_log_callback(enum ggml_log_level level, const char * text, void * user_data) {
+ if (!t_file) {
+ common_log_default_callback(level, text, user_data);
+ return;
+ }
+
+ // gate the capture the way the console would be - the debug chatter of the graph builder is
+ // orders of magnitude more than everything else, a capture never holds less than info
+ if (common_log_get_verbosity(level) <= std::max(common_log_get_verbosity_thold(), LOG_LEVEL_INFO)) {
+ test_log_level(level);
+ fputs(text, t_file);
+ }
+}
+
+void test_log_add(enum ggml_log_level level, const char * fmt, ...) {
+ va_list args;
+ va_start(args, fmt);
+
+ if (t_file) {
+ test_log_level(level);
+ vfprintf(t_file, fmt, args);
+ } else if (level == GGML_LOG_LEVEL_ERROR) {
+ // without a capture the run is followed on the console, a detail line has no place there
+ vfprintf(stderr, fmt, args);
+ fflush(stderr);
+ }
+
+ va_end(args);
+}
+
+bool test_log_dump(const std::string & path) {
+ std::ifstream f(path, std::ios::binary);
+ if (!f) {
+ fprintf(stderr, "failed to read the captured log '%s'\n", path.c_str());
+ return false;
+ }
+
+ const std::string text((std::istreambuf_iterator<char>(f)), std::istreambuf_iterator<char>());
+ f.close();
+
+ size_t skip = 0;
+ if (text.size() > LOG_DUMP_MAX_BYTES) {
+ skip = text.find('\n', text.size() - LOG_DUMP_MAX_BYTES);
+ skip = skip == std::string::npos ? 0 : skip + 1;
+ }
+
+ if (skip > 0) {
+ const size_t n_all = std::count(text.begin(), text.end(), '\n');
+ const size_t n_shown = std::count(text.begin() + skip, text.end(), '\n');
+
+ fprintf(stdout, "... %zu of %zu lines omitted, full log: %s\n", n_all - n_shown, n_all, path.c_str());
+ }
+
+ fwrite(text.data() + skip, 1, text.size() - skip, stdout);
+ fflush(stdout);
+
+ return skip > 0;
+}
+
+std::string test_log_path(const std::string & logs_dir, const std::string & model_name) {
+ const std::string file = "test-llama-context." + model_name + ".log";
+
+ return logs_dir.empty() ? file : (std::filesystem::path(logs_dir) / file).string();
+}
diff --git a/tests/llama-context/test-log.h b/tests/llama-context/test-log.h
new file mode 100644
index 000000000..a847bebe6
--- /dev/null
+++ b/tests/llama-context/test-log.h
@@ -0,0 +1,33 @@
+#pragma once
+
+#include "log.h"
+
+#include <string>
+
+//
+// the log of a model goes to a file of its own, so that a model which failed can be read on its
+// own. a worker thread runs one model and the log callback fires on the thread that produced the
+// message, so a capture belongs to the thread that opened it
+//
+
+// open the log of this thread, an empty path closes it - a capture holds what the verbosity of
+// the process allows, never less than info, the details of a failure are looked for in the file
+void test_log_capture(const std::string & path);
+
+// llama, ggml and the backends log through this: into the file of this thread when it captures,
+// into the default callback otherwise
+void test_log_callback(enum ggml_log_level level, const char * text, void * user_data);
+
+// the lines of the test itself go to the file of this thread, an error goes to stderr as well when
+// there is no capture and the run is followed on the console
+void test_log_add(enum ggml_log_level level, const char * fmt, ...) LOG_ATTRIBUTE_FORMAT(2, 3);
+
+// print a captured log, false when it was printed in full and the file can be removed
+bool test_log_dump(const std::string & path);
+
+std::string test_log_path(const std::string & logs_dir, const std::string & model_name);
+
+// a trace is a detail of a case, it is there to read in the file of the model
+#define TLOG_ERR(...) test_log_add(GGML_LOG_LEVEL_ERROR, __VA_ARGS__)
+#define TLOG_INF(...) test_log_add(GGML_LOG_LEVEL_INFO, __VA_ARGS__)
+#define TLOG_TRC(...) test_log_add(GGML_LOG_LEVEL_DEBUG, __VA_ARGS__)
diff --git a/tests/llama-context/test-runner.cpp b/tests/llama-context/test-runner.cpp
new file mode 100644
index 000000000..6f38d12f3
--- /dev/null
+++ b/tests/llama-context/test-runner.cpp
@@ -0,0 +1,117 @@
+#include "test-runner.h"
+
+#include "test-log.h"
+
+#include <algorithm>
+#include <cstdio>
+#include <cstring>
+#include <regex>
+#include <utility>
+
+std::vector<const ctx_test *> select_cases(const std::string & filter) {
+ std::vector<const ctx_test *> selected;
+
+ if (filter.empty()) {
+ selected.reserve(test_cases().size());
+ for (const auto & tc : test_cases()) {
+ selected.push_back(tc.get());
+ }
+
+ return selected;
+ }
+
+ const std::regex re(filter);
+
+ for (const auto & tc : test_cases()) {
+ if (std::regex_match(tc->name(), re)) {
+ selected.push_back(tc.get());
+ }
+ }
+
+ return selected;
+}
+
+std::vector<case_run> run_cases(model_run & mr, const std::vector<const ctx_test *> & cases, bool report) {
+ // the state cases keep the historical default of a unified KV cache when running a single
+ // sequence, the recurrent cases use whatever was passed on the command line
+ const bool unified_base = mr.params.kv_unified;
+ const bool unified_state = mr.params.n_parallel == 1 ? true : unified_base;
+
+ size_t name_width = 4; // "case"
+ for (const ctx_test * tc : cases) {
+ name_width = std::max(name_width, strlen(tc->name()));
+ }
+
+ std::vector<case_run> runs;
+ runs.reserve(cases.size());
+
+ for (const ctx_test * tc : cases) {
+ mr.params.kv_unified = tc->group() == TEST_GROUP_STATE ? unified_state : unified_base;
+
+ case_run run;
+ run.name = tc->name();
+
+ // the log of a model is read on its own, mark where each case begins
+ TLOG_INF("%s\n", run.name.c_str());
+
+ // a case that throws is a case that failed, and the run goes on with the next one
+ try {
+ if (const std::string why = tc->skip_reason(mr); !why.empty()) {
+ run.status = test_status::SKIP;
+ run.reason = why;
+ } else if (tc->needs_baseline() && !mr.have_baseline) {
+ run.status = test_status::SKIP;
+ run.reason = "no baseline";
+ } else if (!tc->uses_fill()) {
+ mr.fill = 0;
+
+ const case_result res = tc->run(mr);
+
+ run.status = res.status;
+ run.reason = res.reason;
+ } else {
+ // the recurrent cases are also run against a cache pre-filled with a known
+ // pattern, so that a restore reading cells it never wrote cannot pass
+
+ for (const uint8_t fill : { 0, 0x3e }) {
+ TLOG_INF("(fill 0x%02x)\n", fill);
+
+ mr.fill = fill;
+
+ const case_result res = tc->run(mr);
+
+ run.status = merge_status(run.status, res.status);
+ if (run.reason.empty()) {
+ run.reason = res.reason;
+ }
+ if (res.status == test_status::FAIL) {
+ break;
+ }
+ }
+ }
+ } catch (...) {
+ run.status = test_status::FAIL;
+ run.reason = "unhandled exception";
+ }
+
+ if (unified_state) {
+ mr.params.kv_unified = unified_base;
+ }
+
+ // the verdict is one word, what it means goes to the log of the model
+ if (!run.reason.empty()) {
+ TLOG_INF("%s %s: %s\n", run.name.c_str(),
+ run.status == test_status::FAIL ? "failed" : "skipped", run.reason.c_str());
+ }
+
+ if (report) {
+ fprintf(stderr, "%-*s %s%s%s\n", (int) name_width, run.name.c_str(),
+ test_status_str(run.status),
+ run.reason.empty() ? "" : " ", run.reason.c_str());
+ }
+
+ runs.push_back(std::move(run));
+ }
+
+ return runs;
+}
diff --git a/tests/llama-context/test-runner.h b/tests/llama-context/test-runner.h
new file mode 100644
index 000000000..1a6408b85
--- /dev/null
+++ b/tests/llama-context/test-runner.h
@@ -0,0 +1,23 @@
+#pragma once
+
+#include "test-case.h"
+
+#include <string>
+#include <vector>
+
+//
+// the runner: the single place that decides how a case is executed
+//
+
+struct case_run {
+ std::string name;
+ test_status status = test_status::SKIP;
+ std::string reason;
+};
+
+// the cases selected by the filter, all of them when it is empty, in table order
+std::vector<const ctx_test *> select_cases(const std::string & filter);
+
+// run the selected cases for one model and return one verdict per case - with report, print the
+// verdict of every case as it finishes, for single model mode
+std::vector<case_run> run_cases(model_run & mr, const std::vector<const ctx_test *> & cases, bool report);
diff --git a/tests/llama-context/test-util.cpp b/tests/llama-context/test-util.cpp
new file mode 100644
index 000000000..d66f90334
--- /dev/null
+++ b/tests/llama-context/test-util.cpp
@@ -0,0 +1,403 @@
+#include "test-util.h"
+
+#include "test-log.h"
+
+#include "../../src/llama-ext.h" // llama_memory_get_buffers
+
+#include "common.h"
+
+#include <cmath>
+#include <cstdio>
+#include <cstring>
+#include <limits>
+#include <stdexcept>
+
+const double NMSE_THRESHOLD = 1e-5;
+
+//
+// comparison
+//
+
+// the distance of two logits, infinite when either of them is not finite
+static float logit_diff(float a, float b) {
+ return std::isfinite(a) && std::isfinite(b) ? std::fabs(a - b) : std::numeric_limits<float>::infinity();
+}
+
+void logits_diff::add(const float * a, const float * b, size_t n_elem) {
+ for (size_t i = 0; i < n_elem; i++) {
+ const float diff = logit_diff(a[i], b[i]);
+
+ if (diff > 0.0f && i_first < 0) {
+ i_first = (int64_t) (n + i);
+ }
+
+ if (diff > diff_max) {
+ diff_max = diff;
+ }
+
+ if (std::isfinite(a[i]) && std::isfinite(b[i])) {
+ const double d = (double) a[i] - (double) b[i];
+
+ mse_ab += d*d;
+ mse_a0 += (double) a[i] * (double) a[i];
+ } else {
+ n_nan++;
+ mse_ab = std::numeric_limits<double>::infinity();
+ mse_a0 = 1.0;
+ }
+ }
+
+ n += n_elem;
+}
+
+double logits_diff::nmse() const {
+ if (mse_a0 == 0.0) {
+ return mse_ab == 0.0 ? 0.0 : std::numeric_limits<double>::infinity();
+ }
+
+ return mse_ab / mse_a0;
+}
+
+logits_diff compare_logits(const float * a, const float * b, size_t n) {
+ logits_diff res;
+ res.add(a, b, n);
+
+ return res;
+}
+
+//
+// decode
+//
+
+bool get_current_logits(llama_context * ctx, std::vector<float> & out) {
+ const auto * vocab = llama_model_get_vocab(llama_get_model(ctx));
+ const int32_t n_vocab = llama_vocab_n_tokens(vocab);
+ const float * logits = llama_get_logits_ith(ctx, -1);
+ if (logits == nullptr) {
+ return false;
+ }
+ out.assign(logits, logits + n_vocab);
+ return true;
+}
+
+bool decode_seq(llama_context * ctx, const llama_tokens & tokens, llama_pos p_begin, llama_pos p_end,
+ llama_seq_id seq_id, bool last_output) {
+ common_batch batch(ctx);
+ for (llama_pos pos = p_begin; pos < p_end; ++pos) {
+ batch.add(tokens[pos], pos, seq_id, last_output && pos + 1 == p_end);
+ }
+
+ return llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get()) == 0;
+}
+
+bool decode_tokens(llama_context * ctx, const llama_tokens & tokens, llama_seq_id seq_id) {
+ return decode_seq(ctx, tokens, 0, (llama_pos) tokens.size(), seq_id, true);
+}
+
+llama_token gen_token(int32_t n_vocab, llama_seq_id seq_id, llama_pos pos) {
+ return (llama_token) ((7*(uint32_t) pos + 31*(uint32_t) seq_id + 1) % (uint32_t) n_vocab);
+}
+
+bool decode_gen(llama_context * ctx, int32_t n_vocab, llama_seq_id seq_id, llama_pos p_begin, llama_pos p_end) {
+ common_batch batch(ctx);
+ for (llama_pos pos = p_begin; pos < p_end; ++pos) {
+ batch.add(gen_token(n_vocab, seq_id, pos), pos, seq_id, false);
+ }
+
+ return llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get()) == 0;
+}
+
+bool decode_one(llama_context * ctx, llama_token tok, llama_pos pos, llama_seq_id seq_id) {
+ common_batch batch(ctx);
+ batch.add(tok, pos, seq_id, true);
+
+ return llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get()) == 0;
+}
+
+bool generate_tokens(llama_context * ctx, llama_sampler * smpl, int & n_past, int32_t n_predict, llama_seq_id seq_id,
+ generation_result & result) {
+ common_batch batch(ctx);
+
+ // the tokens of the run as a single line - one entry per token would leave the next message
+ // of the log glued to the last one
+ std::string tokens;
+
+ for (int32_t i = 0; i < n_predict; i++) {
+ std::vector<float> logits;
+ if (!get_current_logits(ctx, logits)) {
+ TLOG_ERR("%s: failed to get logits\n", __func__);
+ return false;
+ }
+
+ const auto next_token = llama_sampler_sample(smpl, ctx, -1);
+
+ tokens += std::to_string(next_token) + " ";
+ result.tokens.push_back(next_token);
+ result.logits.push_back(std::move(logits));
+
+ batch.clear();
+ batch.add(next_token, n_past, seq_id, true);
+
+ if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get())) {
+ TLOG_ERR("%s: failed to evaluate\n", __func__);
+ return false;
+ }
+ n_past++;
+ }
+
+ TLOG_INF("%s: %s\n", __func__, tokens.c_str());
+
+ return true;
+}
+
+bool generate_tokens_compare(llama_context * ctx, llama_sampler * smpl, int & n_past, int32_t n_predict, llama_seq_id seq_id,
+ const generation_result & expected, double nmse_eps) {
+ if (expected.tokens.size() != expected.logits.size() || expected.tokens.size() < (size_t) n_predict) {
+ TLOG_ERR("%s: invalid expected generation\n", __func__);
+ return false;
+ }
+
+ common_batch batch(ctx);
+
+ std::string tokens;
+
+ for (int32_t i = 0; i < n_predict; i++) {
+ std::vector<float> logits;
+ if (!get_current_logits(ctx, logits)) {
+ TLOG_ERR("%s: failed to get logits\n", __func__);
+ return false;
+ }
+ if (logits.size() != expected.logits[i].size()) {
+ TLOG_ERR("%s: logits size mismatch at step %d: %zu != %zu\n", __func__, i, logits.size(), expected.logits[i].size());
+ return false;
+ }
+
+ const logits_diff diff = compare_logits(expected.logits[i].data(), logits.data(), logits.size());
+
+ TLOG_TRC("%s: step %d nmse = %.6e\n", __func__, i, diff.nmse());
+ if (diff.nmse() > nmse_eps) {
+ TLOG_ERR("%s: error: NMSE at step %d is %.6e (threshold %.1e)\n", __func__, i, diff.nmse(), nmse_eps);
+ return false;
+ }
+
+ const auto next_token = llama_sampler_sample(smpl, ctx, -1);
+ const auto expected_token = expected.tokens[i];
+
+ tokens += std::to_string(next_token) + " ";
+ if (next_token != expected_token) {
+ TLOG_TRC("%s: sampled token %d differs from expected %d, using expected token\n", __func__, next_token, expected_token);
+ }
+
+ batch.clear();
+ batch.add(expected_token, n_past, seq_id, true);
+
+ if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get())) {
+ TLOG_ERR("%s: failed to evaluate\n", __func__);
+ return false;
+ }
+ n_past++;
+ }
+
+ TLOG_INF("%s: %s\n", __func__, tokens.c_str());
+
+ return true;
+}
+
+//
+// state
+//
+
+bool get_seq_state(llama_context * ctx, llama_seq_id seq_id, uint32_t flags, std::vector<uint8_t> & state) {
+ const size_t state_size = llama_state_seq_get_size_ext(ctx, seq_id, flags);
+ if (state_size == 0) {
+ TLOG_ERR("%s: sequence state is empty\n", __func__);
+ return false;
+ }
+
+ state.resize(state_size);
+ const size_t ncopy = llama_state_seq_get_data_ext(ctx, state.data(), state.size(), seq_id, flags);
+ if (ncopy != state.size()) {
+ TLOG_ERR("%s: sequence state length %zu does not match expected length %zu\n",
+ __func__, ncopy, state.size());
+ return false;
+ }
+
+ return true;
+}
+
+bool set_seq_state(llama_context * ctx, llama_seq_id seq_id, uint32_t flags, const std::vector<uint8_t> & state) {
+ const size_t nset = llama_state_seq_set_data_ext(ctx, state.data(), state.size(), seq_id, flags);
+ if (nset != state.size()) {
+ TLOG_ERR("%s: seq set data length %zu does not match expected length %zu\n", __func__, nset, state.size());
+ return false;
+ }
+
+ return true;
+}
+
+bool seq_rm(llama_context * ctx, llama_seq_id seq_id) {
+ if (!llama_memory_seq_rm(llama_get_memory(ctx), seq_id, -1, -1)) {
+ TLOG_ERR("%s: failed to remove sequence %d\n", __func__, seq_id);
+ return false;
+ }
+
+ return true;
+}
+
+bool load_state_replay_last(llama_context * ctx, const std::string & path, const llama_tokens & tokens, int & n_past) {
+ llama_tokens unused_sts(tokens.size());
+ size_t n_token_count_out = 0;
+
+ if (!llama_state_load_file(ctx, path.c_str(), unused_sts.data(), unused_sts.size(), &n_token_count_out)) {
+ TLOG_ERR("%s: failed to load state\n", __func__);
+ return false;
+ }
+
+ TLOG_TRC("%s: loaded state with %zu tokens\n", __func__, n_token_count_out);
+
+ // the session is saved before the last token, so replay it at the position it was saved at
+ n_past = (int) n_token_count_out - 1;
+ if (!common_replay_last_token(ctx, tokens.back(), n_past)) {
+ return false;
+ }
+ n_past++;
+
+ return true;
+}
+
+bool blobs_equal(const std::vector<uint8_t> & a, const std::vector<uint8_t> & b) {
+ if (a.size() == b.size() && (a.empty() || memcmp(a.data(), b.data(), a.size()) == 0)) {
+ return true;
+ }
+
+ size_t i = 0;
+ while (i < a.size() && i < b.size() && a[i] == b[i]) {
+ ++i;
+ }
+
+ TLOG_ERR("%s: %zu bytes -> %zu bytes, first difference at byte %zu (%02x -> %02x)\n",
+ __func__, a.size(), b.size(), i,
+ i < a.size() ? a[i] : 0,
+ i < b.size() ? b[i] : 0);
+
+ return false;
+}
+
+bool corrupt_state(std::vector<uint8_t> & data) {
+ if (data.size() < 3*4096) {
+ TLOG_ERR("%s: state of %zu bytes is too small to corrupt\n", __func__, data.size());
+ return false;
+ }
+
+ std::fill(data.begin() + 4096, data.end() - data.size()/4, 0xff);
+ return true;
+}
+
+//
+// context
+//
+
+// fill the buffers of the cache before a restore - a restore that reads cells it never wrote
+// then produces garbage instead of silently passing
+llama_context_ptr init_ctx(llama_model * model, llama_context_params cparams, uint8_t fill) {
+ llama_context_ptr ctx{llama_init_from_model(model, cparams)};
+ if (!ctx || fill == 0) {
+ return ctx;
+ }
+
+ // Use a full ubatch so buffer discovery preserves prefill allocation sizes.
+ const uint32_t n_tokens = llama_n_ubatch(ctx.get());
+ if (!decode_tokens(ctx.get(), llama_tokens(n_tokens, 0))) {
+ return nullptr;
+ }
+ llama_synchronize(ctx.get());
+
+ const auto buffers = llama_memory_get_buffers(llama_get_memory(ctx.get()));
+ llama_memory_clear(llama_get_memory(ctx.get()), true);
+ if (buffers.empty()) {
+ TLOG_ERR("%s: no cache buffers found\n", __func__);
+ return nullptr;
+ }
+ for (auto * buffer : buffers) {
+ ggml_backend_buffer_clear(buffer, fill);
+ }
+ return ctx;
+}
+
+llama_context_ptr make_ctx(llama_model * model, const common_params & params, const ctx_opts & opts) {
+ auto cparams = common_context_params_to_llama(params);
+
+ if (opts.n_ctx) {
+ cparams.n_ctx = opts.n_ctx;
+ }
+ if (opts.n_seq_max) {
+ cparams.n_seq_max = opts.n_seq_max;
+ }
+ if (opts.kv_unified >= 0) {
+ cparams.kv_unified = opts.kv_unified != 0;
+ }
+ if (opts.n_rs_seq) {
+ cparams.n_rs_seq = opts.n_rs_seq;
+ cparams.n_batch = std::max(cparams.n_batch, (uint32_t) (opts.n_rs_seq + 1));
+ cparams.n_ubatch = std::max(cparams.n_ubatch, (uint32_t) (opts.n_rs_seq + 1));
+ }
+ if (opts.n_batch) {
+ cparams.n_batch = opts.n_batch;
+ }
+ if (opts.n_ubatch) {
+ cparams.n_ubatch = opts.n_ubatch;
+ }
+ if (opts.fa != LLAMA_FLASH_ATTN_TYPE_AUTO) {
+ cparams.flash_attn_type = opts.fa;
+ }
+ if (opts.type_k != GGML_TYPE_COUNT) {
+ cparams.type_k = opts.type_k;
+ }
+ if (opts.type_v != GGML_TYPE_COUNT) {
+ cparams.type_v = opts.type_v;
+ }
+
+ return init_ctx(model, cparams, opts.fill);
+}
+
+llama_sampler_ptr make_dist_sampler(uint32_t seed) {
+ auto sparams = llama_sampler_chain_default_params();
+ llama_sampler_ptr smpl{llama_sampler_chain_init(sparams)};
+ llama_sampler_chain_add(smpl.get(), llama_sampler_init_dist(seed));
+ return smpl;
+}
+
+bool arch_reserves_single_seq(const llama_model * model) {
+ char arch_str[64] = {};
+ llama_model_meta_val_str(model, "general.architecture", arch_str, sizeof(arch_str));
+
+ return strcmp(arch_str, "kimi-linear") == 0 || strcmp(arch_str, "minimax-01") == 0;
+}
+
+//
+// model_run
+//
+
+llama_context_ptr model_run::ctx(const ctx_opts & opts) const {
+ ctx_opts effective = opts;
+ effective.fill = fill;
+
+ auto ctx = make_ctx(model, params, effective);
+ if (!ctx && !opts.quiet) {
+ TLOG_ERR("%s: failed to create context for '%s'\n", __func__, name.c_str());
+ }
+
+ return ctx;
+}
+
+llama_sampler_ptr model_run::sampler() const {
+ return make_dist_sampler(params.sampling.seed);
+}
+
+bool model_run::is_recurrent() const {
+ return llama_model_is_recurrent(model) || llama_model_is_hybrid(model);
+}
+
+bool model_run::has_vocab() const {
+ return llama_vocab_type(llama_model_get_vocab(model)) != LLAMA_VOCAB_TYPE_NONE;
+}
diff --git a/tests/llama-context/test-util.h b/tests/llama-context/test-util.h
new file mode 100644
index 000000000..b4fc978d8
--- /dev/null
+++ b/tests/llama-context/test-util.h
@@ -0,0 +1,161 @@
+#pragma once
+
+#include "llama-cpp.h"
+#include "llama.h"
+
+#include "common.h"
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+//
+// shared helpers of the context state cases
+//
+
+extern const double NMSE_THRESHOLD;
+
+//
+// comparison
+//
+
+// normalized mean squared error = mse(a, b) / mse(a, 0), infinity if any logit is not finite
+
+// accumulates logit comparisons and derives their NMSE, so that a case can compare an entire
+// multi-seq batch row by row and still have a single verdict
+struct logits_diff {
+ double mse_ab = 0.0;
+ double mse_a0 = 0.0;
+ float diff_max = 0.0f;
+ size_t n_nan = 0;
+ size_t n = 0;
+ int64_t i_first = -1; // flat index of the first differing element
+
+ void add(const float * a, const float * b, size_t n_elem);
+
+ double nmse() const;
+};
+
+logits_diff compare_logits(const float * a, const float * b, size_t n);
+
+//
+// generation
+//
+
+// logits produced by a run, compared against by the replay cases
+struct generation_result {
+ llama_tokens tokens;
+ std::vector<std::vector<float>> logits;
+
+ bool empty() const { return tokens.empty(); }
+};
+
+//
+// decode
+//
+
+bool get_current_logits(llama_context * ctx, std::vector<float> & out);
+
+// decode tokens [p_begin, p_end) of one sequence
+bool decode_seq(llama_context * ctx, const llama_tokens & tokens, llama_pos p_begin, llama_pos p_end,
+ llama_seq_id seq_id, bool last_output);
+
+bool decode_tokens(llama_context * ctx, const llama_tokens & tokens, llama_seq_id seq_id = 0);
+
+// tokens derived from (seq, pos): the recurrent cases do not use the tokenizer, so they do
+// not depend on the prompt or on the prompt length
+llama_token gen_token(int32_t n_vocab, llama_seq_id seq_id, llama_pos pos);
+
+bool decode_gen(llama_context * ctx, int32_t n_vocab, llama_seq_id seq_id, llama_pos p_begin, llama_pos p_end);
+
+bool decode_one(llama_context * ctx, llama_token tok, llama_pos pos, llama_seq_id seq_id = 0);
+
+// greedy-free replay: sample for reporting, but feed the sampled tokens back
+bool generate_tokens(llama_context * ctx, llama_sampler * smpl, int & n_past, int32_t n_predict, llama_seq_id seq_id,
+ generation_result & result);
+
+// as generate_tokens, but feeds the expected tokens back so that the compared logits stay on
+// the reference trajectory
+bool generate_tokens_compare(llama_context * ctx, llama_sampler * smpl, int & n_past, int32_t n_predict, llama_seq_id seq_id,
+ const generation_result & expected, double nmse_eps);
+
+//
+// state
+//
+
+// llama_state_seq_get_size_ext + get_data_ext, with the length verified
+bool get_seq_state(llama_context * ctx, llama_seq_id seq_id, uint32_t flags, std::vector<uint8_t> & state);
+
+// llama_state_seq_set_data_ext, with the length verified
+bool set_seq_state(llama_context * ctx, llama_seq_id seq_id, uint32_t flags, const std::vector<uint8_t> & state);
+
+// drop a sequence entirely
+bool seq_rm(llama_context * ctx, llama_seq_id seq_id);
+
+// load a saved session and replay its last token, leaving n_past ready for the next token
+bool load_state_replay_last(llama_context * ctx, const std::string & path, const llama_tokens & tokens, int & n_past);
+
+// false when the sizes differ or any byte differs, the first difference is logged
+bool blobs_equal(const std::vector<uint8_t> & a, const std::vector<uint8_t> & b);
+
+// overwrite the tensor data with 0xff bytes (NaN when read as f16/f32), so that the restore fails
+bool corrupt_state(std::vector<uint8_t> & data);
+
+//
+// context
+//
+
+// per-test context overrides, sentinel values mean "take the value from common_params"
+struct ctx_opts {
+ uint32_t n_ctx = 0; // 0 = from params
+ uint32_t n_batch = 0; // 0 = from params
+ uint32_t n_ubatch = 0; // 0 = from params
+ uint32_t n_seq_max = 0; // 0 = from params
+ uint32_t n_rs_seq = 0; // 0 = from params, n_batch / n_ubatch are raised to fit it
+ int kv_unified = -1; // -1 = from params, 0 = off, 1 = on
+
+ llama_flash_attn_type fa = LLAMA_FLASH_ATTN_TYPE_AUTO; // AUTO = from params
+
+ ggml_type type_k = GGML_TYPE_COUNT; // COUNT = from params
+ ggml_type type_v = GGML_TYPE_COUNT; // COUNT = from params
+
+ uint8_t fill = 0; // fill the cache buffers with this byte, 0 = disabled
+
+ // a probe that expects the context to be refused should not report it as an error
+ bool quiet = false;
+};
+
+llama_context_ptr init_ctx(llama_model * model, llama_context_params cparams, uint8_t fill);
+llama_context_ptr make_ctx(llama_model * model, const common_params & params, const ctx_opts & opts);
+
+llama_sampler_ptr make_dist_sampler(uint32_t seed);
+
+// these archs reserve the final pp graph with n_seqs = 1, so any multi-seq graph has a
+// different layout and re-reserves by design - they must not be run with a multi-seq batch
+// when GGML_SCHED_DEBUG_REALLOC is enabled
+// see [TAG_RESERVE_DIAG_DECAY] in llama-context.cpp
+bool arch_reserves_single_seq(const llama_model * model);
+
+//
+// the state of a single model run, handed to every case
+//
+
+struct model_run {
+ common_params params; // per-model copy, out_file is scoped to this model
+ std::string name; // model file name, used for the table and for temp files
+ llama_model * model = nullptr;
+ llama_tokens tokens;
+ generation_result baseline;
+ bool have_baseline = false;
+ uint8_t fill = 0; // set by the runner before each case
+
+ // a context with these options, the fill of this run applied; nullptr when it failed to
+ // init, the reason is logged
+ llama_context_ptr ctx(const ctx_opts & opts = {}) const;
+
+ // a seeded distribution sampler
+ llama_sampler_ptr sampler() const;
+
+ bool is_recurrent() const;
+ bool has_vocab() const;
+};
diff --git a/tests/test-llama-context.cpp b/tests/test-llama-context.cpp
new file mode 100644
index 000000000..4ba700d53
--- /dev/null
+++ b/tests/test-llama-context.cpp
@@ -0,0 +1,552 @@
+//
+// test-llama-context - exercises the state handling of llama_context
+//
+
+#include "llama-context/test-case.h"
+#include "llama-context/test-log.h"
+#include "llama-context/test-runner.h"
+#include "llama-context/test-util.h"
+
+#include "arg.h"
+#include "common.h"
+#include "log.h"
+#include "llama-cpp.h"
+
+#include <algorithm>
+#include <atomic>
+#include <clocale>
+#include <cstdio>
+#include <cstring>
+#include <filesystem>
+#include <mutex>
+#include <random>
+#include <string>
+#include <thread>
+#include <vector>
+
+// the state file of a model, named after it so that models running in parallel do not collide
+static std::string test_state_path(const std::string & model_name) {
+ return "dump_state." + model_name + ".bin";
+}
+
+// --name VALUE or --name=VALUE: 0 when this argument is not the option, 1 when it took a value,
+// -1 when the value is missing or empty
+static int test_opt(char ** argv, int & i, int argc, const char * name, std::string & value) {
+ const size_t len = std::strlen(name);
+
+ if (std::strncmp(argv[i], name, len) != 0) {
+ return 0;
+ }
+
+ if (argv[i][len] == '=') {
+ value = argv[i] + len + 1;
+ return value.empty() ? -1 : 1;
+ }
+
+ if (argv[i][len] != '\0') {
+ return 0;
+ }
+
+ if (i + 1 >= argc) {
+ return -1;
+ }
+
+ value = argv[++i];
+ return value.empty() ? -1 : 1;
+}
+
+// tokenize the prompt, or generate random tokens for models without a tokenizer
+static llama_tokens prepare_tokens(llama_model * model, const common_params & params) {
+ llama_tokens tokens;
+
+ if (params.prompt.empty()) {
+ const int n_prompt = params.n_batch;
+
+ // this path is useful for model files that do not have a tokenizer
+ TLOG_INF("%s: no prompt provided, generating %d (n_batch) random tokens\n", __func__, n_prompt);
+
+ const auto * vocab = llama_model_get_vocab(model);
+ const auto n_vocab = llama_vocab_n_tokens(vocab);
+
+ std::mt19937 rng(params.sampling.seed);
+ std::uniform_int_distribution<llama_token> dist(0, n_vocab - 1);
+ for (int i = 0; i < n_prompt; i++) {
+ tokens.push_back(dist(rng));
+ }
+ } else {
+ TLOG_INF("%s: tokenizing prompt '%s'\n", __func__, params.prompt.c_str());
+
+ auto ctx = llama_context_ptr{llama_init_from_model(model, common_context_params_to_llama(params))};
+ tokens = common_tokenize(ctx.get(), params.prompt, true);
+ }
+
+ TLOG_INF("%s: the input prompt is %d tokens\n", __func__, (int) tokens.size());
+
+ return tokens;
+}
+
+static std::vector<std::string> collect_models(const std::string & models_dir) {
+ std::vector<std::string> models;
+
+ for (const auto & entry : std::filesystem::directory_iterator(models_dir)) {
+ if (entry.is_regular_file() && entry.path().extension() == ".gguf") {
+ models.push_back(entry.path().string());
+ }
+ }
+ std::sort(models.begin(), models.end());
+
+ return models;
+}
+
+// one row of the result table
+struct model_result {
+ std::string name;
+ std::string log_path;
+ std::vector<case_run> runs;
+ bool failed = false;
+};
+
+// run the selected cases for one model and return one verdict per case
+//
+// safe to call from a worker thread: the log capture and the state file are either per thread or
+// named after the model
+//
+// report prints the verdict of every case as it finishes, capture_log keeps the log of the run in
+// a file - table mode always captures, so that a failing model can be dumped after the table;
+// single model mode only does when --logs asks for it
+static model_result run_model(const std::string & model_path,
+ const common_params & params,
+ const std::vector<const ctx_test *> & cases,
+ const std::string & logs_dir,
+ bool report,
+ bool capture_log) {
+ model_result res;
+
+ res.name = std::filesystem::path(model_path).filename().string();
+ res.log_path = test_log_path(logs_dir, res.name);
+
+ if (capture_log) {
+ // capture the model from the moment it is loaded, so that even a load failure has a log
+ test_log_capture(res.log_path);
+ }
+
+ struct common_params model_params = params;
+ model_params.model.path = model_path;
+
+ // a model that cannot be loaded is a failure of every case, not a skip - the cases never
+ // run, so their verdicts are recorded here
+ auto llama_init = common_init_from_params(model_params, true);
+ if (llama_init->model() == nullptr) {
+ TLOG_ERR("%s: failed to init model '%s'\n", __func__, model_path.c_str());
+
+ for (const auto * tc : cases) {
+ case_run run;
+ run.name = tc->name();
+ run.status = test_status::FAIL;
+ run.reason = "the model failed to load";
+ res.runs.push_back(std::move(run));
+ }
+ } else {
+ GGML_ASSERT(llama_init->context() == nullptr);
+
+ model_run mr;
+ mr.model = llama_init->model();
+ mr.name = res.name;
+ // the cases share a working directory, so state files are named after the model
+ mr.params = std::move(model_params);
+ mr.params.out_file = test_state_path(res.name);
+ mr.tokens = prepare_tokens(mr.model, mr.params);
+
+ res.runs = run_cases(mr, cases, report);
+ }
+
+ if (capture_log) {
+ test_log_capture("");
+ }
+
+ std::remove(test_state_path(res.name).c_str()); // the cases leave the session file behind
+
+ res.failed = std::any_of(res.runs.begin(), res.runs.end(),
+ [](const case_run & run) { return run.status == test_status::FAIL; });
+
+ return res;
+}
+
+static void print_usage(int /* argc */, char ** argv) {
+ LOG("\nexample usage:\n");
+ LOG("\n %s -m your_model.gguf\n", argv[0]);
+ LOG("\n %s --models tests/test-models\n", argv[0]);
+ LOG("\n %s --models tests/test-models --test 'rollback|replay'\n", argv[0]);
+ LOG("\n %s --models tests/test-models -j 8\n", argv[0]);
+ LOG("\n %s -m your_model.gguf -lv 5\n", argv[0]);
+ LOG("\nspecial flags:\n");
+ LOG("\n --models DIR run every .gguf in DIR and print a result table\n");
+ LOG("\n --test REGEX run only the cases whose name matches REGEX\n");
+ LOG("\n --list list the registered cases and exit\n");
+ LOG("\n -j N, --threads N run N models at a time in --models mode (default: min(nproc, 4))\n");
+ LOG("\n --logs DIR keep the log of each model in DIR instead of discarding it\n");
+ LOG("\n --logs-on-fail always after the table, print the log of every failed model (default)\n");
+ LOG("\n --logs-on-fail never never print the logs of the failed models\n");
+ LOG("\n");
+}
+
+static void print_cases() {
+ printf("\ncases:\n");
+
+ for (const auto & tc : test_cases()) {
+ printf(" %-11s %-5s %s\n", tc->name(), tc->group() == TEST_GROUP_RS ? "[rs]" : "[kv]", tc->description());
+ }
+
+ printf("\nthe [rs] cases are also run against a cache pre-filled with a known pattern\n\n");
+}
+
+int main(int argc, char ** argv) {
+ std::setlocale(LC_NUMERIC, "C");
+
+ common_params params;
+ params.prompt = "";
+ params.n_batch = 100;
+ params.out_file = "dump_state.bin";
+ params.sampling.seed = 1234;
+
+ common_init();
+
+ // extract our own options before handing the rest to the common parser
+ std::string models_dir;
+ std::string logs_dir;
+ std::string test_filter;
+
+ bool list_cases = false;
+ bool dump_logs_on_fail = true;
+
+ int parallel_workers = 1; // 0 = auto, same flag as in test-backend-ops
+
+ auto parse_parallel_workers = [&](const char * flag, const char * value) {
+ try {
+ parallel_workers = std::stoi(value);
+ } catch (const std::exception &) {
+ fprintf(stderr, "%s: expected a number, got '%s'\n", flag, value);
+ return false;
+ }
+ if (parallel_workers < 1) {
+ fprintf(stderr, "%s: expected a positive number, got %d\n", flag, parallel_workers);
+ return false;
+ }
+
+ return true;
+ };
+
+ std::vector<char *> filtered_argv;
+ filtered_argv.push_back(argv[0]);
+ for (int i = 1; i < argc; i++) {
+ std::string value;
+
+ const int res_models = test_opt(argv, i, argc, "--models", value);
+ if (res_models != 0) {
+ if (res_models < 0) {
+ fprintf(stderr, "--models requires a directory argument\n");
+ return 1;
+ }
+ models_dir = value;
+ continue;
+ }
+
+ const int res_logs = test_opt(argv, i, argc, "--logs", value);
+ if (res_logs != 0) {
+ if (res_logs < 0) {
+ fprintf(stderr, "--logs requires a directory argument\n");
+ return 1;
+ }
+ logs_dir = value;
+ continue;
+ }
+
+ const int res_test = test_opt(argv, i, argc, "--test", value);
+ if (res_test != 0) {
+ if (res_test < 0) {
+ fprintf(stderr, "--test requires a regular expression argument\n");
+ return 1;
+ }
+ test_filter = value;
+ continue;
+ }
+
+ if (test_opt(argv, i, argc, "--logs-on-fail", value) != 0) {
+ if (value == "always") {
+ dump_logs_on_fail = true;
+ } else if (value == "never") {
+ dump_logs_on_fail = false;
+ } else {
+ fprintf(stderr, "--logs-on-fail: expected 'always' or 'never', got '%s'\n", value.c_str());
+ return 1;
+ }
+ continue;
+ }
+
+ // -j is the flag of test-backend-ops, --threads is its long form
+ if (std::strcmp(argv[i], "-j") == 0) {
+ if (++i >= argc) {
+ fprintf(stderr, "-j requires a number argument\n");
+ return 1;
+ }
+ if (!parse_parallel_workers("-j", argv[i])) {
+ return 1;
+ }
+ continue;
+ }
+
+ const int res_threads = test_opt(argv, i, argc, "--threads", value);
+ if (res_threads != 0) {
+ if (res_threads < 0) {
+ fprintf(stderr, "--threads requires a number argument\n");
+ return 1;
+ }
+ if (!parse_parallel_workers("--threads", value.c_str())) {
+ return 1;
+ }
+ continue;
+ }
+
+ if (std::strcmp(argv[i], "--list") == 0) {
+ list_cases = true;
+ continue;
+ }
+
+ filtered_argv.push_back(argv[i]);
+ }
+ filtered_argv.push_back(nullptr);
+ const int fargc = (int) filtered_argv.size() - 1;
+
+ if (list_cases) {
+ print_cases();
+ return 0;
+ }
+
+ // in --models mode there is no single model; set a placeholder so the common parser's
+ // "--model is required" check passes (each model is set individually inside the loop)
+ if (!models_dir.empty()) {
+ params.model.path = models_dir;
+ }
+
+ if (!common_params_parse(fargc, filtered_argv.data(), params, LLAMA_EXAMPLE_COMMON, print_usage)) {
+ return 1;
+ }
+
+ // one model per thread: fitting the params to the free memory would swap in a log callback
+ // that points to the stack of the thread doing the fit, which is only safe while no other
+ // thread is inside llama - and the models of this test are far too small to need it
+ params.fit_params = false;
+
+ const std::vector<const ctx_test *> cases = select_cases(test_filter);
+ if (cases.empty()) {
+ fprintf(stderr, "--test '%s' selected no cases, see --list\n", test_filter.c_str());
+ return 1;
+ }
+
+ if (params.n_predict < 0) {
+ params.n_predict = 16;
+ }
+
+ // capture the logs of each model to a file, so that they can be dumped when a model fails:
+ // the table stays clean, but the details of a failure stay close at hand
+ if (!logs_dir.empty() && !std::filesystem::create_directories(logs_dir)) {
+ fprintf(stderr, "failed to create the log directory '%s'\n", logs_dir.c_str());
+ return 1;
+ }
+
+ // route the log of every thread to its own capture - without a capture of its own, a thread
+ // writes to the console as usual
+ llama_log_set(test_log_callback, nullptr);
+
+ llama_backend_init();
+
+ // the threads a context runs on, resolved by the common parser - in table mode every worker
+ // builds a context of its own, so the run uses n_pool times this many
+ const int32_t n_threads = params.cpuparams.n_threads;
+
+ // run the full suite for a single model
+ if (models_dir.empty()) {
+ LOG("threads: %d\n", n_threads);
+
+ const model_result res = run_model(params.model.path, params, cases, logs_dir,
+ true /* report */, !logs_dir.empty());
+
+ size_t n_pass = 0;
+ size_t n_skip = 0;
+ size_t n_fail = 0;
+
+ for (const auto & run : res.runs) {
+ switch (run.status) {
+ case test_status::PASS: n_pass++; break;
+ case test_status::SKIP: n_skip++; break;
+ case test_status::FAIL: n_fail++; break;
+ }
+ }
+
+ fprintf(stderr, "\ncases : %zu\n", res.runs.size());
+ fprintf(stderr, "passed : %zu\n", n_pass);
+ fprintf(stderr, "skipped : %zu\n", n_skip);
+ fprintf(stderr, "failed : %zu\n", n_fail);
+
+ return n_fail == 0 ? 0 : 1;
+ }
+
+ // table mode: run the suite for every model in the directory
+ if (!std::filesystem::exists(models_dir) || !std::filesystem::is_directory(models_dir)) {
+ TLOG_ERR("%s: models directory '%s' does not exist\n", __func__, models_dir.c_str());
+ return 1;
+ }
+
+ const auto models = collect_models(models_dir);
+ if (models.empty()) {
+ TLOG_ERR("%s: no .gguf models found in '%s'\n", __func__, models_dir.c_str());
+ return 1;
+ }
+
+ // the dummy models are tiny, but the same flag is meant to run sweeps of real checkpoints,
+ // where every worker holds the weights it loaded - so do not scale with the cores
+ const unsigned n_hw = std::thread::hardware_concurrency();
+ const size_t n_auto = n_hw > 0 ? std::min<size_t>(n_hw, 4) : 1;
+ const size_t n_pool = parallel_workers > 0
+ ? std::min<size_t>((size_t) parallel_workers, models.size())
+ : n_auto;
+
+ // a column is as wide as its case name, but never narrower than the widest verdict
+ const size_t verdict_width = std::max({ test_status_width(test_status::PASS),
+ test_status_width(test_status::FAIL),
+ test_status_width(test_status::SKIP) });
+
+ auto col_width = [&](const char * name) { return (int) std::max(strlen(name), verdict_width); };
+
+ size_t name_width = 5; // "Model"
+ for (const auto & model_path : models) {
+ name_width = std::max(name_width, std::filesystem::path(model_path).filename().string().size());
+ }
+
+ // silence everything but the table itself (LOG has verbosity LOG_LEVEL_OUTPUT = 0)
+ common_log_set_verbosity_thold(0);
+
+ if (!test_filter.empty()) {
+ LOG("filter: %s\n", test_filter.c_str());
+ }
+ LOG("models in parallel: %zu, threads per model: %d\n", n_pool, n_threads);
+ common_log_flush(common_log_main());
+
+ LOG("%-*s", (int) name_width, "Model");
+ for (const auto & tc : cases) {
+ LOG(" %-*s", col_width(tc->name()), tc->name());
+ }
+ LOG("\n");
+ common_log_flush(common_log_main());
+
+ std::vector<size_t> n_pass(cases.size(), 0);
+ std::vector<size_t> n_skip(cases.size(), 0);
+ std::vector<size_t> n_fail(cases.size(), 0);
+
+ size_t n_model_fail = 0;
+
+ // the models are independent of each other, so they run on several threads - but the table
+ // is printed in model order, with the log of a failed model right after its row, exactly as
+ // if the models had run one by one: a worker stores its result and then prints every row
+ // that became printable in order, under the print lock
+ std::vector<model_result> results(models.size());
+ std::vector<char> ready(models.size(), 0);
+
+ std::mutex pool_mutex;
+ size_t next_print = 0;
+
+ auto print_ready = [&]() {
+ std::lock_guard<std::mutex> lock(pool_mutex);
+
+ while (next_print < results.size() && ready[next_print]) {
+ const model_result & res = results[next_print];
+
+ LOG("%-*s", (int) name_width, res.name.c_str());
+
+ for (size_t i = 0; i < res.runs.size(); i++) {
+ const int pad = (int) (col_width(cases[i]->name()) - test_status_width(res.runs[i].status));
+
+ LOG(" %s%*s", test_status_str(res.runs[i].status), pad, "");
+
+ switch (res.runs[i].status) {
+ case test_status::PASS: n_pass[i]++; break;
+ case test_status::FAIL: n_fail[i]++; break;
+ case test_status::SKIP: n_skip[i]++; break;
+ }
+ }
+ LOG("\n");
+ common_log_flush(common_log_main());
+
+ // a log is discarded when it was printed in full and nobody asked to keep it
+ bool keep_log = !logs_dir.empty();
+
+ if (res.failed) {
+ n_model_fail++;
+
+ // dump the log of a failed model once its row is printed - interleaving it with
+ // the rows would make the table unreadable
+ if (dump_logs_on_fail) {
+ LOG("\n=== logs: %s (%s) ===\n", res.name.c_str(), res.log_path.c_str());
+ common_log_flush(common_log_main());
+
+ // a dump that had to be cut short leaves the file behind, the table cannot
+ // show the whole log
+ keep_log = keep_log || !test_log_dump(res.log_path);
+ }
+ }
+
+ if (!keep_log) {
+ std::remove(res.log_path.c_str());
+ }
+
+ ++next_print;
+ }
+ };
+
+ std::atomic<size_t> next_model(0);
+
+ auto worker = [&]() {
+ // no verdict lines here - the table carries the verdicts, the reasons are logged to the
+ // file of the model
+ for (size_t i = next_model.fetch_add(1); i < results.size(); i = next_model.fetch_add(1)) {
+ results[i] = run_model(models[i], params, cases, logs_dir,
+ false /* report */, true /* capture_log */);
+
+ {
+ std::lock_guard<std::mutex> lock(pool_mutex);
+ ready[i] = 1;
+ }
+
+ print_ready();
+ }
+ };
+
+ std::vector<std::thread> pool;
+ pool.reserve(n_pool > 0 ? n_pool - 1 : 0);
+
+ for (size_t i = 1; i < n_pool; i++) {
+ pool.emplace_back(worker);
+ }
+
+ worker(); // the main thread runs models like any other worker
+
+ for (auto & t : pool) {
+ t.join();
+ }
+
+ print_ready(); // the last worker can have stored its result while the others were joining
+
+ common_log_set_verbosity_thold(LOG_DEFAULT_LLAMA);
+ common_log_flush(common_log_main());
+
+ for (size_t i = 0; i < cases.size(); i++) {
+ if (n_fail[i] == 0 && n_skip[i] == 0) {
+ continue;
+ }
+ LOG_INF("%s: %-13s %zu passed, %zu skipped, %zu failed (of %zu)\n",
+ __func__, cases[i]->name(), n_pass[i], n_skip[i], n_fail[i], models.size());
+ }
+ LOG_INF("%s: models: %zu passed, %zu failed (of %zu)\n", __func__, models.size() - n_model_fail, n_model_fail, models.size());
+
+ return n_model_fail == 0 ? 0 : 1;
+}
diff --git a/tests/test-recurrent-state-rollback.cpp b/tests/test-recurrent-state-rollback.cpp
deleted file mode 100644
index 006486b26..000000000
--- a/tests/test-recurrent-state-rollback.cpp
+++ /dev/null
@@ -1,739 +0,0 @@
-// TODO: merge with test-save-load-state.cpp
-// TODO: merge with test-state-restore-fragmented.cpp
-
-#include "arg.h"
-#include "common.h"
-#include "ggml-backend.h"
-#include "log.h"
-#include "llama-cpp.h"
-#include "llama.h"
-
-#include "../src/llama-io.h"
-#include "../src/llama-memory.h"
-
-#include <algorithm>
-#include <clocale>
-#include <cmath>
-#include <cstdio>
-#include <cstring>
-#include <filesystem>
-#include <limits>
-#include <set>
-#include <string>
-#include <vector>
-
-enum class test_status {
- PASS,
- FAIL,
- SKIP,
-};
-
-static const char * test_status_str(test_status status) {
- switch (status) {
- case test_status::PASS: return "\033[1;32mPASS\033[m";
- case test_status::FAIL: return "\033[1;31mFAIL\033[m";
- case test_status::SKIP: return "\033[1;33mSKIP\033[m";
- }
- return "";
-}
-
-static bool decode_tokens(llama_context * ctx, const std::vector<llama_token> & tokens) {
- common_batch batch(ctx);
- for (uint32_t pos = 0; pos < tokens.size(); ++pos) {
- batch.add(tokens[pos], pos, 0, pos + 1 == tokens.size());
- }
- return llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get()) == 0;
-}
-
-static bool decode_one(llama_context * ctx, llama_token tok, llama_pos pos, llama_seq_id seq = 0) {
- common_batch batch(ctx);
- batch.add(tok, pos, seq, true);
- return llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get()) == 0;
-}
-
-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_ptr init_ctx(llama_model * model, llama_context_params cparams, uint8_t fill) {
- llama_context_ptr ctx{llama_init_from_model(model, cparams)};
- if (!ctx || fill == 0) {
- return ctx;
- }
-
- // Use a full ubatch so buffer discovery preserves prefill allocation sizes.
- const uint32_t n_tokens = llama_n_ubatch(ctx.get());
- if (!decode_tokens(ctx.get(), std::vector<llama_token>(n_tokens, 0))) {
- return nullptr;
- }
- llama_synchronize(ctx.get());
- cache_buffer_collector collector;
- llama_get_memory(ctx.get())->state_write(collector);
- llama_memory_clear(llama_get_memory(ctx.get()), true);
- if (collector.buffers.empty()) {
- LOG_ERR("%s: no cache buffers found\n", __func__);
- return nullptr;
- }
- for (auto * buffer : collector.buffers) {
- ggml_backend_buffer_clear(buffer, fill);
- }
- return ctx;
-}
-
-static llama_context_ptr 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();
-}
-
-static double nmse(const float * a, const float * b, int n) {
- double mse_ab = 0.0;
- double mse_a0 = 0.0;
- for (int i = 0; i < n; i++) {
- if (!std::isfinite(a[i]) || !std::isfinite(b[i])) {
- return std::numeric_limits<double>::infinity();
- }
- const double diff = (double) a[i] - b[i];
- mse_ab += diff*diff;
- mse_a0 += (double) a[i]*a[i];
- }
- return mse_a0 == 0.0 ? (mse_ab == 0.0 ? 0.0 : std::numeric_limits<double>::infinity()) : mse_ab/mse_a0;
-}
-
-// 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 test_status test_multi_seq_split_replay(const common_params & params, llama_model * model, uint8_t fill) {
- const int n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model));
-
- constexpr uint32_t n_seqs = 2;
- constexpr uint32_t n_ubatch = 16;
- constexpr uint32_t n_prompt = 20;
- constexpr uint32_t n_rollback = 3;
- constexpr uint32_t n_replay = 40; // > n_ubatch so each seq spans multiple ubatches
- constexpr llama_pos p0 = n_prompt - n_rollback;
-
- const auto make_ctx_multi = [&]() {
- auto cparams = common_context_params_to_llama(params);
- cparams.n_seq_max = n_seqs;
- cparams.n_rs_seq = 8;
- cparams.n_ctx = 256;
- cparams.n_batch = 256;
- cparams.n_ubatch = n_ubatch;
- cparams.kv_unified = false;
- return init_ctx(model, cparams, fill);
- };
-
- llama_context_ptr ctx_roll = make_ctx_multi();
- llama_context_ptr ctx_ref = make_ctx_multi();
- if (!ctx_roll || !ctx_ref) {
- LOG_ERR("%s: failed to init multi-seq contexts\n", __func__);
- return test_status::FAIL;
- }
-
- if (llama_n_rs_seq(ctx_roll.get()) < n_rollback) {
- LOG_INF("%s: skipping because n_rs_seq is too small\n", __func__);
- return test_status::SKIP;
- }
-
- const auto tok = [&](uint32_t seq, llama_pos pos) {
- return (llama_token) ((7*(uint32_t) pos + 31*seq + 1) % (uint32_t) n_vocab);
- };
-
- bool ok = true;
-
- // decode tokens [p_begin, p_end) of seq s, a batch belongs to one context so it is built per call
- const auto decode_range = [&](llama_context * ctx, uint32_t s, llama_pos p_begin, llama_pos p_end) {
- common_batch batch(ctx);
- for (llama_pos pos = p_begin; pos < p_end; ++pos) {
- batch.add(tok(s, pos), pos, (llama_seq_id) s, false);
- }
- return llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get()) == 0;
- };
-
- // Keep the anchor at p0 - 1 when rolling back the tail.
- for (uint32_t s = 0; s < n_seqs && ok; ++s) {
- ok = ok && decode_range(ctx_roll.get(), s, 0, (llama_pos) p0 - 1);
- ok = ok && decode_range(ctx_ref.get(), s, 0, (llama_pos) p0 - 1);
-
- ok = ok && decode_range(ctx_roll.get(), s, (llama_pos) p0 - 1, (llama_pos) n_prompt);
- ok = ok && decode_range(ctx_ref.get(), s, (llama_pos) p0 - 1, (llama_pos) p0);
-
- ok = ok && llama_memory_seq_rm(llama_get_memory(ctx_roll.get()), (llama_seq_id) s, p0, -1);
-
- // a second partial removal while one is pending must be refused
- ok = ok && !llama_memory_seq_rm(llama_get_memory(ctx_roll.get()), (llama_seq_id) s, p0 - 1, -1);
- }
- if (!ok) {
- LOG_ERR("%s: multi-seq prefill/rollback failed\n", __func__);
- return test_status::FAIL;
- }
-
- // all seqs replay in a single batch
- const auto decode_replay = [&](llama_context * ctx) {
- common_batch batch(ctx);
- for (uint32_t s = 0; s < n_seqs; ++s) {
- for (uint32_t i = 0; i < n_replay; ++i) {
- const llama_pos pos = p0 + (llama_pos) i;
- batch.add(tok(s, pos), pos, (llama_seq_id) s, true);
- }
- }
- return llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get()) == 0;
- };
- ok = decode_replay(ctx_roll.get());
- ok = ok && decode_replay(ctx_ref.get());
- if (!ok) {
- LOG_ERR("%s: multi-seq replay decode failed\n", __func__);
- return test_status::FAIL;
- }
-
- // both contexts decode identical batches, so the logits should match;
- // random dummy models can still drift up to ~1.7e-5, so the bound is 1e-4
- constexpr float nmse_eps = 1e-4f;
-
- float diff_max = 0.0f;
- uint32_t seq_first = 0;
- int32_t pos_first = -1;
- double nmse_ab = 0.0;
- double nmse_a0 = 0.0;
- for (uint32_t i = 0; i < n_seqs*n_replay; ++i) {
- const float * l_roll = llama_get_logits_ith(ctx_roll.get(), i);
- const float * l_ref = llama_get_logits_ith(ctx_ref.get(), i);
- if (l_roll == nullptr || l_ref == nullptr) {
- LOG_ERR("%s: missing multi-seq logits at index %u\n", __func__, i);
- return test_status::FAIL;
- }
- for (int t = 0; t < n_vocab; ++t) {
- const float r = l_roll[t];
- const float f = l_ref[t];
- const float diff = logit_diff(r, f);
- if (diff > 0.0f && pos_first < 0) {
- seq_first = i/n_replay;
- pos_first = p0 + (int32_t) (i%n_replay);
- }
- diff_max = std::max(diff_max, diff);
- if (std::isfinite(r) && std::isfinite(f)) {
- const double d = (double) r - f;
- nmse_ab += d*d;
- nmse_a0 += (double) r*r;
- } else {
- nmse_ab = std::numeric_limits<double>::infinity();
- nmse_a0 = 1.0;
- }
- }
- }
- const double nmse_val = nmse_a0 == 0.0 ? (nmse_ab == 0.0 ? 0.0 : std::numeric_limits<double>::infinity()) : nmse_ab/nmse_a0;
-
- if (nmse_val > nmse_eps) {
- LOG_ERR("%s: multi-seq split replay logits mismatch (max diff %g, nmse %g, first at seq %u pos %d)\n",
- __func__, (double) diff_max, nmse_val, seq_first, pos_first);
- return test_status::FAIL;
- }
-
- LOG_INF("%s: multi-seq split replay matched (max diff %g, nmse %g)\n", __func__, (double) diff_max, nmse_val);
-
- // seq-1-only decodes must be independent of seq 0's content: diverge seq 0
- // in ctx_ref only, then compare identical seq-1-only continuations bitwise
- constexpr uint32_t n_tail = 4;
-
- {
- common_batch batch_tail(ctx_ref.get());
- for (uint32_t i = 0; i < n_tail; ++i) {
- const llama_pos pos = p0 + (llama_pos) (n_replay + i);
- batch_tail.add(tok(0, pos + 7), pos, 0, false);
- }
- ok = llama_process(ctx_ref.get(), LLAMA_PROCESS_TYPE_DECODE, batch_tail.get()) == 0;
- }
-
- float diff_tail = 0.0f;
- double nmse_tail_ab = 0.0;
- double nmse_tail_a0 = 0.0;
- for (uint32_t i = 0; i < n_tail && ok; ++i) {
- const llama_pos pos = p0 + (llama_pos) (n_replay + i);
- ok = decode_one(ctx_roll.get(), tok(1, pos), pos, 1);
- ok = ok && decode_one(ctx_ref.get(), tok(1, pos), pos, 1);
- if (!ok) {
- break;
- }
-
- const float * l_roll = llama_get_logits_ith(ctx_roll.get(), 0);
- const float * l_ref = llama_get_logits_ith(ctx_ref.get(), 0);
- ok = l_roll != nullptr && l_ref != nullptr;
- for (int t = 0; ok && t < n_vocab; ++t) {
- const float r = l_roll[t];
- const float f = l_ref[t];
- diff_tail = std::max(diff_tail, logit_diff(r, f));
- if (std::isfinite(r) && std::isfinite(f)) {
- const double d = (double) r - f;
- nmse_tail_ab += d*d;
- nmse_tail_a0 += (double) r*r;
- } else {
- nmse_tail_ab = std::numeric_limits<double>::infinity();
- nmse_tail_a0 = 1.0;
- }
- }
- }
- const double nmse_tail = nmse_tail_a0 == 0.0 ? (nmse_tail_ab == 0.0 ? 0.0 : std::numeric_limits<double>::infinity()) : nmse_tail_ab/nmse_tail_a0;
-
- if (!ok || nmse_tail > nmse_eps) {
- LOG_ERR("%s: seq-1-only decode leaked seq 0 state (ok=%d, max diff %g, nmse %g)\n",
- __func__, ok ? 1 : 0, (double) diff_tail, nmse_tail);
- return test_status::FAIL;
- }
-
- LOG_INF("%s: seq-1-only decode independent of seq 0 (max diff %g, nmse %g)\n", __func__, (double) diff_tail, nmse_tail);
- return test_status::PASS;
-}
-
-// Save a rolled-back single-seq state, restore it into fresh and dirty
-// contexts, and verify exact logit matches on replay.
-static test_status 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_ptr ctx_src = make_ctx(params, model, fill);
- llama_context_ptr ctx_dst = make_ctx(params, model, fill);
- if (!ctx_src || !ctx_dst) {
- LOG_ERR("%s: failed to init contexts\n", __func__);
- return test_status::FAIL;
- }
-
- if (llama_n_rs_seq(ctx_src.get()) == 0) {
- LOG_INF("%s: skipping because n_rs_seq is disabled\n", __func__);
- return test_status::SKIP;
- }
-
- std::vector<llama_token> tokens;
- if (llama_vocab_type(vocab) == LLAMA_VOCAB_TYPE_NONE) {
- tokens = { 1, 2, 3, 4, 5, 6, 7, 8, 9 };
- } else {
- tokens = common_tokenize(ctx_src.get(), "The quick brown fox jumps over the lazy dog", true);
- }
- const uint32_t n_rs_seq = llama_n_rs_seq(ctx_src.get());
- constexpr uint32_t n_rollback = 3;
- if (n_rs_seq < n_rollback) {
- LOG_INF("%s: skipping because n_rs_seq is too small\n", __func__);
- return test_status::SKIP;
- }
- if (tokens.empty()) {
- LOG_ERR("%s: not enough prompt tokens\n", __func__);
- return test_status::FAIL;
- }
- tokens.resize(n_rs_seq + 1, tokens.back());
-
- const uint32_t n_tokens = tokens.size();
- const llama_pos rollback_pos = (llama_pos) n_tokens - n_rollback;
-
- // Decode the full prompt on the source, then roll back three positions.
- // Replaying them crosses DSV4's ratio-4 compressor boundary.
- // Rollback leaves the recurrent memory in a snapshot state (rs_idx != 0).
- if (!decode_tokens(ctx_src.get(), tokens)) {
- LOG_ERR("%s: failed to decode prompt\n", __func__);
- return test_status::FAIL;
- }
- if (!llama_memory_seq_rm(llama_get_memory(ctx_src.get()), 0, rollback_pos, -1)) {
- LOG_ERR("%s: rollback failed\n", __func__);
- return test_status::FAIL;
- }
-
- // Save the rolled-back state and restore it into a fresh context.
- common_prompt_checkpoint ckpt;
- ckpt.update_tgt(ctx_src.get(), 0, 0);
- ckpt.load_tgt(ctx_dst.get(), 0, 0);
-
- constexpr float nmse_eps = 0.0;
- std::vector<std::vector<float>> logits_src_replay(n_rollback);
- const auto replay_and_compare = [&](const char * mode) {
- for (uint32_t i = 0; i < n_rollback; ++i) {
- const llama_pos pos = rollback_pos + i;
- if (!decode_one(ctx_src.get(), tokens[pos], pos) ||
- !decode_one(ctx_dst.get(), tokens[pos], pos)) {
- LOG_ERR("%s: %s replay failed at position %d\n", __func__, mode, pos);
- return false;
- }
-
- const float * logits_src = llama_get_logits_ith(ctx_src.get(), 0);
- const float * logits_dst = llama_get_logits_ith(ctx_dst.get(), 0);
- if (logits_src == nullptr || logits_dst == nullptr) {
- LOG_ERR("%s: missing %s logits at position %d\n", __func__, mode, pos);
- return false;
- }
-
- logits_src_replay[i].assign(logits_src, logits_src + n_vocab);
- const double nmse_val = nmse(logits_src, logits_dst, n_vocab);
- int token_first = -1;
- for (int token = 0; token < n_vocab; ++token) {
- if (logit_diff(logits_src[token], logits_dst[token]) > 0.0f && token_first < 0) {
- token_first = token;
- }
- }
- if (nmse_val > nmse_eps) {
- LOG_ERR("%s: %s logits mismatch at position %d, first token %d, nmse %g\n",
- __func__, mode, pos, token_first, nmse_val);
- return false;
- }
- }
- return true;
- };
- if (!replay_and_compare("full")) {
- return test_status::FAIL;
- }
-
- // TODO: this test is invalid because RS rollback is only correct once after a ubatch with more than n_rs_seq tokens
- // this is not the case here. add asserts and guardrails to prevent such attempts
- //if (!llama_memory_seq_rm(llama_get_memory(ctx_src), 0, rollback_pos, -1) ||
- // !llama_memory_seq_rm(llama_get_memory(ctx_dst), 0, rollback_pos, -1)) {
- // fprintf(stderr, "%s : partial rollback failed\n", __func__);
- // return 1;
- //}
-
- //constexpr llama_state_seq_flags partial_flags = LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY;
- //common_prompt_checkpoint ckpt_partial;
- //ckpt_partial.update_tgt(ctx_src, 0, partial_flags);
- //ckpt_partial.load_tgt(ctx_dst, 0, partial_flags);
-
- //if (!replay_and_compare("partial")) {
- // return 1;
- //}
-
- // 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_ptr ctx_dirty = make_ctx(params, model, fill);
- if (!ctx_dirty) {
- LOG_ERR("%s: failed to init dirty ctx\n", __func__);
- return test_status::FAIL;
- }
-
- std::vector<llama_token> noise = tokens;
- for (auto & t : noise) {
- t = (t + 1) % n_vocab;
- if (t < 0) {
- t = 0;
- }
- }
- if (!decode_tokens(ctx_dirty.get(), noise)) {
- LOG_ERR("%s: dirty prompt decode failed\n", __func__);
- return test_status::FAIL;
- }
- if (!llama_memory_seq_rm(llama_get_memory(ctx_dirty.get()), 0, rollback_pos, -1)) {
- LOG_ERR("%s: dirty rollback failed\n", __func__);
- return test_status::FAIL;
- }
-
- ckpt.load_tgt(ctx_dirty.get(), 0, 0);
-
- for (uint32_t i = 0; i < n_rollback; ++i) {
- const llama_pos pos = rollback_pos + i;
- if (!decode_one(ctx_dirty.get(), tokens[pos], pos)) {
- LOG_ERR("%s: dirty replay failed at position %d\n", __func__, pos);
- return test_status::FAIL;
- }
-
- const float * logits_dirty = llama_get_logits_ith(ctx_dirty.get(), 0);
- if (logits_dirty == nullptr) {
- LOG_ERR("%s: missing dirty logits at position %d\n", __func__, pos);
- return test_status::FAIL;
- }
-
- const double nmse_dirty = nmse(logits_src_replay[i].data(), logits_dirty, n_vocab);
- int token_first = -1;
- for (int token = 0; token < n_vocab; ++token) {
- if (logit_diff(logits_src_replay[i][token], logits_dirty[token]) > 0.0f && token_first < 0) {
- token_first = token;
- }
- }
- if (nmse_dirty > nmse_eps) {
- LOG_ERR("%s: dirty-ctx logits mismatch at position %d, first token %d, nmse %g\n",
- __func__, pos, token_first, nmse_dirty);
- return test_status::FAIL;
- }
- }
-
- LOG_INF("%s: recurrent rollback checkpoint restored successfully\n", __func__);
- return test_status::PASS;
-}
-
-// Decode a prompt into seq 0, share its cells with a second sequence, then keep
-// decoding both.
-static test_status test_shared_seq_reserve(const common_params & params, llama_model * model, uint8_t fill) {
- const int n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model));
-
- // these archs reserve the final pp graph with n_seqs = 1, so every multi-seq
- // graph has a different layout and re-reserves by design
- // see [TAG_RESERVE_DIAG_DECAY] in llama-context.cpp
- char arch_str[64] = {};
- llama_model_meta_val_str(model, "general.architecture", arch_str, sizeof(arch_str));
- if (strcmp(arch_str, "kimi-linear") == 0 || strcmp(arch_str, "minimax-01") == 0) {
- LOG_INF("%s: skipping %s, its reserve uses n_seqs = 1\n", __func__, arch_str);
- return test_status::SKIP;
- }
-
- constexpr uint32_t n_seqs = 2;
- constexpr uint32_t n_prompt = 128;
- constexpr uint32_t n_continue = 32;
-
- auto cparams = common_context_params_to_llama(params);
- cparams.n_seq_max = n_seqs;
- cparams.n_ctx = 512;
- cparams.n_batch = 256;
- cparams.n_ubatch = 64;
- cparams.kv_unified = true; // only a unified cache shares cells on seq_cp
-
- llama_context_ptr ctx = init_ctx(model, cparams, fill);
- if (!ctx) {
- LOG_ERR("%s: failed to init context\n", __func__);
- return test_status::FAIL;
- }
-
- const auto tok = [&](uint32_t seq, llama_pos pos) {
- return (llama_token) ((7*(uint32_t) pos + 31*seq + 1) % (uint32_t) n_vocab);
- };
-
- {
- common_batch batch(ctx.get());
- for (llama_pos pos = 0; pos < (llama_pos) n_prompt; ++pos) {
- batch.add(tok(0, pos), pos, 0, false);
- }
- if (llama_process(ctx.get(), LLAMA_PROCESS_TYPE_DECODE, batch.get()) != 0) {
- LOG_ERR("%s: prompt decode failed\n", __func__);
- return test_status::FAIL;
- }
- }
-
- // this is what llama-batched-bench does for -pps
- llama_memory_seq_cp(llama_get_memory(ctx.get()), 0, 1, -1, -1);
-
- for (uint32_t i = 0; i < n_continue; ++i) {
- const llama_pos pos = (llama_pos) (n_prompt + i);
-
- common_batch batch(ctx.get());
- for (uint32_t s = 0; s < n_seqs; ++s) {
- batch.add(tok(s, pos), pos, (llama_seq_id) s, true);
- }
- if (llama_process(ctx.get(), LLAMA_PROCESS_TYPE_DECODE, batch.get()) != 0) {
- LOG_ERR("%s: shared-seq decode failed at step %u\n", __func__, i);
- return test_status::FAIL;
- }
-
- for (uint32_t s = 0; s < n_seqs; ++s) {
- const float * logits = llama_get_logits_ith(ctx.get(), (int) s);
- if (logits == nullptr) {
- LOG_ERR("%s: missing shared-seq logits at index %u\n", __func__, s);
- return test_status::FAIL;
- }
- for (int t = 0; t < n_vocab; ++t) {
- if (!std::isfinite(logits[t])) {
- LOG_ERR("%s: non-finite shared-seq logit at step %u, seq %u, index %d\n", __func__, i, s, t);
- return test_status::FAIL;
- }
- }
- }
- }
-
- LOG_INF("%s: shared-seq decode succeeded (%u tokens after seq_cp)\n", __func__, n_continue*n_seqs);
- return test_status::PASS;
-}
-
-static test_status merge_status(test_status a, test_status b) {
- if (a == test_status::FAIL || b == test_status::FAIL) {
- return test_status::FAIL;
- }
- if (a == test_status::PASS || b == test_status::PASS) {
- return test_status::PASS;
- }
- return test_status::SKIP;
-}
-
-struct test_results {
- test_status rollback = test_status::SKIP;
- test_status replay = test_status::SKIP;
- test_status shared = test_status::SKIP;
-};
-
-// Run every test for an initialized model over both cache fills.
-static test_results run_tests(const common_params & params, llama_model * model) {
- test_results res;
- for (uint8_t fill : { 0, 0x3e }) {
- LOG_INF("%s: testing with cache fill 0x%02x\n", __func__, fill);
- const test_status rb = test_rollback(params, model, fill);
- const test_status rp = test_multi_seq_split_replay(params, model, fill);
- const test_status ss = test_shared_seq_reserve(params, model, fill);
- res.rollback = merge_status(res.rollback, rb);
- res.replay = merge_status(res.replay, rp);
- res.shared = merge_status(res.shared, ss);
- if (rb == test_status::FAIL || rp == test_status::FAIL || ss == test_status::FAIL) {
- break;
- }
- }
- return res;
-}
-
-// Run the tests for a single model file.
-// Returns the per-test statuses.
-static test_results run_tests_for_model(const std::string & model_path, const struct common_params & base_params) {
- struct common_params params = base_params;
- params.model.path = model_path;
-
- auto llama_init = common_init_from_params(params, true);
- auto * model = llama_init->model();
-
- if (model == nullptr) {
- LOG_ERR("%s: failed to init model '%s'\n", __func__, model_path.c_str());
- // a model that cannot be loaded is a failure, not a skip
- return { test_status::FAIL, test_status::FAIL, test_status::FAIL };
- }
-
- if (!llama_model_is_recurrent(model) && !llama_model_is_hybrid(model)) {
- LOG_INF("%s: skipping for non-recurrent model\n", __func__);
- return {};
- }
-
- return run_tests(params, model);
-}
-
-static void print_usage(int /* argc */, char ** argv) {
- LOG("\nexample usage:\n");
- LOG("\n %s -m your_model.gguf\n", argv[0]);
- LOG("\n %s --models tests/test-models\n", argv[0]);
- LOG("\n");
-}
-
-int main(int argc, char ** argv) {
- std::setlocale(LC_NUMERIC, "C");
-
- common_params params;
- params.sampling.seed = 1234;
- params.n_predict = 1;
-
- common_init();
-
- // extract our own --models DIR option before handing the rest to the common arg parser
- std::string models_dir;
- std::vector<char *> filtered_argv;
- filtered_argv.push_back(argv[0]);
- for (int i = 1; i < argc; i++) {
- if (strcmp(argv[i], "--models") == 0) {
- if (i + 1 >= argc) {
- LOG_ERR("%s: --models requires a directory argument\n", __func__);
- return 1;
- }
- models_dir = argv[i + 1];
- i++;
- } else {
- filtered_argv.push_back(argv[i]);
- }
- }
- filtered_argv.push_back(nullptr);
- const int fargc = (int)filtered_argv.size() - 1;
-
- // in --models mode there is no single model; set a placeholder so the common parser's
- // "--model is required" check passes (each model is set individually inside the loop)
- if (!models_dir.empty()) {
- params.model.path = models_dir;
- }
-
- if (!common_params_parse(fargc, filtered_argv.data(), params, LLAMA_EXAMPLE_COMMON, print_usage)) {
- return 1;
- }
-
- llama_backend_init();
-
- if (!models_dir.empty()) {
- // run every test over each dummy model in the directory
- if (!std::filesystem::exists(models_dir) || !std::filesystem::is_directory(models_dir)) {
- LOG_ERR("%s: models directory '%s' does not exist\n", __func__, models_dir.c_str());
- return 1;
- }
-
- std::vector<std::string> models;
- for (const auto & entry : std::filesystem::directory_iterator(models_dir)) {
- if (entry.is_regular_file() && entry.path().extension() == ".gguf") {
- models.push_back(entry.path().string());
- }
- }
- std::sort(models.begin(), models.end());
-
- if (models.empty()) {
- LOG_ERR("%s: no .gguf models found in '%s'\n", __func__, models_dir.c_str());
- return 1;
- }
-
- size_t name_width = 5; // "Model"
- for (const auto & model_path : models) {
- name_width = std::max(name_width, std::filesystem::path(model_path).filename().string().size());
- }
-
- // silence everything but the table itself (LOG has verbosity LOG_LEVEL_OUTPUT = 0)
- common_log_set_verbosity_thold(0);
-
- LOG("%-*s %-8s %-11s %s\n", (int) name_width, "Model", "rollback", "split replay", "shared seq");
- common_log_flush(common_log_main());
-
- size_t n_pass[3] = { 0, 0, 0 };
- size_t n_skip[3] = { 0, 0, 0 };
- size_t n_fail[3] = { 0, 0, 0 };
- for (const auto & model_path : models) {
- const auto name = std::filesystem::path(model_path).filename().string();
-
- LOG("%-*s", (int) name_width, name.c_str());
-
- const test_results res = run_tests_for_model(model_path, params);
-
- // all status strings have the same raw length, so the columns line up;
- // pad the first status to the width of the "rollback" header + separator
- LOG(" %s %s %s", test_status_str(res.rollback), test_status_str(res.replay), test_status_str(res.shared));
- LOG("\n");
- common_log_flush(common_log_main());
-
- const test_status all[3] = { res.rollback, res.replay, res.shared };
- for (int t = 0; t < 3; ++t) {
- switch (all[t]) {
- case test_status::PASS: n_pass[t]++; break;
- case test_status::FAIL: n_fail[t]++; break;
- case test_status::SKIP: n_skip[t]++; break;
- }
- }
- }
-
- common_log_set_verbosity_thold(LOG_DEFAULT_LLAMA);
- common_log_flush(common_log_main());
-
- LOG_INF("%s: rollback: %zu passed, %zu skipped, %zu failed (of %zu)\n",
- __func__, n_pass[0], n_skip[0], n_fail[0], models.size());
- LOG_INF("%s: split replay: %zu passed, %zu skipped, %zu failed (of %zu)\n",
- __func__, n_pass[1], n_skip[1], n_fail[1], models.size());
- LOG_INF("%s: shared seq: %zu passed, %zu skipped, %zu failed (of %zu)\n",
- __func__, n_pass[2], n_skip[2], n_fail[2], models.size());
-
- return (n_fail[0] + n_fail[1] + n_fail[2]) == 0 ? 0 : 1;
- }
-
- // single-model mode
- const test_results res = run_tests_for_model(params.model.path, params);
-
- return (res.rollback == test_status::FAIL || res.replay == test_status::FAIL || res.shared == test_status::FAIL) ? 1 : 0;
-}
diff --git a/tests/test-save-load-state.cpp b/tests/test-save-load-state.cpp
deleted file mode 100644
index 43a513ae5..000000000
--- a/tests/test-save-load-state.cpp
+++ /dev/null
@@ -1,1069 +0,0 @@
-// TODO: merge with test-recurrent-state-rollback.cpp
-// TODO: merge with test-state-restore-fragmented.cpp
-
-#include "arg.h"
-#include "common.h"
-#include "log.h"
-#include "llama-cpp.h"
-
-#include <algorithm>
-#include <clocale>
-#include <cmath>
-#include <cstdio>
-#include <cstring>
-#include <filesystem>
-#include <fstream>
-#include <functional>
-#include <iterator>
-#include <random>
-#include <string>
-#include <utility>
-#include <vector>
-
-constexpr double NMSE_THRESHOLD = 1e-5;
-
-enum class test_status {
- PASS,
- FAIL,
- SKIP,
-};
-
-static const char * test_status_str(test_status status) {
- switch (status) {
- case test_status::PASS: return "\033[1;32mPASS\033[0m";
- case test_status::FAIL: return "\033[1;31mFAIL\033[0m";
- case test_status::SKIP: return "\033[1;33mSKIP\033[0m";
- }
- return "";
-}
-
-// normalized mean squared error = mse(a, b) / mse(a, 0)
-static double nmse(const std::vector<float> & a, const std::vector<float> & b) {
- GGML_ASSERT(a.size() == b.size());
- double mse_a_b = 0.0;
- double mse_a_0 = 0.0;
-
- for (size_t i = 0; i < a.size(); i++) {
- const float a_i = a[i];
- const float b_i = b[i];
-
- mse_a_b += (double) (a_i - b_i) * (a_i - b_i);
- mse_a_0 += (double) a_i * a_i;
- }
-
- return mse_a_b / mse_a_0;
-}
-
-struct generation_result {
- llama_tokens tokens;
- std::vector<std::vector<float>> logits;
-
- bool empty() const { return tokens.empty(); }
-};
-
-static bool get_current_logits(llama_context * ctx, std::vector<float> & out) {
- const auto * vocab = llama_model_get_vocab(llama_get_model(ctx));
- const int32_t n_vocab = llama_vocab_n_tokens(vocab);
- const float * logits = llama_get_logits_ith(ctx, -1);
- if (logits == nullptr) {
- return false;
- }
- out.assign(logits, logits + n_vocab);
- return true;
-}
-
-static generation_result generate_tokens(llama_context * ctx, llama_sampler * smpl, int & n_past, int32_t n_predict, llama_seq_id seq_id) {
- generation_result result;
- common_batch batch(ctx);
-
- for (int i = 0; i < n_predict; i++) {
- std::vector<float> logits;
- if (!get_current_logits(ctx, logits)) {
- LOG_ERR("\n%s: failed to get logits\n", __func__);
- return {};
- }
-
- auto next_token = llama_sampler_sample(smpl, ctx, -1);
-
- // LOG_LEVEL_INFO gate: visible in single-model mode, silenced in --models table mode
- LOGV(LOG_LEVEL_INFO, "%d ", next_token);
- result.tokens.push_back(next_token);
- result.logits.push_back(std::move(logits));
-
- batch.clear();
- batch.add(next_token, n_past, seq_id, true);
-
- if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get())) {
- LOG_ERR("\n%s: failed to evaluate\n", __func__);
- return {};
- }
- n_past++;
- }
-
- return result;
-}
-
-static bool generate_tokens_compare(
- llama_context * ctx, llama_sampler * smpl, int & n_past, int32_t n_predict, llama_seq_id seq_id,
- const generation_result & expected) {
- if (expected.tokens.size() != expected.logits.size() || expected.tokens.size() < (size_t) n_predict) {
- LOG_ERR("\n%s: invalid expected generation\n", __func__);
- return false;
- }
-
- common_batch batch(ctx);
-
- for (int i = 0; i < n_predict; i++) {
- std::vector<float> logits;
- if (!get_current_logits(ctx, logits)) {
- LOG_ERR("\n%s: failed to get logits\n", __func__);
- return false;
- }
- if (logits.size() != expected.logits[i].size()) {
- LOG_ERR("\n%s: logits size mismatch at step %d: %zu != %zu\n", __func__, i, logits.size(), expected.logits[i].size());
- return false;
- }
-
- const double nmse_val = nmse(expected.logits[i], logits);
- LOG_TRC("%s: step %d nmse = %.6e\n", __func__, i, nmse_val);
- if (nmse_val > NMSE_THRESHOLD) {
- LOG_ERR("\n%s: error: NMSE at step %d is %.6e (threshold %.1e)\n", __func__, i, nmse_val, NMSE_THRESHOLD);
- return false;
- }
-
- const auto next_token = llama_sampler_sample(smpl, ctx, -1);
- const auto expected_token = expected.tokens[i];
-
- LOGV(LOG_LEVEL_INFO, "%d ", next_token);
- if (next_token != expected_token) {
- LOG_TRC("%s: sampled token %d differs from expected %d, using expected token\n", __func__, next_token, expected_token);
- }
-
- batch.clear();
- batch.add(expected_token, n_past, seq_id, true);
-
- if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get())) {
- LOG_ERR("\n%s: failed to evaluate\n", __func__);
- return false;
- }
- n_past++;
- }
-
- return true;
-}
-
-// Test 1: baseline
-// - decode all but the last token
-// - save state to disk
-// - decode the last token
-// - generate n_predict tokens
-static generation_result test_baseline(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens) {
- auto params_ctx = common_context_params_to_llama(params);
- params_ctx.n_seq_max = 2;
- auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)};
-
- auto sparams = llama_sampler_chain_default_params();
- auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)};
- llama_sampler_chain_add(smpl.get(), llama_sampler_init_dist(params.sampling.seed));
-
- auto n_past = 0;
- if (!common_prompt_batch_decode(ctx.get(), tokens, (int)tokens.size(), n_past, params.n_batch, params.out_file, true)) {
- LOG_ERR("%s: failed to decode prompt\n", __func__);
- return {};
- }
-
- LOGV(LOG_LEVEL_INFO, "\n=== Test 1: baseline ===\n");
-
- auto result = generate_tokens(ctx.get(), smpl.get(), n_past, params.n_predict, 0);
- if (result.empty()) {
- return {};
- }
-
- LOGV(LOG_LEVEL_INFO, "\n");
-
- return result;
-}
-
-
-// Test 2: sequence removal isolation
-// - decode the same prefix into two sequences
-// - remove sequence 0
-// - verify that sequence 1 remains unchanged
-static bool test_seq_rm_isolated(
- struct llama_model * model,
- const struct common_params & params,
- const llama_tokens & tokens) {
- auto params_ctx = common_context_params_to_llama(params);
- params_ctx.n_ctx = 256;
- params_ctx.n_seq_max = 2;
- params_ctx.kv_unified = true;
-
- auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)};
- if (!ctx) {
- LOG_ERR("%s: failed to create context\n", __func__);
- return false;
- }
-
- LOGV(LOG_LEVEL_INFO, "\n=== Test 2: sequence removal isolation ===\n");
-
- const size_t n_tokens = tokens.size() < 128 ? tokens.size() : 128;
- for (llama_seq_id seq_id = 0; seq_id < 2; ++seq_id) {
- common_batch batch(ctx.get());
- for (size_t i = 0; i < n_tokens; ++i) {
- batch.add(tokens[i], i, seq_id, i == n_tokens - 1);
- }
-
- if (llama_process(ctx.get(), LLAMA_PROCESS_TYPE_DECODE, batch.get())) {
- LOG_ERR("%s: failed to decode prompt for sequence %d\n", __func__, seq_id);
- return false;
- }
- }
-
- const auto get_seq_state = [&](llama_seq_id seq_id, std::vector<uint8_t> & state) {
- const size_t state_size = llama_state_seq_get_size(ctx.get(), seq_id);
- if (state_size == 0) {
- LOG_ERR("%s: sequence state is empty\n", __func__);
- return false;
- }
-
- state.resize(state_size);
- const size_t ncopy = llama_state_seq_get_data(ctx.get(), state.data(), state.size(), seq_id);
- if (ncopy != state.size()) {
- LOG_ERR("%s: sequence state length %zu does not match expected length %zu\n",
- __func__, ncopy, state.size());
- return false;
- }
-
- return true;
- };
-
- std::vector<uint8_t> state_before;
- if (!get_seq_state(1, state_before)) {
- return false;
- }
-
- if (!llama_memory_seq_rm(llama_get_memory(ctx.get()), 0, -1, -1)) {
- LOG_ERR("%s: failed to remove sequence 0\n", __func__);
- return false;
- }
-
- std::vector<uint8_t> state_after;
- if (!get_seq_state(1, state_after)) {
- return false;
- }
-
- if (state_before != state_after) {
- LOG_ERR("%s: removing sequence 0 changed sequence 1\n", __func__);
- return false;
- }
-
- LOGV(LOG_LEVEL_INFO, "PASS\n");
- return true;
-}
-
-
-// Test 3: state load
-// - create a new context
-// - load state from file
-// - replay the last prompt token
-// - generate n_predict tokens and compare against expected result
-static bool test_state_load(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens, const generation_result & expected_result) {
- auto params_ctx = common_context_params_to_llama(params);
- params_ctx.n_seq_max = 2;
- auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)};
-
- auto sparams = llama_sampler_chain_default_params();
- auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)};
- llama_sampler_chain_add(smpl.get(), llama_sampler_init_dist(params.sampling.seed));
-
- LOGV(LOG_LEVEL_INFO, "\n=== Test 3: state load ===\n");
-
- // Load state from file
- llama_tokens unused_sts(tokens.size());
- size_t n_token_count_out = 0;
-
- if (!llama_state_load_file(ctx.get(), params.out_file.data(), unused_sts.data(), unused_sts.size(), &n_token_count_out)) {
- LOG_ERR("\n%s: failed to load state\n", __func__);
- return false;
- }
-
- LOG_TRC("%s: loaded state with %zu tokens\n", __func__, n_token_count_out);
-
- // Replay last token
- int n_past = (int) n_token_count_out - 1;
- if (!common_replay_last_token(ctx.get(), tokens.back(), n_past)) {
- return false;
- }
- n_past++;
-
- // Generate tokens and compare logits against the baseline
- if (!generate_tokens_compare(ctx.get(), smpl.get(), n_past, params.n_predict, 0, expected_result)) {
- return false;
- }
-
- LOGV(LOG_LEVEL_INFO, "\nPASS\n");
- return true;
-}
-
-
-// Test 4: seq copy (host)
-// - create a multi-seq context
-// - load state from file
-// - replay the last prompt token
-// - migrate KV cache from seq 0 to seq 1 via the CPU path
-// - generate n_predict tokens on seq 1 and compare against expected result
-static bool test_seq_cp_host(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens, const generation_result & expected_result) {
- auto params_ctx = common_context_params_to_llama(params);
- params_ctx.n_seq_max = 2;
- auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)};
-
- auto sparams = llama_sampler_chain_default_params();
- auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)};
- llama_sampler_chain_add(smpl.get(), llama_sampler_init_dist(params.sampling.seed));
-
- LOGV(LOG_LEVEL_INFO, "\n=== Test 4: seq copy (host) ===\n");
-
- // Load state from file
- llama_tokens unused_sts(tokens.size());
- size_t n_token_count_out = 0;
-
- if (!llama_state_load_file(ctx.get(), params.out_file.data(), unused_sts.data(), unused_sts.size(), &n_token_count_out)) {
- LOG_ERR("\n%s: failed to load state\n", __func__);
- return false;
- }
-
- LOG_TRC("%s: loaded state with %zu tokens\n", __func__, n_token_count_out);
-
- // Replay last token
- int n_past = (int) n_token_count_out - 1;
- if (!common_replay_last_token(ctx.get(), tokens.back(), n_past)) {
- return false;
- }
- n_past++;
-
- // Migrate KV cache from seq 0 to seq 1 (CPU path)
- {
- std::vector<uint8_t> seq_store(llama_state_seq_get_size(ctx.get(), 0));
- const size_t ncopy = llama_state_seq_get_data(ctx.get(), seq_store.data(), seq_store.size(), 0);
- if (ncopy != seq_store.size()) {
- LOG_ERR("\n%s: seq copy data length %zd does not match expected length %zd\n", __func__, ncopy, seq_store.size());
- return false;
- }
- LOG_TRC("%s: seq 0 copied, %zd bytes\n", __func__, ncopy);
-
- llama_memory_clear(llama_get_memory(ctx.get()), true);
- LOG_TRC("%s: kv cache cleared\n", __func__);
-
- const size_t nset = llama_state_seq_set_data(ctx.get(), seq_store.data(), seq_store.size(), 1);
- if (nset != seq_store.size()) {
- LOG_ERR("\n%s: seq set data length %zd does not match expected length %zd\n", __func__, nset, seq_store.size());
- return false;
- }
- LOG_TRC("%s: seq 1 restored, %zd bytes\n", __func__, nset);
- }
-
- // Generate tokens and compare logits against the baseline
- if (!generate_tokens_compare(ctx.get(), smpl.get(), n_past, params.n_predict, 1, expected_result)) {
- return false;
- }
-
- LOGV(LOG_LEVEL_INFO, "\nPASS\n");
- return true;
-}
-
-
-// Test 5: seq copy (device)
-// - create a multi-seq context
-// - load state from file
-// - replay the last prompt token
-// - migrate KV cache from seq 0 to seq 1 via the on-device path
-// - generate n_predict tokens on seq 1 and compare against expected result
-static bool test_seq_cp_device(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens, const generation_result & expected_result) {
- auto params_ctx = common_context_params_to_llama(params);
- params_ctx.n_seq_max = 2;
- auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)};
-
- auto sparams = llama_sampler_chain_default_params();
- auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)};
- llama_sampler_chain_add(smpl.get(), llama_sampler_init_dist(params.sampling.seed));
-
- LOGV(LOG_LEVEL_INFO, "\n=== Test 5: seq copy (device) ===\n");
-
- // Load state from file
- llama_tokens unused_sts(tokens.size());
- size_t n_token_count_out = 0;
-
- if (!llama_state_load_file(ctx.get(), params.out_file.data(), unused_sts.data(), unused_sts.size(), &n_token_count_out)) {
- LOG_ERR("\n%s: failed to load state\n", __func__);
- return false;
- }
-
- LOG_TRC("%s: loaded state with %zu tokens\n", __func__, n_token_count_out);
-
- // Replay last token
- int n_past = (int) n_token_count_out - 1;
- if (!common_replay_last_token(ctx.get(), tokens.back(), n_past)) {
- return false;
- }
- n_past++;
-
- // Migrate KV cache from seq 0 to seq 1 (on-device path)
- {
- std::vector<uint8_t> seq_store(llama_state_seq_get_size_ext(ctx.get(), 0, LLAMA_STATE_SEQ_FLAGS_ON_DEVICE));
- const size_t ncopy = llama_state_seq_get_data_ext(ctx.get(), seq_store.data(), seq_store.size(), 0, LLAMA_STATE_SEQ_FLAGS_ON_DEVICE);
- if (ncopy != seq_store.size()) {
- LOG_ERR("\n%s: seq copy data length %zd does not match expected length %zd\n", __func__, ncopy, seq_store.size());
- return false;
- }
- LOG_TRC("%s: seq 0 copied, %zd bytes\n", __func__, ncopy);
-
- llama_memory_clear(llama_get_memory(ctx.get()), true);
- LOG_TRC("%s: kv cache cleared\n", __func__);
-
- const size_t nset = llama_state_seq_set_data_ext(ctx.get(), seq_store.data(), seq_store.size(), 1, LLAMA_STATE_SEQ_FLAGS_ON_DEVICE);
- if (nset != seq_store.size()) {
- LOG_ERR("\n%s: seq set data length %zd does not match expected length %zd\n", __func__, nset, seq_store.size());
- return false;
- }
- LOG_TRC("%s: seq 1 restored, %zd bytes\n", __func__, nset);
- }
-
- // Generate tokens and compare logits against the baseline
- if (!generate_tokens_compare(ctx.get(), smpl.get(), n_past, params.n_predict, 1, expected_result)) {
- return false;
- }
-
- LOGV(LOG_LEVEL_INFO, "\nPASS\n");
- return true;
-}
-
-
-// Test 6/7: seq copy (scatter)
-// - decode the same prefix on two sequences, interleaving seq 0 cells between the seq 1 cells
-// - save the seq 1 state, free the interleaved seq 0 cells, and restore via the given io path
-// - the restore destination is non-contiguous: scatter reads are batched per contiguous run
-// - save again on the host and compare the two blobs byte for byte
-static bool test_seq_cp_scatter(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens, int test_num, bool on_device) {
- auto params_ctx = common_context_params_to_llama(params);
- params_ctx.n_ctx = 256;
- params_ctx.n_seq_max = 2;
- params_ctx.kv_unified = true;
- auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)};
-
- LOGV(LOG_LEVEL_INFO, "\n=== Test %d: seq copy (%s, scatter) ===\n", test_num, on_device ? "device" : "host");
-
- const uint32_t flags = on_device ? LLAMA_STATE_SEQ_FLAGS_ON_DEVICE : LLAMA_STATE_SEQ_FLAGS_NONE;
-
- auto decode_one = [&](llama_token tok, int pos, llama_seq_id seq) {
- common_batch batch(ctx.get());
- batch.add(tok, pos, seq, true);
- return llama_process(ctx.get(), LLAMA_PROCESS_TYPE_DECODE, batch.get()) == 0;
- };
-
- // seq 0 cells 0,1,4 interleave the seq 1 cells 2,3,5
- if (!decode_one(tokens[0], 0, 0) ||
- !decode_one(tokens[1], 1, 0) ||
- !decode_one(tokens[0], 0, 1) ||
- !decode_one(tokens[1], 1, 1) ||
- !decode_one(tokens[2], 2, 0) ||
- !decode_one(tokens[2], 2, 1)) {
- LOG_ERR("%s: failed to build interleaved state\n", __func__);
- return false;
- }
-
- const auto get_seq_state = [&](llama_seq_id seq_id, uint32_t fl, std::vector<uint8_t> & state) {
- const size_t state_size = llama_state_seq_get_size_ext(ctx.get(), seq_id, fl);
- if (state_size == 0) {
- LOG_ERR("%s: sequence state is empty\n", __func__);
- return false;
- }
-
- state.resize(state_size);
- const size_t ncopy = llama_state_seq_get_data_ext(ctx.get(), state.data(), state.size(), seq_id, fl);
- if (ncopy != state.size()) {
- LOG_ERR("%s: sequence state length %zu does not match expected length %zu\n",
- __func__, ncopy, state.size());
- return false;
- }
-
- return true;
- };
-
- // host blob: contains the KV data, used for the byte-for-byte comparison
- std::vector<uint8_t> state_before;
- if (!get_seq_state(1, LLAMA_STATE_SEQ_FLAGS_NONE, state_before)) {
- return false;
- }
-
- // save via the io path under test
- std::vector<uint8_t> state_save;
- if (!get_seq_state(1, flags, state_save)) {
- return false;
- }
- LOG_TRC("%s: seq 1 saved via %s, %zu bytes\n", __func__, on_device ? "device" : "host", state_save.size());
-
- // free seq 0's cells so the ring is fragmented: the restore destination (seq 1's interleaved cells) stays non-contiguous
- if (!llama_memory_seq_rm(llama_get_memory(ctx.get()), 0, -1, -1)) {
- LOG_ERR("%s: failed to remove sequence 0\n", __func__);
- return false;
- }
-
- // restore via the io path under test
- const size_t nset = llama_state_seq_set_data_ext(ctx.get(), state_save.data(), state_save.size(), 1, flags);
- if (nset != state_save.size()) {
- LOG_ERR("%s: seq set data length %zu does not match expected length %zu\n", __func__, nset, state_save.size());
- return false;
- }
- LOG_TRC("%s: seq 1 restored via %s, %zu bytes\n", __func__, on_device ? "device" : "host", nset);
-
- std::vector<uint8_t> state_after;
- if (!get_seq_state(1, LLAMA_STATE_SEQ_FLAGS_NONE, state_after)) {
- return false;
- }
-
- // the blob is serialized in sequence cell order, so identical bytes iff the restore wrote the same KV
- if (state_before.size() != state_after.size() || memcmp(state_before.data(), state_after.data(), state_before.size()) != 0) {
- LOG_ERR("\n%s: error: restored KV state is not byte-identical to the saved state\n", __func__);
- return false;
- }
-
- LOGV(LOG_LEVEL_INFO, "\nPASS\n");
- return true;
-}
-
-
-// Test 8: state blob round-trip
-// compares blobs rather than generated text: a partially restored cell still decodes to plausible tokens
-static bool test_state_roundtrip(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens) {
- auto params_ctx = common_context_params_to_llama(params);
- auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)};
-
- LOGV(LOG_LEVEL_INFO, "\n=== Test 8: state blob round-trip ===\n");
-
- common_batch batch = common_batch_get_one(ctx.get(), tokens);
- if (llama_process(ctx.get(), LLAMA_PROCESS_TYPE_DECODE, batch.get())) {
- LOG_ERR("\n%s: failed to decode prompt\n", __func__);
- return false;
- }
-
- std::vector<uint8_t> blob_a(llama_state_seq_get_size(ctx.get(), 0));
- const size_t n_a = llama_state_seq_get_data(ctx.get(), blob_a.data(), blob_a.size(), 0);
- if (n_a != blob_a.size()) {
- LOG_ERR("\n%s: saved %zu bytes, expected %zu\n", __func__, n_a, blob_a.size());
- return false;
- }
-
- if (!llama_memory_seq_rm(llama_get_memory(ctx.get()), 0, -1, -1)) {
- LOG_ERR("\n%s: failed to erase seq 0\n", __func__);
- return false;
- }
-
- if (llama_state_seq_set_data(ctx.get(), blob_a.data(), blob_a.size(), 0) != blob_a.size()) {
- LOG_ERR("\n%s: failed to restore seq 0\n", __func__);
- return false;
- }
-
- std::vector<uint8_t> blob_b(llama_state_seq_get_size(ctx.get(), 0));
- const size_t n_b = llama_state_seq_get_data(ctx.get(), blob_b.data(), blob_b.size(), 0);
- if (n_b != n_a) {
- LOG_ERR("\n%s: re-saved %zu bytes, expected %zu\n", __func__, n_b, n_a);
- return false;
- }
-
- size_t n_diff = 0;
- size_t i_diff = 0;
- for (size_t i = 0; i < n_a; i++) {
- if (blob_a[i] != blob_b[i]) {
- if (n_diff == 0) {
- i_diff = i;
- }
- n_diff++;
- }
- }
-
- if (n_diff > 0) {
- LOG_ERR("\n%s: state changed across a restore: %zu of %zu bytes differ, first at offset %zu\n",
- __func__, n_diff, n_a, i_diff);
- return false;
- }
-
- LOGV(LOG_LEVEL_INFO, "\nPASS\n");
- return true;
-}
-
-
-// overwrite the tensor data with 0xff bytes (NaN when read as f16/f32), so that the restore fails
-static bool corrupt_state(std::vector<uint8_t> & data) {
- if (data.size() < 3*4096) {
- LOG_ERR("%s: state of %zu bytes is too small to corrupt\n", __func__, data.size());
- return false;
- }
-
- std::fill(data.begin() + 4096, data.end() - data.size()/4, 0xff);
- return true;
-}
-
-
-// Test 9: state restore failure
-// a failed restore must leave the sequence empty and must not change the logits of other sequences
-static bool test_state_restore_failure(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens) {
- auto params_ctx = common_context_params_to_llama(params);
- params_ctx.n_ctx = 256;
- params_ctx.n_seq_max = 4;
- params_ctx.kv_unified = true;
-
- // without flash attention, corrupted data left behind by the restore shows up as NaN logits on the other sequences
- params_ctx.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED;
-
- auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)};
- if (!ctx) {
- LOG_ERR("%s: failed to create context\n", __func__);
- return false;
- }
-
- LOGV(LOG_LEVEL_INFO, "\n=== Test 9: state restore failure ===\n");
-
- llama_memory_t mem = llama_get_memory(ctx.get());
- if (mem == nullptr) {
- LOGV(LOG_LEVEL_INFO, "PASS (model has no memory)\n");
- return true;
- }
-
- const auto decode = [&](const llama_tokens & inp, llama_seq_id seq_id, std::vector<float> * logits_out) {
- common_batch batch(ctx.get());
- for (size_t i = 0; i < inp.size(); ++i) {
- batch.add(inp[i], i, seq_id, i == inp.size() - 1);
- }
-
- if (llama_process(ctx.get(), LLAMA_PROCESS_TYPE_DECODE, batch.get())) {
- LOG_ERR("%s: failed to decode on sequence %d\n", __func__, seq_id);
- return false;
- }
-
- if (logits_out && !get_current_logits(ctx.get(), *logits_out)) {
- LOG_ERR("%s: failed to get logits\n", __func__);
- return false;
- }
-
- return true;
- };
-
- const llama_tokens tokens_save (tokens.begin(), tokens.begin() + std::min<size_t>(24, tokens.size()));
- const llama_tokens tokens_verify(tokens.end() - std::min<size_t>(8, tokens.size()), tokens.end());
-
- // the registered tests share a working directory, so the state file is named after the model
- const std::string path = "state-restore-failure." + std::filesystem::path(params.model.path).filename().string() + ".tmp.bin";
-
- llama_memory_clear(mem, true);
-
- std::vector<float> baseline;
- if (!decode(tokens_verify, 1, &baseline)) {
- return false;
- }
-
- const std::vector<std::pair<const char *, std::function<bool()>>> cases = {
- { "buffer", [&]() {
- std::vector<uint8_t> state(llama_state_seq_get_size(ctx.get(), 0));
- GGML_ASSERT(llama_state_seq_get_data(ctx.get(), state.data(), state.size(), 0) == state.size());
- llama_memory_seq_rm(mem, 0, -1, -1);
-
- if (!corrupt_state(state)) {
- return false;
- }
-
- return llama_state_seq_set_data(ctx.get(), state.data(), state.size(), 0) == 0;
- }},
- { "file", [&]() {
- GGML_ASSERT(llama_state_seq_save_file(ctx.get(), path.c_str(), 0, tokens_save.data(), tokens_save.size()) > 0);
- llama_memory_seq_rm(mem, 0, -1, -1);
-
- std::vector<uint8_t> data;
- {
- std::ifstream f(path, std::ios::binary);
- data.assign(std::istreambuf_iterator<char>(f), std::istreambuf_iterator<char>());
- }
-
- if (!corrupt_state(data)) {
- std::remove(path.c_str());
- return false;
- }
-
- {
- std::ofstream f(path, std::ios::binary);
- f.write((const char *) data.data(), data.size());
- }
-
- llama_tokens tokens_out(tokens_save.size());
- size_t n_token_count = 0;
- const size_t nread = llama_state_seq_load_file(ctx.get(), path.c_str(), 0, tokens_out.data(), tokens_out.size(), &n_token_count);
- std::remove(path.c_str());
-
- return nread == 0;
- }},
- };
-
- for (const auto & [name, restore_failed] : cases) {
- llama_memory_clear(mem, true);
-
- if (!decode(tokens_save, 0, nullptr)) {
- return false;
- }
-
- if (!restore_failed()) {
- LOG_ERR("%s: %s: restoring a corrupted state did not fail\n", __func__, name);
- return false;
- }
-
- if (llama_memory_seq_pos_max(mem, 0) != -1) {
- LOG_ERR("%s: %s: sequence not empty after failed restore\n", __func__, name);
- return false;
- }
-
- std::vector<float> logits;
- if (!decode(tokens_verify, 1, &logits)) {
- return false;
- }
-
- float diff_max = 0.0f;
- size_t n_nan = 0;
- for (size_t i = 0; i < logits.size(); ++i) {
- if (std::isnan(logits[i]) || std::isnan(baseline[i])) {
- n_nan++;
- } else {
- diff_max = std::max(diff_max, std::fabs(logits[i] - baseline[i]));
- }
- }
-
- if (n_nan > 0 || diff_max > 1e-6f) {
- LOG_ERR("%s: %s: logits changed after failed restore (max diff = %g, nan = %zu)\n", __func__, name, diff_max, n_nan);
- return false;
- }
-
- LOG_TRC("%s: %s: logits match (max diff = %g)\n", __func__, name, diff_max);
- }
-
- LOGV(LOG_LEVEL_INFO, "\nPASS\n");
- return true;
-}
-
-
-// Test 10: state rotation
-// a KV state saved with attention rotation enabled must restore only into a context with the same setting;
-// note: rotation is only active for quantized KV caches with a head size that is a multiple of 64,
-// for other models the restore into the rotation-disabled context is valid and the test passes vacuously
-static bool test_state_rotation(struct llama_model * model, const struct common_params & params) {
- LOGV(LOG_LEVEL_INFO, "\n=== Test 10: state rotation ===\n");
-
- const std::string attn_rot_disable = common_get_env("LLAMA_ATTN_ROT_DISABLE");
- const auto make_context = [&](ggml_type type_k, ggml_type type_v, bool disable_rotation) {
- common_set_env("LLAMA_ATTN_ROT_DISABLE", disable_rotation ? "1" : "0");
- auto params_ctx = common_context_params_to_llama(params);
- params_ctx.n_ctx = 32;
- params_ctx.n_batch = 1;
- params_ctx.n_ubatch = 1;
- params_ctx.type_k = type_k;
- params_ctx.type_v = type_v;
- params_ctx.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED;
- return llama_context_ptr(llama_init_from_model(model, params_ctx));
- };
-
- std::vector<std::pair<ggml_type, ggml_type>> type_pairs;
- for (const auto & types : { std::pair{GGML_TYPE_Q8_0, GGML_TYPE_Q8_0} }) {
- if (make_context(types.first, types.second, false)) {
- type_pairs.push_back(types);
- }
- }
- if (type_pairs.empty()) {
- LOG_WRN("%s: no supported quantized KV cache type combination - skipping\n", __func__);
- return true;
- }
-
- bool success = true;
- for (const auto & types : type_pairs) {
- auto src = make_context(types.first, types.second, false);
- if (!src) {
- LOG_ERR("%s: failed to create source context\n", __func__);
- success = false;
- break;
- }
-
- llama_token token = 0;
- if (llama_decode(src.get(), llama_batch_get_one(&token, 1))) {
- LOG_ERR("%s: failed to decode token\n", __func__);
- success = false;
- break;
- }
-
- const size_t state_size = llama_state_seq_get_size(src.get(), 0);
- if (state_size == 0) {
- continue; // no KV state to test
- }
-
- std::vector<uint8_t> state(state_size);
- if (llama_state_seq_get_data(src.get(), state.data(), state.size(), 0) != state.size()) {
- LOG_ERR("%s: failed to save sequence state\n", __func__);
- success = false;
- break;
- }
-
- auto matching = make_context(types.first, types.second, false);
- if (!matching || llama_state_seq_set_data(matching.get(), state.data(), state.size(), 0) != state.size()) {
- LOG_ERR("%s: failed to restore matching rotation\n", __func__);
- success = false;
- break;
- }
-
- auto mismatched = make_context(types.first, types.second, true);
- if (!mismatched) {
- LOG_ERR("%s: failed to create mismatched rotation context\n", __func__);
- success = false;
- break;
- }
- if (llama_state_seq_set_data(mismatched.get(), state.data(), state.size(), 0) != 0) {
- LOG_TRC("%s: state restored into rotation-disabled context, model does not use attention rotation\n", __func__);
- }
- }
- common_set_env("LLAMA_ATTN_ROT_DISABLE", attn_rot_disable);
-
- if (!success) {
- return false;
- }
-
- LOGV(LOG_LEVEL_INFO, "\nPASS\n");
- return true;
-}
-
-struct test_suite {
- std::vector<test_status> results;
-
- bool all_passed() const {
- return std::all_of(results.begin(), results.end(), [](test_status s) { return s == test_status::PASS; });
- }
-};
-
-// column headers for the --models table, one per test, in the order they are run
-static const std::vector<const char *> test_names = {
- "baseline", "seq_rm", "state_load", "cp_h", "cp_d", "cp_h_s", "cp_d_s", "rt", "rf", "rot",
-};
-
-// Run the full save/load test suite (tests 1-10) for a single model.
-// Returns the per-test results.
-static test_suite run_save_load_tests_for_model(const std::string & model_path, const struct common_params & base_params) {
- test_suite suite;
-
- struct common_params params = base_params;
- params.model.path = model_path;
-
- auto llama_init = common_init_from_params(params, true);
- auto * model = llama_init->model();
-
- if (model == nullptr) {
- LOG_ERR("%s: failed to init model '%s'\n", __func__, model_path.c_str());
- suite.results.assign(test_names.size(), test_status::SKIP);
- return suite;
- }
-
- GGML_ASSERT(llama_init->context() == nullptr);
-
- // Tokenize prompt or generate random tokens
- llama_tokens tokens;
- if (params.prompt.empty()) {
- const int n_prompt = params.n_batch;
-
- // this path is useful for model files that do not have a tokenizer
- LOG_INF("%s: no prompt provided, generating %d (n_batch) random tokens\n", __func__, n_prompt);
-
- const auto * vocab = llama_model_get_vocab(model);
- const auto n_vocab = llama_vocab_n_tokens(vocab);
-
- std::mt19937 rng(params.sampling.seed);
- std::uniform_int_distribution<llama_token> dist(0, n_vocab - 1);
- for (int i = 0; i < n_prompt; i++) {
- tokens.push_back(dist(rng));
- }
- } else {
- LOG_INF("%s: tokenizing prompt '%s'\n", __func__, params.prompt.c_str());
-
- auto ctx = llama_context_ptr{llama_init_from_model(model, common_context_params_to_llama(params))};
- tokens = common_tokenize(ctx.get(), params.prompt, true);
- }
-
- LOG_INF("%s: the input prompt is %d tokens\n", __func__, (int)tokens.size());
-
- // Test 1: baseline (saves state to disk)
- auto result_baseline = test_baseline(model, params, tokens);
- suite.results.push_back(result_baseline.empty() ? test_status::FAIL : test_status::PASS);
-
- // Test 2: sequence removal isolation
- suite.results.push_back(test_seq_rm_isolated(model, params, tokens) ? test_status::PASS : test_status::FAIL);
-
- if (!result_baseline.empty()) {
- // Test 3: state load
- suite.results.push_back(test_state_load(model, params, tokens, result_baseline) ? test_status::PASS : test_status::FAIL);
-
- // Test 4: seq copy (host)
- suite.results.push_back(test_seq_cp_host(model, params, tokens, result_baseline) ? test_status::PASS : test_status::FAIL);
-
- // Test 5: seq copy (device)
- suite.results.push_back(test_seq_cp_device(model, params, tokens, result_baseline) ? test_status::PASS : test_status::FAIL);
- } else {
- // tests 3-5 depend on the baseline result and the state file it saves
- suite.results.push_back(test_status::SKIP);
- suite.results.push_back(test_status::SKIP);
- suite.results.push_back(test_status::SKIP);
- }
-
- // Test 6: seq copy (host, scatter)
- suite.results.push_back(test_seq_cp_scatter(model, params, tokens, 6, false) ? test_status::PASS : test_status::FAIL);
-
- // Test 7: seq copy (device, scatter)
- suite.results.push_back(test_seq_cp_scatter(model, params, tokens, 7, true) ? test_status::PASS : test_status::FAIL);
-
- // Test 8: state blob round-trip
- suite.results.push_back(test_state_roundtrip(model, params, tokens) ? test_status::PASS : test_status::FAIL);
-
- // Test 9: state restore failure
- suite.results.push_back(test_state_restore_failure(model, params, tokens) ? test_status::PASS : test_status::FAIL);
-
- // Test 10: state rotation
- suite.results.push_back(test_state_rotation(model, params) ? test_status::PASS : test_status::FAIL);
-
- return suite;
-}
-
-
-static void print_usage(int /* argc */, char ** argv) {
- LOG("\nexample usage:\n");
- LOG("\n %s -m your_model.gguf\n", argv[0]);
- LOG("\n %s --models tests/test-models\n", argv[0]);
- LOG("\n %s -m your_model.gguf -lv 5\n", argv[0]);
- LOG("\n");
-}
-
-int main(int argc, char ** argv) {
- std::setlocale(LC_NUMERIC, "C");
-
- common_params params;
- params.prompt = "";
- params.n_batch = 100;
- params.out_file = "dump_state.bin";
- params.sampling.seed = 1234;
-
- common_init();
-
- // extract our own --models DIR option before handing the rest to the common arg parser
- std::string models_dir;
- std::vector<char *> filtered_argv;
- filtered_argv.push_back(argv[0]);
- for (int i = 1; i < argc; i++) {
- if (strcmp(argv[i], "--models") == 0) {
- if (i + 1 >= argc) {
- LOG_ERR("%s: --models requires a directory argument\n", __func__);
- return 1;
- }
- models_dir = argv[i + 1];
- i++;
- } else {
- filtered_argv.push_back(argv[i]);
- }
- }
- filtered_argv.push_back(nullptr);
- const int fargc = (int)filtered_argv.size() - 1;
-
- // in --models mode there is no single model; set a placeholder so the common parser's
- // "--model is required" check passes (each model is set individually inside the loop)
- if (!models_dir.empty()) {
- params.model.path = models_dir;
- }
-
- if (!common_params_parse(fargc, filtered_argv.data(), params, LLAMA_EXAMPLE_COMMON, print_usage)) {
- return 1;
- }
-
- if (params.n_parallel == 1) {
- LOG_TRC("%s: n_parallel == 1, enabling unified kv cache\n", __func__);
- params.kv_unified = true;
- }
-
- if (params.n_predict < 0) {
- params.n_predict = 16;
- }
-
- llama_backend_init();
-
- if (!models_dir.empty()) {
- // run the suite over every dummy model in the directory
- if (!std::filesystem::exists(models_dir) || !std::filesystem::is_directory(models_dir)) {
- LOG_ERR("%s: models directory '%s' does not exist\n", __func__, models_dir.c_str());
- return 1;
- }
-
- std::vector<std::string> models;
- for (const auto & entry : std::filesystem::directory_iterator(models_dir)) {
- if (entry.is_regular_file() && entry.path().extension() == ".gguf") {
- models.push_back(entry.path().string());
- }
- }
- std::sort(models.begin(), models.end());
-
- if (models.empty()) {
- LOG_ERR("%s: no .gguf models found in '%s'\n", __func__, models_dir.c_str());
- return 1;
- }
-
- auto col_width = [](const char * name) { return (int) std::max(strlen(name), (size_t) 4); };
-
- size_t name_width = 5; // "Model"
- for (const auto & model_path : models) {
- name_width = std::max(name_width, std::filesystem::path(model_path).filename().string().size());
- }
-
- // silence everything but the table itself (LOG has verbosity LOG_LEVEL_OUTPUT = 0)
- common_log_set_verbosity_thold(0);
-
- LOG("%-*s", (int) name_width, "Model");
- for (const auto & name : test_names) {
- LOG(" %-*s", col_width(name), name);
- }
- LOG("\n");
- common_log_flush(common_log_main());
-
- size_t n_pass = 0;
- size_t n_fail = 0;
- for (const auto & model_path : models) {
- const auto name = std::filesystem::path(model_path).filename().string();
-
- LOG("%-*s", (int) name_width, name.c_str());
- common_log_flush(common_log_main());
-
- const test_suite suite = run_save_load_tests_for_model(model_path, params);
-
- for (size_t i = 0; i < suite.results.size(); i++) {
- LOG(" %s%*s", test_status_str(suite.results[i]), col_width(test_names[i]) - 4, "");
- }
- LOG("\n");
- common_log_flush(common_log_main());
-
- if (suite.all_passed()) {
- n_pass++;
- } else {
- n_fail++;
- }
- }
-
- common_log_set_verbosity_thold(LOG_DEFAULT_LLAMA);
- common_log_flush(common_log_main());
-
- LOG_INF("%s: summary: %zu passed, %zu failed (of %zu)\n", __func__, n_pass, n_fail, models.size());
-
- return n_fail == 0 ? 0 : 1;
- }
-
- // single-model mode
- const test_suite suite = run_save_load_tests_for_model(params.model.path, params);
- const bool all_passed = suite.all_passed();
- if (all_passed) {
- LOG("\nAll tests passed.\n");
- }
- return all_passed ? 0 : 1;
-}
diff --git a/tests/test-state-restore-fragmented.cpp b/tests/test-state-restore-fragmented.cpp
deleted file mode 100644
index d20c231d7..000000000
--- a/tests/test-state-restore-fragmented.cpp
+++ /dev/null
@@ -1,125 +0,0 @@
-// Test for state restore with fragmented KV cache
-// This tests the fix for: https://github.com/ggml-org/llama.cpp/issues/17527
-// The issue was that state restore required contiguous KV cache slots,
-// which fails when the cache is fragmented.
-//
-// The fix changes find_slot(ubatch, true) to find_slot(ubatch, false)
-// in state_read_meta(), allowing non-contiguous slot allocation.
-
-// TODO: merge with test-save-load-state.cpp
-// TODO: merge with test-recurrent-state-rollback.cpp
-
-#include "arg.h"
-#include "common.h"
-#include "llama.h"
-
-#include <vector>
-#include <cstdio>
-#include <cstring>
-
-int main(int argc, char ** argv) {
- common_params params;
-
- params.sampling.seed = 1234;
- params.kv_unified = true;
- params.n_parallel = 3;
- params.n_ctx = 256;
-
- common_init();
-
- if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_COMMON)) {
- return 1;
- }
-
- // init
-
- llama_backend_init();
-
- common_init_result_ptr llama_init = common_init_from_params(params);
-
- llama_model * model = llama_init->model();
- llama_context * ctx = llama_init->context();
-
- if (model == nullptr || ctx == nullptr) {
- fprintf(stderr, "%s : failed to init\n", __func__);
- return 1;
- }
-
- GGML_UNUSED(model);
-
- // tokenize prompt
- std::vector<llama_token> tokens(70, 1);
-
- // interleave the 3 sequences:
- // 01201230123...
- common_batch batch(ctx);
- for (size_t i = 0; i < tokens.size(); i++) {
- for (int s = 0; s < params.n_parallel; ++s) {
- batch.add(tokens[i], i, s, false);
- }
- }
- batch.set_output(batch.size() - 1, true);
-
- if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get())) {
- fprintf(stderr, "%s : failed to decode seq 0\n", __func__);
- return 1;
- }
-
- fprintf(stderr, "%s : processed prompt on seq 0, 1, 2 (%zu tokens each)\n", __func__, tokens.size());
-
- // Save state of seq 1
- std::vector<uint8_t> seq_state(llama_state_seq_get_size(ctx, 1));
- const size_t ncopy = llama_state_seq_get_data(ctx, seq_state.data(), seq_state.size(), 1);
- if (ncopy != seq_state.size()) {
- fprintf(stderr, "%s : failed to save seq 1 state\n", __func__);
- return 1;
- }
- fprintf(stderr, "%s : saved seq 1 state, %zu bytes\n", __func__, ncopy);
-
- // clear seq 1 to create a "hole" in the KV cache (fragmentation)
- // 0.20.20.20.2....
- llama_memory_t mem = llama_get_memory(ctx);
- llama_memory_seq_rm(mem, 1, -1, -1);
- fprintf(stderr, "%s : cleared seq 1 to create fragmentation\n", __func__);
-
- // Now the cache has holes where seq 1 was
- // This creates fragmentation - there's no contiguous block large enough
- // for the seq 1 state if we only look for contiguous slots
-
- // Restore seq 1 state into seq 1 (should work with non-contiguous allocation)
- // We use seq 1 since it's a valid sequence ID (0 to n_parallel-1)
- // Before the fix, this would fail with "failed to find available cells in kv cache"
- const size_t nset = llama_state_seq_set_data(ctx, seq_state.data(), seq_state.size(), 1);
- if (nset != seq_state.size()) {
- fprintf(stderr, "%s : FAILED to restore seq state into fragmented cache (got %zu, expected %zu)\n",
- __func__, nset, seq_state.size());
- fprintf(stderr, "%s : This is the bug - state restore fails with fragmented KV cache\n", __func__);
- return 1;
- }
- fprintf(stderr, "%s : restored state into seq 1, %zu bytes\n", __func__, nset);
-
- // Verify we can decode with the restored state
- // Generate one token to verify the restored state is usable
- auto sparams = llama_sampler_chain_default_params();
- llama_sampler * smpl = llama_sampler_chain_init(sparams);
- llama_sampler_chain_add(smpl, llama_sampler_init_dist(params.sampling.seed));
-
- auto next_token = llama_sampler_sample(smpl, ctx, -1);
- auto next_token_str = common_token_to_piece(ctx, next_token);
-
- batch.clear();
- batch.add(next_token, (int)tokens.size(), 1, true);
-
- if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get())) {
- fprintf(stderr, "%s : failed to decode with restored state\n", __func__);
- llama_sampler_free(smpl);
- return 1;
- }
-
- fprintf(stderr, "%s : successfully decoded with restored state, generated: '%s'\n", __func__, next_token_str.c_str());
- fprintf(stderr, "%s : SUCCESS - state restore works with fragmented KV cache\n", __func__);
-
- llama_sampler_free(smpl);
-
- return 0;
-}