diff --git a/src/llama-context.cpp b/src/llama-context.cpp index fb88919f9d..179c526c29 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -661,11 +661,19 @@ void llama_context::sched_reserve() { // reserve again with pp graph to avoid ggml-alloc reallocations during inference { - // TODO: not sure if the following graph would be worst case for multi-stream KV caches: - // - // auto * gf = graph_reserve(n_tokens, 1, n_tokens, mctx.get()); - // - auto * gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get(), model.hparams.no_alloc); + // TODO: the worst case graph is not always reached for `n_seqs > 1` + // need to implement a more robust mechanism that tries a few different inputs and analyzes the results + ggml_cgraph * gf = nullptr; + switch (model.arch) { + case LLM_ARCH_MINIMAX_01: + // the `inp_diag_decay` tensor size scales with `n_seq_tokens^2` which + // makes `n_seqs == 1` use more memory for the compute graph compared to `n_seqs > 1` + gf = graph_reserve(n_tokens, 1, n_outputs_pp, mctx.get(), model.hparams.no_alloc); + break; + default: + gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get(), model.hparams.no_alloc); + }; + if (!gf) { throw std::runtime_error("failed to allocate compute pp buffers"); } @@ -2892,13 +2900,83 @@ public: for (auto & [buft, mbuf] : mbufs_new) { const auto & mbuf_cur = mbufs.at(buft); - if (!mbuf_cur.buf || mbuf_cur.n_tensors != mbuf.n_tensors || mbuf_cur.total_size != mbuf.total_size) { + if (!mbuf_cur.buf || mbuf_cur.total_size != mbuf.total_size) { GGML_ABORT("%s: memory buffer mismatch\n", __func__); } - for (size_t i = 0; i < mbuf_cur.org.size(); ++i) { - ggml_backend_tensor_copy(mbuf_cur.cpy[i], mbuf.org[i]); + if (mbuf_cur.n_tensors == mbuf.n_tensors) { + // same chunking: copy 1:1 by index + for (size_t i = 0; i < mbuf_cur.org.size(); ++i) { + GGML_ASSERT(ggml_nbytes(mbuf_cur.cpy[i]) == ggml_nbytes(mbuf.org[i])); + ggml_backend_tensor_copy(mbuf_cur.cpy[i], mbuf.org[i]); + } + continue; } + + // different chunking: copy the write-side data (mbuf_cur.cpy) into the read-side targets (mbuf.org) + // with a byte cursor. Write and read enumerate the same logical data in the same order but may chunk + // it differently, so copy across tensor boundaries rather than 1:1 by index. + const size_t total = mbuf_cur.total_size; + + ggml_init_params params_scratch = { + /*.mem_size =*/ 2*(mbuf_cur.cpy.size() + mbuf.org.size())*ggml_tensor_overhead(), + /*.mem_buffer =*/ NULL, + /*.no_alloc =*/ true, + }; + ggml_context * ctx_scratch = ggml_init(params_scratch); + + size_t src_pos = 0; + size_t dst_pos = 0; + size_t src_j = 0; + size_t dst_i = 0; + size_t src_base = 0; + size_t dst_base = 0; + + while (src_pos < total) { + const auto & src_t = mbuf_cur.cpy[src_j]; + const auto & dst_t = mbuf.org[dst_i]; + + const size_t src_size = ggml_nbytes(src_t); + const size_t dst_size = ggml_nbytes(dst_t); + + const size_t src_off = src_pos - src_base; + const size_t dst_off = dst_pos - dst_base; + + const size_t n_copy = std::min(src_size - src_off, dst_size - dst_off); + + const size_t el = ggml_element_size(src_t); + const int64_t n_el = (int64_t) (n_copy / el); + + auto * src_v = ggml_view_1d(ctx_scratch, src_t, n_el, src_off); + ggml_backend_view_init(src_v); + auto * dst_v = ggml_view_1d(ctx_scratch, dst_t, n_el, dst_off); + ggml_backend_view_init(dst_v); + + ggml_backend_tensor_copy(src_v, dst_v); + + src_pos += n_copy; + dst_pos += n_copy; + + if (src_pos - src_base == src_size) { + src_base = src_pos; + ++src_j; + } + if (dst_pos - dst_base == dst_size) { + dst_base = dst_pos; + ++dst_i; + } + } + + GGML_ASSERT(src_pos == total && dst_pos == total); + // any tensors left unvisited hold no data + for (size_t i = src_j; i < mbuf_cur.cpy.size(); ++i) { + GGML_ASSERT(ggml_nbytes(mbuf_cur.cpy[i]) == 0); + } + for (size_t i = dst_i; i < mbuf.org.size(); ++i) { + GGML_ASSERT(ggml_nbytes(mbuf.org[i]) == 0); + } + + ggml_free(ctx_scratch); } GGML_ASSERT(buf_size == 0); diff --git a/src/models/minimax-01.cpp b/src/models/minimax-01.cpp index f14626b2c7..361114acc3 100644 --- a/src/models/minimax-01.cpp +++ b/src/models/minimax-01.cpp @@ -181,7 +181,7 @@ public: return res; } - const llama_hparams & hparams; + const llama_hparams hparams; ggml_tensor * inp_slopes = nullptr; // F32 [n_head] ggml_tensor * inp_q_decay = nullptr; // F32 [1, n_head, n_batch] diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index b9f9d4b78a..fe3d14ffc5 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -149,6 +149,7 @@ if (LLAMA_LLGUIDANCE) endif () llama_build(test-recurrent-state-rollback.cpp) +llama_build(test-save-load-state.cpp) 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) @@ -237,6 +238,14 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) set_tests_properties(test-recurrent-state-rollback-dsv4 PROPERTIES FIXTURES_REQUIRED generate-models ) + + # 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) endif() llama_build_and_test(test-chat-peg-parser.cpp peg-parser/simple-tokenize.cpp) @@ -299,10 +308,6 @@ llama_build_and_test(test-backend-sampler.cpp LABEL "model") 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) -# Test state save/load functionality -llama_build_and_test(test-save-load-state.cpp LABEL "model" ARGS -m "${MODEL_DEST}") -set_tests_properties(test-save-load-state PROPERTIES FIXTURES_REQUIRED test-download-model) - if (APPLE) llama_build(test-rset-release.cpp) endif() diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index d58d90952e..35a3286e4a 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -65,7 +65,7 @@ static void set_tensor_data(struct ggml_tensor * tensor, void * userdata) { } static void usage(char ** argv) { - printf("Usage: %s [-a/--arch arch] [-s/--seed seed] [-v/--verbose]\n", argv[0]); + printf("Usage: %s [-a/--arch arch] [-s/--seed seed] [-o/--out dir] [-v/--verbose] [-h/--help]\n", argv[0]); } static std::vector get_tokens(const uint32_t n_tokens, const uint32_t n_vocab, const size_t seed){ @@ -82,7 +82,7 @@ static std::vector get_tokens(const uint32_t n_tokens, const uint32 static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { gguf_context_ptr ret(gguf_init_empty()); llama_model_saver ms(arch, ret.get()); - const uint32_t n_ctx = 128; + const uint32_t n_ctx = 256; uint32_t n_vocab = 128; uint32_t n_embd = 256; @@ -256,10 +256,12 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS, std::vector(n_layer, 4)); } - ms.add_kv(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, arch == LLM_ARCH_MINIMAX_M3 || arch == LLM_ARCH_DEEPSEEK4 ? n_head : uint32_t(1)); + // minimax-m3 keeps one indexer head per GQA head; the rest use a fixed 64 to match the fused + ms.add_kv(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, arch == LLM_ARCH_MINIMAX_M3 ? n_head : uint32_t(64)); // qwen4exp ropes indexer keys with the main rotary width, so its head can't be < n_rot ms.add_kv(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, - arch == LLM_ARCH_QWEN4EXP ? n_embd_head : uint32_t(64)); + arch == LLM_ARCH_QWEN4EXP ? n_embd_head : uint32_t(128)); + ms.add_kv(LLM_KV_ATTENTION_INDEXER_TOP_K, uint32_t(8)); ms.add_kv(LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, uint32_t(4)); ms.add_kv(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, uint32_t(1)); @@ -762,6 +764,10 @@ int main(int argc, char ** argv) { std::string out; for (int i = 1; i < argc; i++) { + if (strcmp(argv[i], "-h") == 0 || strcmp(argv[i], "--help") == 0) { + usage(argv); + return 0; + } if (strcmp(argv[i], "-a") == 0 || strcmp(argv[i], "--arch") == 0) { if (i + 1 < argc) { const std::string arch_name = argv[++i]; diff --git a/tests/test-save-load-state.cpp b/tests/test-save-load-state.cpp index 6e93ce6fb8..0ceab7c545 100644 --- a/tests/test-save-load-state.cpp +++ b/tests/test-save-load-state.cpp @@ -3,8 +3,12 @@ #include "log.h" #include "llama-cpp.h" +#include #include +#include +#include #include +#include #include struct llama_batch_ptr { @@ -53,7 +57,9 @@ static llama_tokens generate_tokens(llama_context * ctx, llama_sampler * smpl, i // - decode the last token // - generate n_predict tokens static llama_tokens test_baseline(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens) { - auto ctx = llama_context_ptr{llama_init_from_model(model, common_context_params_to_llama(params))}; + 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)}; @@ -161,7 +167,9 @@ static bool test_seq_rm_isolated( // - 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 llama_tokens & expected_result) { - auto ctx = llama_context_ptr{llama_init_from_model(model, common_context_params_to_llama(params))}; + 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)}; @@ -347,38 +355,18 @@ static bool test_seq_cp_device(struct llama_model * model, const struct common_p } -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(); - - if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_COMMON)) { - 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; - } - - ggml_backend_load_all(); +// Run the full save/load test suite (tests 1-5) for a single model. +// Returns true if all tests pass, false otherwise. +static bool run_save_load_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\n", __func__); - return 1; + LOG_ERR("%s: failed to init model '%s'\n", __func__, model_path.c_str()); + return false; } GGML_ASSERT(llama_init->context() == nullptr); @@ -411,30 +399,127 @@ int main(int argc, char ** argv) { // Test 1: baseline (saves state to disk) auto result_baseline = test_baseline(model, params, tokens); if (result_baseline.empty()) { - return 1; + return false; } // Test 2: sequence removal isolation if (!test_seq_rm_isolated(model, params, tokens)) { - return 1; + return false; } // Test 3: state load if (!test_state_load(model, params, tokens, result_baseline)) { - return 1; + return false; } // Test 4: seq copy (host) if (!test_seq_cp_host(model, params, tokens, result_baseline)) { - return 1; + return false; } // Test 5: seq copy (device) if (!test_seq_cp_device(model, params, tokens, result_baseline)) { - return 1; + return false; } LOG("\nAll tests passed.\n"); - return 0; + return true; +} + + +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)) { + 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; + } + + ggml_backend_load_all(); + + 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; + } + + LOG_INF("%s: running save/load tests over %zu models in '%s'\n", __func__, models.size(), models_dir.c_str()); + + size_t n_pass = 0; + size_t n_fail = 0; + for (const auto & model_path : models) { + LOG("\n================================================================\n"); + LOG_INF("%s: model %s\n", __func__, model_path.c_str()); + + if (run_save_load_tests_for_model(model_path, params)) { + n_pass++; + } else { + n_fail++; + } + } + + LOG("\n================================================================\n"); + 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 + return run_save_load_tests_for_model(params.model.path, params) ? 0 : 1; }