Files
llama.cpp/tests/test-save-load-state.cpp
T
Georgi Gerganov a2878d30df metal : single-source fusion table + fusion debug rework (#28164)
* metal : rework fusion patterns into a single table

All fusable op patterns for the Metal backend are now declared once in a
fusion table (ggml-metal-fuse.cpp) and consumed by both the graph optimizer
(ggml_metal_fuse_max, packing) and the op encoders (ggml_metal_fuse_next,
compute). The two phases share the same pattern table plus ggml_can_fuse_subgraph_ext
for the structural checks, and differ only in the mode used for the pattern
check (STRUCTURAL at optimize time, since tensors are not allocated yet, and
FULL at compute time, including Metal buffer placement). This also protects the
snake activation (MUL + SIN + SQR + MUL + ADD) from being reordered during graph
optimization, which was previously unprotected.

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-0731

* metal : fix absolute output indices in fusion patterns

ggml_can_fuse_subgraph_ext expects the outputs array to contain absolute graph
node indices (it indexes cgraph->nodes[outputs[i]]), but the fusion table query
was passing a relative index (n_ops - 1). As a result the last node of every
pattern was not recognized as an output and was subjected to the elidable
use-count check, which failed for essentially all fusions. This silently
disabled the norm/MUL fusion and caused a ~5% token-generation regression.

Pass the absolute graph index of the last node instead.

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-0731

* metal : fuse gated_delta_net with cache cpy

Add GGML_METAL_FUSE_GDN_CACHE to the fusion table: when the gated_delta_net
kernel is followed by a cpy that scatters its recurrent state snapshots into
the KV cache, the kernel writes the snapshots straight into the cache buffer
and the trailing cpy is elided.

The gdn output has other consumers (the attn scores view), so unlike the
elision-chain patterns this is not a simple chain: a 'raw' flag on the fusion
pattern skips the generic chain/shape and ggml_can_fuse_subgraph_ext checks,
making the pattern-specific check callback the sole validator. Packing
(ggml_metal_fuse_max) now matches on the same view-transparent node sequence
that the compute phase uses, so the gdn + cache cpy group is packed along with
any intermediate views and stays adjacent through the reorder.

The fused cpy is a view consumer of the gdn (it writes the cache directly),
so its mem-range is skipped in the encoder; the skip is restricted to CPY
nodes consuming the previous fused node through a view so other fusions are
unaffected.

Add test_gated_delta_net_cache_fusion and register 5 cases.

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-0731

* metal : drop is_view_consumer mem-range skip

The is_view_consumer skip was carried over from the upstream gated_delta_net
cache-fusion draft, but it is not needed: keeping the elided cpy's mem-range in
the concurrency tracker only ever adds a (conservative) memory barrier at the
fusion point. It can never remove a barrier, so it cannot introduce a race. The
worst case is one spurious barrier per gdn+cache-cpy fusion, which is within
run-to-run noise on Qwen3.5-0.8B Q8_0.

Dropping the check keeps the mem-range loop uniform for all fused groups.

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-0731

* metal : rename gated_delta_net fused state output args

Rename the fused cache-write kernel argument to match the rest of the kargs:
state_out_stride -> nb_out (and widen it to uint64_t), and the local buffer id
bid_state_out -> bid_out.

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-0731

* metal : rename raw fusion flag to unsafe

raw did not convey that the flag opts a fusion pattern out of the generic
elision-chain safety net (ggml_can_fuse_subgraph_ext + chain/shape checks).
rename it to 'unsafe' to make explicit that the pattern's check callback is the
sole validator and must re-establish the safety guarantees itself.

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-0731

* metal : tidy fusion pattern checks and table

- const-correct ggml_metal_fuse_outputs buffer
- annotate unused check-callback parameters
- drop a redundant size_t cast
- align the ops/table initializers and add blank-line separation

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-0731

* metal : add generic fusion stats via ad-hoc proc-address API

Add a device-owned fusion context that lets a test tool count how many
times each fusion pattern fires and toggle fusion. It is exposed through
the ad-hoc ggml_backend_reg_get_proc_address mechanism with generic names
so the testing tool is backend-agnostic:

- ggml_backend_fusion_stats_init: start collecting fusion stats; when a
  context is created afterwards it registers the labels/counters and
  encodes single-threaded (n_cb == 0) so the counters are race-free
- ggml_backend_fusion_stats_reset / _get_stats / _set_enabled

The context lives on the metal device (not on the last backend context),
so counters accumulate across contexts and reads are always consistent.
The enable/disable toggle is initialized from GGML_METAL_FUSION_DISABLE
and can be overridden by the test through set_enabled. Labels are
synthesized from the fuse table via ggml_metal_fuse_label (e.g.
"GATED_DELTA_NET+CPY").

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-0731

* tests : add fusion count regression test with per-backend baseline

test-fusion runs every dummy model generated by test-llama-archs on a
single backend (single-threaded encoding, n_cb == 0) with fusion enabled
and disabled, and for each mode (prefill / decode) reports the per-fusion
counters and the NMSE between the fused and unfused logits, plus the NMSE
against a CPU reference.

A fusion pattern that silently stops matching (or fires when it should
not) is caught as a regression by comparing the counters against a
committed per-backend TSV baseline:

- --record writes the golden baseline, --check (default) validates it
- the unfused run doubles as a control: its counters must be all-zero
- NMSE is skipped when it is NaN or the arch is already broken on the
  device (e.g. plamo2 on Metal), so the count check is the hard gate
- baseline counts depend only on graph structure, not weights (verified
  stable across weight seeds)
- the fusion stats API is resolved through the ad-hoc get_proc_address
  mechanism with generic names; a backend that does not export it makes
  the test fail with an error

The committed MTL0.tsv baseline covers 110 dummy archs (298 rows).

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-0731

* tests : rename fusion api helpers to match stats_init signature

Align the test with the ad-hoc fusion stats API: fusion_stats_init no
longer takes an enable bool (stats are turned on by calling it), so the
proc-address wrappers and typedefs are renamed to the api_* convention.

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-0731

* tests : rename backend to device in fusion test CLI

The fusion test operates on a compute device (e.g. MTL0), not a backend,
so rename the --backend argument to --device and the backend_name
variable to device_name. Keep "backend" where it refers to the ggml
backend interface (the ad-hoc proc-address mechanism).

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-0731

* tests : add --model and --help to fusion test

--model FILE runs the fusion regression test over a single model file
instead of enumerating a --models DIR. --models and --model are mutually
exclusive. Also add a --help/-h option that prints the usage.

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-0731

* tests : use backend base name for fusion baseline output

The fusion test is invoked with a specific device name (e.g. MTL0), but
its output - the recorded baseline and the header it writes - should be
named after the backend base name (e.g. MTL, via ggml_backend_reg_name),
since the counters depend on the backend, not on the specific device
index. Rename the committed baseline to MTL.tsv.

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-0731

* tests : run fusion test from ci instead of ctest

The fusion test needs Metal and generates a lot of dummy models, so it
does not belong in the generic ctest suite. Move it to ci/run.sh as
gg_run_test_fusion, gated on GG_BUILD_METAL like
gg_run_test_llama_archs_tensor_split: it generates the dummy models with
test-llama-archs -o and then validates the fusion counts against the
committed baseline. test-fusion.cpp is still built (llama_build) but no
longer registered as a ctest.

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-0731

* tests : align fusion baseline TSV columns

Pad the TSV fields to fixed widths so the columns line up regardless of
the variable arch and fusion-label lengths, and trim each field on parse
so the padded file is still accepted. Regenerate the committed MTL.tsv
baseline in the padded format (data unchanged, verified identical modulo
padding).

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-0731

* tests : widen label column and align fusion TSV header

Give the label column more room (28 chars) and fix the column header
widths so they match the data rows (moe/mode/label), keeping the header
aligned with the values. Regenerate the MTL.tsv baseline in the new
format (data unchanged).

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-0731

* tests : switch fusion baseline from TSV to CSV

Use comma-separated values like the rest of the project, keeping the
padded, aligned columns. Split on ',' and trim on parse. Rename the
committed baseline to MTL.csv (data unchanged, verified identical modulo
padding/separator). Update the ci/run.sh check path accordingly.

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-0731

* cont : rebase + update MTL stats

* tests : avoid graph reallocations for some archs

* metal : tidy fusion debugging context and op init

- simplify the shared fusion debugging context comments
- shorten the ggml_metal_fusion struct comment
- align the ggml_metal_fuse struct fields and comments
- move the fusion parameter of ggml_metal_op_init right after dev

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-0731

* tests : dedup fusion baseline into any mode

prefill and decode always produce the same per-graph fusion count, so
store a single row per label with mode = "any" and the per-graph count
instead of two rows. this halves the baseline size and keeps the check
stable.

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-0731

* ci : move fusion model generation to a separate step

the dummy models generated by test-llama-archs are reused by other tests,
so generate them once in their own step instead of inside test_fusion.

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-0731

* tests : bump nmse thold

* models : fix plamo2 graph

* tests : remove "skip" logic from test-fusion

* tests : set qwen3tts dummy vocab to codec head size

the dummy qwen3tts model used a vocab of 4096 while the codec head is
3072, so the graph padded the output with -inf which made the NMSE in
test-fusion produce NaN. use the exact codec head size instead so the
padding is not generated at all.

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-0731

* tests : regen fusion baseline

reflect the plamo2 graph fix, which changed its fusion pattern split
(RMS_NORM+MUL 11->10, RMS_NORM+MUL+ADD 3->4; same total).

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-0731

* ci : skip dummy model generation on OpenVINO

test-llama-archs does not build on the OpenVINO platform, so do not try
to generate the dummy models there.

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-0731

* cont : minor

* tests : enable test-llama-archs on windows

* cont : disable on windows + workaround

* metal : naming nits

* test-fusion : add instructions to update baseline

* context : fix Kimi-K3 graph reserve

* fusion : update MTL

* cont : fix naming

* metal : rework fusion info storage

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* metal : align fusion info API

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* metal : use opaque fusion handle in ad-hoc API

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* ci : move fusion test to dedicated workflow

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* cont : run only on ggml changes

* cont : simplify

* fusion : remove multi-output stuff for now

* ci : fix typo
2026-09-11 12:41:54 +03:00

694 lines
24 KiB
C++

#include "arg.h"
#include "common.h"
#include "log.h"
#include "llama-cpp.h"
#include <algorithm>
#include <clocale>
#include <cstring>
#include <filesystem>
#include <random>
#include <string>
#include <vector>
struct llama_batch_ptr {
llama_batch batch;
llama_batch_ptr(int32_t n_tokens, int32_t embd, int32_t n_seq_max)
: batch{llama_batch_init(n_tokens, embd, n_seq_max)} {}
~llama_batch_ptr() { llama_batch_free(batch); }
llama_batch_ptr(const llama_batch_ptr &) = delete;
llama_batch_ptr & operator=(const llama_batch_ptr &) = delete;
llama_batch_ptr(llama_batch_ptr &&) = default;
llama_batch_ptr & operator=(llama_batch_ptr &&) = default;
llama_batch & get() { return batch; }
const llama_batch & get() const { return batch; }
};
static llama_tokens generate_tokens(llama_context * ctx, llama_sampler * smpl, int & n_past, int32_t n_predict, llama_seq_id seq_id) {
llama_tokens result;
llama_batch_ptr batch(1, 0, 1);
for (int i = 0; i < n_predict; i++) {
auto next_token = llama_sampler_sample(smpl, ctx, -1);
LOG("%d ", next_token);
result.push_back(next_token);
common_batch_clear(batch.get());
common_batch_add(batch.get(), next_token, n_past, {seq_id}, true);
if (llama_decode(ctx, batch.get())) {
LOG_ERR("\n%s: failed to evaluate\n", __func__);
return {};
}
n_past++;
}
return result;
}
// Test 1: baseline
// - decode all but the last token
// - save state to disk
// - 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 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 {};
}
LOG("\n=== Test 1: baseline ===\n");
auto result = generate_tokens(ctx.get(), smpl.get(), n_past, params.n_predict, 0);
if (result.empty()) {
return {};
}
LOG("\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;
}
LOG("\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) {
llama_batch_ptr batch(n_tokens, 0, 1);
for (size_t i = 0; i < n_tokens; ++i) {
common_batch_add(batch.get(), tokens[i], i, { seq_id }, i == n_tokens - 1);
}
if (llama_decode(ctx.get(), batch.get())) {
LOG_ERR("%s: failed to decode prompt for sequence %d\n", __func__, seq_id);
return false;
}
}
const auto get_seq_state = [&](llama_seq_id seq_id, std::vector<uint8_t> & state) {
const size_t state_size = llama_state_seq_get_size(ctx.get(), seq_id);
if (state_size == 0) {
LOG_ERR("%s: sequence state is empty\n", __func__);
return false;
}
state.resize(state_size);
const size_t ncopy = llama_state_seq_get_data(ctx.get(), state.data(), state.size(), seq_id);
if (ncopy != state.size()) {
LOG_ERR("%s: sequence state length %zu does not match expected length %zu\n",
__func__, ncopy, state.size());
return false;
}
return true;
};
std::vector<uint8_t> state_before;
if (!get_seq_state(1, state_before)) {
return false;
}
if (!llama_memory_seq_rm(llama_get_memory(ctx.get()), 0, -1, -1)) {
LOG_ERR("%s: failed to remove sequence 0\n", __func__);
return false;
}
std::vector<uint8_t> state_after;
if (!get_seq_state(1, state_after)) {
return false;
}
if (state_before != state_after) {
LOG_ERR("%s: removing sequence 0 changed sequence 1\n", __func__);
return false;
}
LOG("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 llama_tokens & 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));
LOG("\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
auto result = generate_tokens(ctx.get(), smpl.get(), n_past, params.n_predict, 0);
if (result.empty()) {
return false;
}
if (result != expected_result) {
LOG_ERR("\n%s: error: generation differs from expected\n", __func__);
return false;
}
LOG("\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 llama_tokens & 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));
LOG("\n=== Test 4: seq copy (host) ===\n");
// Load state from file
llama_tokens unused_sts(tokens.size());
size_t n_token_count_out = 0;
if (!llama_state_load_file(ctx.get(), params.out_file.data(), unused_sts.data(), unused_sts.size(), &n_token_count_out)) {
LOG_ERR("\n%s: failed to load state\n", __func__);
return false;
}
LOG_TRC("%s: loaded state with %zu tokens\n", __func__, n_token_count_out);
// Replay last token
int n_past = (int) n_token_count_out - 1;
if (!common_replay_last_token(ctx.get(), tokens.back(), n_past)) {
return false;
}
n_past++;
// Migrate KV cache from seq 0 to seq 1 (CPU path)
{
std::vector<uint8_t> seq_store(llama_state_seq_get_size(ctx.get(), 0));
const size_t ncopy = llama_state_seq_get_data(ctx.get(), seq_store.data(), seq_store.size(), 0);
if (ncopy != seq_store.size()) {
LOG_ERR("\n%s: seq copy data length %zd does not match expected length %zd\n", __func__, ncopy, seq_store.size());
return false;
}
LOG_TRC("%s: seq 0 copied, %zd bytes\n", __func__, ncopy);
llama_memory_clear(llama_get_memory(ctx.get()), true);
LOG_TRC("%s: kv cache cleared\n", __func__);
const size_t nset = llama_state_seq_set_data(ctx.get(), seq_store.data(), seq_store.size(), 1);
if (nset != seq_store.size()) {
LOG_ERR("\n%s: seq set data length %zd does not match expected length %zd\n", __func__, nset, seq_store.size());
return false;
}
LOG_TRC("%s: seq 1 restored, %zd bytes\n", __func__, nset);
}
// Generate tokens on seq 1
auto result = generate_tokens(ctx.get(), smpl.get(), n_past, params.n_predict, 1);
if (result.empty()) {
return false;
}
if (result != expected_result) {
LOG_ERR("\n%s: error: generation differs from expected\n", __func__);
return false;
}
LOG("\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 llama_tokens & 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));
LOG("\n=== Test 5: seq copy (device) ===\n");
// Load state from file
llama_tokens unused_sts(tokens.size());
size_t n_token_count_out = 0;
if (!llama_state_load_file(ctx.get(), params.out_file.data(), unused_sts.data(), unused_sts.size(), &n_token_count_out)) {
LOG_ERR("\n%s: failed to load state\n", __func__);
return false;
}
LOG_TRC("%s: loaded state with %zu tokens\n", __func__, n_token_count_out);
// Replay last token
int n_past = (int) n_token_count_out - 1;
if (!common_replay_last_token(ctx.get(), tokens.back(), n_past)) {
return false;
}
n_past++;
// Migrate KV cache from seq 0 to seq 1 (on-device path)
{
std::vector<uint8_t> seq_store(llama_state_seq_get_size_ext(ctx.get(), 0, LLAMA_STATE_SEQ_FLAGS_ON_DEVICE));
const size_t ncopy = llama_state_seq_get_data_ext(ctx.get(), seq_store.data(), seq_store.size(), 0, LLAMA_STATE_SEQ_FLAGS_ON_DEVICE);
if (ncopy != seq_store.size()) {
LOG_ERR("\n%s: seq copy data length %zd does not match expected length %zd\n", __func__, ncopy, seq_store.size());
return false;
}
LOG_TRC("%s: seq 0 copied, %zd bytes\n", __func__, ncopy);
llama_memory_clear(llama_get_memory(ctx.get()), true);
LOG_TRC("%s: kv cache cleared\n", __func__);
const size_t nset = llama_state_seq_set_data_ext(ctx.get(), seq_store.data(), seq_store.size(), 1, LLAMA_STATE_SEQ_FLAGS_ON_DEVICE);
if (nset != seq_store.size()) {
LOG_ERR("\n%s: seq set data length %zd does not match expected length %zd\n", __func__, nset, seq_store.size());
return false;
}
LOG_TRC("%s: seq 1 restored, %zd bytes\n", __func__, nset);
}
// Generate tokens on seq 1
auto result = generate_tokens(ctx.get(), smpl.get(), n_past, params.n_predict, 1);
if (result.empty()) {
return false;
}
if (result != expected_result) {
LOG_ERR("\n%s: error: generation differs from expected\n", __func__);
return false;
}
LOG("\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)};
LOG("\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) {
llama_batch_ptr batch(1, 0, 1);
common_batch_add(batch.get(), tok, pos, { seq }, true);
return llama_decode(ctx.get(), batch.get()) == 0;
};
// seq 0 cells 0,1,4 interleave the seq 1 cells 2,3,5
if (!decode_one(tokens[0], 0, 0) ||
!decode_one(tokens[1], 1, 0) ||
!decode_one(tokens[0], 0, 1) ||
!decode_one(tokens[1], 1, 1) ||
!decode_one(tokens[2], 2, 0) ||
!decode_one(tokens[2], 2, 1)) {
LOG_ERR("%s: failed to build interleaved state\n", __func__);
return false;
}
const auto get_seq_state = [&](llama_seq_id seq_id, uint32_t fl, std::vector<uint8_t> & state) {
const size_t state_size = llama_state_seq_get_size_ext(ctx.get(), seq_id, fl);
if (state_size == 0) {
LOG_ERR("%s: sequence state is empty\n", __func__);
return false;
}
state.resize(state_size);
const size_t ncopy = llama_state_seq_get_data_ext(ctx.get(), state.data(), state.size(), seq_id, fl);
if (ncopy != state.size()) {
LOG_ERR("%s: sequence state length %zu does not match expected length %zu\n",
__func__, ncopy, state.size());
return false;
}
return true;
};
// host blob: contains the KV data, used for the byte-for-byte comparison
std::vector<uint8_t> state_before;
if (!get_seq_state(1, LLAMA_STATE_SEQ_FLAGS_NONE, state_before)) {
return false;
}
// save via the io path under test
std::vector<uint8_t> state_save;
if (!get_seq_state(1, flags, state_save)) {
return false;
}
LOG_TRC("%s: seq 1 saved via %s, %zu bytes\n", __func__, on_device ? "device" : "host", state_save.size());
// free seq 0's cells so the ring is fragmented: the restore destination (seq 1's interleaved cells) stays non-contiguous
if (!llama_memory_seq_rm(llama_get_memory(ctx.get()), 0, -1, -1)) {
LOG_ERR("%s: failed to remove sequence 0\n", __func__);
return false;
}
// restore via the io path under test
const size_t nset = llama_state_seq_set_data_ext(ctx.get(), state_save.data(), state_save.size(), 1, flags);
if (nset != state_save.size()) {
LOG_ERR("%s: seq set data length %zu does not match expected length %zu\n", __func__, nset, state_save.size());
return false;
}
LOG_TRC("%s: seq 1 restored via %s, %zu bytes\n", __func__, on_device ? "device" : "host", nset);
std::vector<uint8_t> state_after;
if (!get_seq_state(1, LLAMA_STATE_SEQ_FLAGS_NONE, state_after)) {
return false;
}
// the blob is serialized in sequence cell order, so identical bytes iff the restore wrote the same KV
if (state_before.size() != state_after.size() || memcmp(state_before.data(), state_after.data(), state_before.size()) != 0) {
LOG_ERR("\n%s: error: restored KV state is not byte-identical to the saved state\n", __func__);
return false;
}
LOG("\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)};
LOG("\n=== Test 8: state blob round-trip ===\n");
if (llama_decode(ctx.get(), llama_batch_get_one(const_cast<llama_token *>(tokens.data()), (int32_t) tokens.size()))) {
LOG_ERR("\n%s: failed to decode prompt\n", __func__);
return false;
}
std::vector<uint8_t> blob_a(llama_state_seq_get_size(ctx.get(), 0));
const size_t n_a = llama_state_seq_get_data(ctx.get(), blob_a.data(), blob_a.size(), 0);
if (n_a != blob_a.size()) {
LOG_ERR("\n%s: saved %zu bytes, expected %zu\n", __func__, n_a, blob_a.size());
return false;
}
if (!llama_memory_seq_rm(llama_get_memory(ctx.get()), 0, -1, -1)) {
LOG_ERR("\n%s: failed to erase seq 0\n", __func__);
return false;
}
if (llama_state_seq_set_data(ctx.get(), blob_a.data(), blob_a.size(), 0) != blob_a.size()) {
LOG_ERR("\n%s: failed to restore seq 0\n", __func__);
return false;
}
std::vector<uint8_t> blob_b(llama_state_seq_get_size(ctx.get(), 0));
const size_t n_b = llama_state_seq_get_data(ctx.get(), blob_b.data(), blob_b.size(), 0);
if (n_b != n_a) {
LOG_ERR("\n%s: re-saved %zu bytes, expected %zu\n", __func__, n_b, n_a);
return false;
}
size_t n_diff = 0;
size_t i_diff = 0;
for (size_t i = 0; i < n_a; i++) {
if (blob_a[i] != blob_b[i]) {
if (n_diff == 0) {
i_diff = i;
}
n_diff++;
}
}
if (n_diff > 0) {
LOG_ERR("\n%s: state changed across a restore: %zu of %zu bytes differ, first at offset %zu\n",
__func__, n_diff, n_a, i_diff);
return false;
}
LOG("\nPASS\n");
return true;
}
// Run the full save/load test suite (tests 1-8) 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 model '%s'\n", __func__, model_path.c_str());
return false;
}
GGML_ASSERT(llama_init->context() == nullptr);
// Tokenize prompt or generate random tokens
llama_tokens tokens;
if (params.prompt.empty()) {
const int n_prompt = params.n_batch;
// this path is useful for model files that do not have a tokenizer
LOG_INF("%s: no prompt provided, generating %d (n_batch) random tokens\n", __func__, n_prompt);
const auto * vocab = llama_model_get_vocab(model);
const auto n_vocab = llama_vocab_n_tokens(vocab);
std::mt19937 rng(params.sampling.seed);
std::uniform_int_distribution<llama_token> dist(0, n_vocab - 1);
for (int i = 0; i < n_prompt; i++) {
tokens.push_back(dist(rng));
}
} else {
LOG_INF("%s: tokenizing prompt '%s'\n", __func__, params.prompt.c_str());
auto ctx = llama_context_ptr{llama_init_from_model(model, common_context_params_to_llama(params))};
tokens = common_tokenize(ctx.get(), params.prompt, true);
}
LOG_INF("%s: the input prompt is %d tokens\n", __func__, (int)tokens.size());
// Test 1: baseline (saves state to disk)
auto result_baseline = test_baseline(model, params, tokens);
if (result_baseline.empty()) {
return false;
}
// Test 2: sequence removal isolation
if (!test_seq_rm_isolated(model, params, tokens)) {
return false;
}
// Test 3: state load
if (!test_state_load(model, params, tokens, result_baseline)) {
return false;
}
// Test 4: seq copy (host)
if (!test_seq_cp_host(model, params, tokens, result_baseline)) {
return false;
}
// Test 5: seq copy (device)
if (!test_seq_cp_device(model, params, tokens, result_baseline)) {
return false;
}
// Test 6: seq copy (host, scatter)
if (!test_seq_cp_scatter(model, params, tokens, 6, false)) {
return false;
}
// Test 7: seq copy (device, scatter)
if (!test_seq_cp_scatter(model, params, tokens, 7, true)) {
return false;
}
// Test 8: state blob round-trip
if (!test_state_roundtrip(model, params, tokens)) {
return false;
}
LOG("\nAll tests passed.\n");
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<char *> filtered_argv;
filtered_argv.push_back(argv[0]);
for (int i = 1; i < argc; i++) {
if (strcmp(argv[i], "--models") == 0) {
if (i + 1 >= argc) {
LOG_ERR("%s: --models requires a directory argument\n", __func__);
return 1;
}
models_dir = argv[i + 1];
i++;
} else {
filtered_argv.push_back(argv[i]);
}
}
filtered_argv.push_back(nullptr);
const int fargc = (int)filtered_argv.size() - 1;
// in --models mode there is no single model; set a placeholder so the common parser's
// "--model is required" check passes (each model is set individually inside the loop)
if (!models_dir.empty()) {
params.model.path = models_dir;
}
if (!common_params_parse(fargc, filtered_argv.data(), params, LLAMA_EXAMPLE_COMMON)) {
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<std::string> models;
for (const auto & entry : std::filesystem::directory_iterator(models_dir)) {
if (entry.is_regular_file() && entry.path().extension() == ".gguf") {
models.push_back(entry.path().string());
}
}
std::sort(models.begin(), models.end());
if (models.empty()) {
LOG_ERR("%s: no .gguf models found in '%s'\n", __func__, models_dir.c_str());
return 1;
}
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;
}