From a17e94d229fc7fed1567a799baebd11b771fe5af Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Tue, 13 Oct 2026 16:21:55 +0300 Subject: [PATCH] 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 --- .github/workflows/build-apple.yml | 2 +- .github/workflows/build-openvino.yml | 2 +- ci/run.sh | 10 +- common/log.cpp | 18 +- common/log.h | 1 + ggml/src/ggml.c | 15 +- src/llama-ext.h | 8 + src/llama-kv-cache.cpp | 10 +- src/llama-memory.cpp | 33 + src/llama-model.cpp | 4 + src/llama-model.h | 3 + tests/.gitignore | 1 + tests/CMakeLists.txt | 39 +- tests/llama-context/cases-kv.cpp | 694 +++++++++++++++ tests/llama-context/cases-rs.cpp | 404 +++++++++ tests/llama-context/test-case.cpp | 69 ++ tests/llama-context/test-case.h | 83 ++ tests/llama-context/test-log.cpp | 119 +++ tests/llama-context/test-log.h | 33 + tests/llama-context/test-runner.cpp | 117 +++ tests/llama-context/test-runner.h | 23 + tests/llama-context/test-util.cpp | 403 +++++++++ tests/llama-context/test-util.h | 161 ++++ tests/test-llama-context.cpp | 552 ++++++++++++ tests/test-recurrent-state-rollback.cpp | 739 ---------------- tests/test-save-load-state.cpp | 1069 ----------------------- tests/test-state-restore-fragmented.cpp | 125 --- 27 files changed, 2772 insertions(+), 1965 deletions(-) create mode 100644 tests/llama-context/cases-kv.cpp create mode 100644 tests/llama-context/cases-rs.cpp create mode 100644 tests/llama-context/test-case.cpp create mode 100644 tests/llama-context/test-case.h create mode 100644 tests/llama-context/test-log.cpp create mode 100644 tests/llama-context/test-log.h create mode 100644 tests/llama-context/test-runner.cpp create mode 100644 tests/llama-context/test-runner.h create mode 100644 tests/llama-context/test-util.cpp create mode 100644 tests/llama-context/test-util.h create mode 100644 tests/test-llama-context.cpp delete mode 100644 tests/test-recurrent-state-rollback.cpp delete mode 100644 tests/test-save-load-state.cpp delete mode 100644 tests/test-state-restore-fragmented.cpp 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 #include +#include // 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 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 +#include + 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 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 llama_memory_get_buffers(const llama_memory_t mem) { + llama_memory_buffers_writer writer; + + mem->state_write(writer); + + return std::vector(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 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 +#include +#include +#include +#include +#include +#include +#include +#include + +// +// 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(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 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 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 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 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 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 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 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 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 * 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(24, mr.tokens.size())); + const llama_tokens tokens_verify(mr.tokens.end() - std::min(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 baseline; + if (!decode(tokens_verify, 1, &baseline)) { + return case_result::fail(); + } + + const std::vector>> sub_cases = { + { "buffer", [&]() { + std::vector 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 data; + { + std::ifstream f(path, std::ios::binary); + data.assign(std::istreambuf_iterator(f), std::istreambuf_iterator()); + } + + 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 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> 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 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(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 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 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 make_kv_cases() { + std::vector cases; + + cases.push_back(std::make_unique()); + cases.push_back(std::make_unique()); + cases.push_back(std::make_unique()); + cases.push_back(std::make_unique(false)); + cases.push_back(std::make_unique(true)); + cases.push_back(std::make_unique(false)); + cases.push_back(std::make_unique(true)); + cases.push_back(std::make_unique()); + cases.push_back(std::make_unique()); + cases.push_back(std::make_unique()); + cases.push_back(std::make_unique()); + + 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 +#include +#include +#include +#include + +// 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> 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 make_rs_cases() { + std::vector cases; + + cases.push_back(std::make_unique()); + cases.push_back(std::make_unique()); + cases.push_back(std::make_unique()); + + 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 + +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 & test_cases() { + static const std::vector cases = [] { + std::vector 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 +#include +#include + +// +// 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; + +// the cases of one source file, in table order +std::vector make_kv_cases(); // cases-kv.cpp +std::vector make_rs_cases(); // cases-rs.cpp + +// all cases, in table order +const std::vector & 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 +#include +#include +#include +#include +#include + +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(f)), std::istreambuf_iterator()); + 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 + +// +// 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 +#include +#include +#include +#include + +std::vector select_cases(const std::string & filter) { + std::vector 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 run_cases(model_run & mr, const std::vector & 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 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 +#include + +// +// 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 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 run_cases(model_run & mr, const std::vector & 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 +#include +#include +#include +#include + +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::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::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::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 & 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 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 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 & 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 & 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 & a, const std::vector & 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 & 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 +#include +#include + +// +// 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> logits; + + bool empty() const { return tokens.empty(); } +}; + +// +// decode +// + +bool get_current_logits(llama_context * ctx, std::vector & 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 & 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 & 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 & a, const std::vector & b); + +// overwrite the tensor data with 0xff bytes (NaN when read as f16/f32), so that the restore fails +bool corrupt_state(std::vector & 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 +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// 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 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 collect_models(const std::string & models_dir) { + std::vector 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 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 & 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 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 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(n_hw, 4) : 1; + const size_t n_pool = parallel_workers > 0 + ? std::min((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 n_pass(cases.size(), 0); + std::vector n_skip(cases.size(), 0); + std::vector 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 results(models.size()); + std::vector ready(models.size(), 0); + + std::mutex pool_mutex; + size_t next_print = 0; + + auto print_ready = [&]() { + std::lock_guard 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 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 lock(pool_mutex); + ready[i] = 1; + } + + print_ready(); + } + }; + + std::vector 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 -#include -#include -#include -#include -#include -#include -#include -#include -#include - -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 & 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 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(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::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::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::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::infinity(); - nmse_a0 = 1.0; - } - } - } - const double nmse_val = nmse_a0 == 0.0 ? (nmse_ab == 0.0 ? 0.0 : std::numeric_limits::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::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::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 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> 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 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 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 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 -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -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 & a, const std::vector & 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> logits; - - bool empty() const { return tokens.empty(); } -}; - -static bool get_current_logits(llama_context * ctx, std::vector & 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 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 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 & 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 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 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 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 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 & 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 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 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 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 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 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 & 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 * 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(24, tokens.size())); - const llama_tokens tokens_verify(tokens.end() - std::min(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 baseline; - if (!decode(tokens_verify, 1, &baseline)) { - return false; - } - - const std::vector>> cases = { - { "buffer", [&]() { - std::vector 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 data; - { - std::ifstream f(path, std::ios::binary); - data.assign(std::istreambuf_iterator(f), std::istreambuf_iterator()); - } - - 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 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> 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 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 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 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 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 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 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 -#include -#include - -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 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 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; -}