Georgi Gerganov
6666a7330c
* tests : fix multi-seq split replay against trained models The rollback snapshot is captured by the graph that decoded the rolled-back tail, while the reference context decodes the same prefix with different batch shapes. The two states differ in the last bits and trained models amplify the difference with every replayed token. Compare only the first replayed token of each seq, then restore the rollback state into the reference context before the seq-1-only tail. Assisted-by: dsh:llama.cpp/ggml-org/Qwen3.8-Flash-Next-GGUF (MacBook, M5 Max, 128GB) * tests : require bit-exact seq-1 isolation in split replay The seq state save is not always bit-lossless - SWA-masked cells are not saved - so reload the restored state back into the rollback context. That way both contexts enter the seq-1-only tail with byte-identical memory, including the cell layout, and the tail logits must match bit-exact. Also assert that the restored state round-trips bit-identical. Assisted-by: dsh:llama.cpp/ggml-org/Qwen3.8-Flash-Next-GGUF (MacBook, M5 Max, 128GB) * tests : drop the n_swa partial-state check from ps The check keys off llama_model_n_swa(), which returns 0 for dsv4 even though its raw KV cache has a 128-cell sliding window. The partial state grows with the tokens there by design, so the flatness check is wrong. Keep the partial <= full invariant. Assisted-by: dsh:llama.cpp/ggml-org/Qwen3.8-Flash-Next-GGUF (MacBook, M5 Max, 128GB)
784 lines
29 KiB
C++
784 lines
29 KiB
C++
// KV cache state cases
|
|
//
|
|
#include "test-case.h"
|
|
|
|
#include "test-log.h"
|
|
#include "test-util.h"
|
|
|
|
#include "common.h"
|
|
|
|
#include "../../src/llama-ext.h" // llama_model_set_attn_rot
|
|
|
|
#include <algorithm>
|
|
#include <cstdio>
|
|
#include <cstring>
|
|
#include <fstream>
|
|
#include <functional>
|
|
#include <iterator>
|
|
#include <string>
|
|
#include <utility>
|
|
#include <vector>
|
|
|
|
//
|
|
// reference run
|
|
//
|
|
|
|
// - decode all but the last token, saving the state to disk
|
|
// - decode the last token
|
|
// - generate n_predict tokens
|
|
// the logits it produces are the reference of every case with needs_baseline()
|
|
struct case_baseline : ctx_test {
|
|
const char * name() const override { return "baseline"; }
|
|
|
|
const char * description() const override {
|
|
return "reference generation over the whole prompt, saves the state before the last token";
|
|
}
|
|
|
|
case_result run(model_run & mr) const override {
|
|
ctx_opts opts;
|
|
opts.n_seq_max = 2;
|
|
|
|
auto ctx = mr.ctx(opts);
|
|
auto smpl = mr.sampler();
|
|
if (!ctx) {
|
|
return case_result::fail();
|
|
}
|
|
|
|
auto n_past = 0;
|
|
if (!common_prompt_batch_decode(ctx.get(), mr.tokens, (int) mr.tokens.size(), n_past, mr.params.n_batch, mr.params.out_file, true)) {
|
|
TLOG_ERR("%s: failed to decode the prompt\n", __func__);
|
|
return case_result::fail();
|
|
}
|
|
|
|
generation_result result;
|
|
if (!generate_tokens(ctx.get(), smpl.get(), n_past, mr.params.n_predict, 0, result)) {
|
|
return case_result::fail();
|
|
}
|
|
if (result.empty()) {
|
|
return case_result::fail();
|
|
}
|
|
|
|
mr.baseline = std::move(result);
|
|
mr.have_baseline = true;
|
|
|
|
return case_result::pass();
|
|
}
|
|
};
|
|
|
|
//
|
|
// sequences
|
|
//
|
|
|
|
// - decode the same prefix into two sequences
|
|
// - remove sequence 0
|
|
// - verify that sequence 1 remains unchanged
|
|
struct case_seq_rm_isolated : ctx_test {
|
|
const char * name() const override { return "seq_rm"; }
|
|
|
|
const char * description() const override {
|
|
return "removing one sequence leaves the state of the other sequences untouched";
|
|
}
|
|
|
|
case_result run(model_run & mr) const override {
|
|
ctx_opts opts;
|
|
opts.n_ctx = 256;
|
|
opts.n_seq_max = 2;
|
|
opts.kv_unified = 1;
|
|
|
|
auto ctx = mr.ctx(opts);
|
|
if (!ctx) {
|
|
return case_result::fail();
|
|
}
|
|
|
|
const size_t n_tokens = std::min<size_t>(mr.tokens.size(), 128);
|
|
for (llama_seq_id seq_id = 0; seq_id < 2; ++seq_id) {
|
|
if (!decode_seq(ctx.get(), mr.tokens, 0, (llama_pos) n_tokens, seq_id, true)) {
|
|
TLOG_ERR("%s: failed to decode the prompt for sequence %d\n", __func__, seq_id);
|
|
return case_result::fail();
|
|
}
|
|
}
|
|
|
|
std::vector<uint8_t> before;
|
|
if (!get_seq_state(ctx.get(), 1, LLAMA_STATE_SEQ_FLAGS_NONE, before)) {
|
|
return case_result::fail();
|
|
}
|
|
|
|
if (!seq_rm(ctx.get(), 0)) {
|
|
return case_result::fail();
|
|
}
|
|
|
|
std::vector<uint8_t> after;
|
|
if (!get_seq_state(ctx.get(), 1, LLAMA_STATE_SEQ_FLAGS_NONE, after)) {
|
|
return case_result::fail();
|
|
}
|
|
|
|
CTX_REQUIRE(blobs_equal(before, after), "removing seq 0 changed the state of seq 1");
|
|
|
|
return case_result::pass();
|
|
}
|
|
};
|
|
|
|
// - load the state saved by the baseline case
|
|
// - replay its last token
|
|
// - generate again and compare against the baseline logits
|
|
struct case_state_load : ctx_test {
|
|
const char * name() const override { return "state_load"; }
|
|
|
|
const char * description() const override {
|
|
return "a saved session replays to the same logits as the run that saved it";
|
|
}
|
|
|
|
bool needs_baseline() const override { return true; }
|
|
|
|
case_result run(model_run & mr) const override {
|
|
ctx_opts opts;
|
|
opts.n_seq_max = 2;
|
|
|
|
auto ctx = mr.ctx(opts);
|
|
auto smpl = mr.sampler();
|
|
if (!ctx) {
|
|
return case_result::fail();
|
|
}
|
|
|
|
int n_past = 0;
|
|
if (!load_state_replay_last(ctx.get(), mr.params.out_file, mr.tokens, n_past)) {
|
|
return case_result::fail();
|
|
}
|
|
|
|
if (!generate_tokens_compare(ctx.get(), smpl.get(), n_past, mr.params.n_predict, 0, mr.baseline, NMSE_THRESHOLD)) {
|
|
return case_result::fail();
|
|
}
|
|
|
|
return case_result::pass();
|
|
}
|
|
};
|
|
|
|
// - load the state saved by the baseline case and replay its last token
|
|
// - copy sequence 0 to sequence 1 over the given io path, clearing the cache in between
|
|
// - continue on sequence 1 and compare against the baseline logits
|
|
struct case_seq_cp : ctx_test {
|
|
explicit case_seq_cp(bool on_device_) : on_device(on_device_) {}
|
|
|
|
const char * name() const override { return on_device ? "cp_d" : "cp_h"; }
|
|
|
|
const char * description() const override {
|
|
return on_device ? "copying a sequence over the device io path keeps its state"
|
|
: "copying a sequence over the host io path keeps its state";
|
|
}
|
|
|
|
bool needs_baseline() const override { return true; }
|
|
|
|
case_result run(model_run & mr) const override {
|
|
ctx_opts opts;
|
|
opts.n_seq_max = 2;
|
|
|
|
auto ctx = mr.ctx(opts);
|
|
auto smpl = mr.sampler();
|
|
if (!ctx) {
|
|
return case_result::fail();
|
|
}
|
|
|
|
const uint32_t flags = on_device ? LLAMA_STATE_SEQ_FLAGS_ON_DEVICE : LLAMA_STATE_SEQ_FLAGS_NONE;
|
|
TLOG_INF("io path: %s\n", on_device ? "device" : "host");
|
|
|
|
int n_past = 0;
|
|
if (!load_state_replay_last(ctx.get(), mr.params.out_file, mr.tokens, n_past)) {
|
|
return case_result::fail();
|
|
}
|
|
|
|
{
|
|
std::vector<uint8_t> seq_store;
|
|
if (!get_seq_state(ctx.get(), 0, flags, seq_store)) {
|
|
return case_result::fail();
|
|
}
|
|
TLOG_TRC("%s: seq 0 copied, %zu bytes\n", __func__, seq_store.size());
|
|
|
|
llama_memory_clear(llama_get_memory(ctx.get()), true);
|
|
TLOG_TRC("%s: kv cache cleared\n", __func__);
|
|
|
|
if (!set_seq_state(ctx.get(), 1, flags, seq_store)) {
|
|
return case_result::fail();
|
|
}
|
|
TLOG_TRC("%s: seq 1 restored, %zu bytes\n", __func__, seq_store.size());
|
|
}
|
|
|
|
if (!generate_tokens_compare(ctx.get(), smpl.get(), n_past, mr.params.n_predict, 1, mr.baseline, NMSE_THRESHOLD)) {
|
|
return case_result::fail();
|
|
}
|
|
|
|
return case_result::pass();
|
|
}
|
|
|
|
private:
|
|
bool on_device;
|
|
};
|
|
|
|
// - decode the same prefix on two sequences, interleaving seq 0 cells between the seq 1 cells
|
|
// - save the seq 1 state, free the interleaved seq 0 cells, and restore via the given io path
|
|
// - the restore destination is non-contiguous: scatter reads are batched per contiguous run
|
|
// - save again on the host and compare the two blobs byte for byte
|
|
struct case_seq_cp_scatter : ctx_test {
|
|
explicit case_seq_cp_scatter(bool on_device_) : on_device(on_device_) {}
|
|
|
|
const char * name() const override { return on_device ? "cp_d_s" : "cp_h_s"; }
|
|
|
|
const char * description() const override {
|
|
return "saving and restoring a sequence with a fragmented (scattered) cell range";
|
|
}
|
|
|
|
case_result run(model_run & mr) const override {
|
|
ctx_opts opts;
|
|
opts.n_ctx = 256;
|
|
opts.n_seq_max = 2;
|
|
opts.kv_unified = 1;
|
|
|
|
auto ctx = mr.ctx(opts);
|
|
if (!ctx) {
|
|
return case_result::fail();
|
|
}
|
|
|
|
const uint32_t flags = on_device ? LLAMA_STATE_SEQ_FLAGS_ON_DEVICE : LLAMA_STATE_SEQ_FLAGS_NONE;
|
|
TLOG_INF("io path: %s\n", on_device ? "device" : "host");
|
|
|
|
// seq 0 cells 0,1,4 interleave the seq 1 cells 2,3,5
|
|
if (!decode_one(ctx.get(), mr.tokens[0], 0, 0) ||
|
|
!decode_one(ctx.get(), mr.tokens[1], 1, 0) ||
|
|
!decode_one(ctx.get(), mr.tokens[0], 0, 1) ||
|
|
!decode_one(ctx.get(), mr.tokens[1], 1, 1) ||
|
|
!decode_one(ctx.get(), mr.tokens[2], 2, 0) ||
|
|
!decode_one(ctx.get(), mr.tokens[2], 2, 1)) {
|
|
TLOG_ERR("%s: failed to build the interleaved state\n", __func__);
|
|
return case_result::fail();
|
|
}
|
|
|
|
// host blob: contains the KV data, used for the byte-for-byte comparison
|
|
std::vector<uint8_t> state_before;
|
|
if (!get_seq_state(ctx.get(), 1, LLAMA_STATE_SEQ_FLAGS_NONE, state_before)) {
|
|
return case_result::fail();
|
|
}
|
|
|
|
// save via the io path under test
|
|
std::vector<uint8_t> state_save;
|
|
if (!get_seq_state(ctx.get(), 1, flags, state_save)) {
|
|
return case_result::fail();
|
|
}
|
|
TLOG_TRC("%s: seq 1 saved via %s, %zu bytes\n", __func__, on_device ? "device" : "host", state_save.size());
|
|
|
|
// free seq 0's cells so the ring is fragmented: the restore destination (seq 1's interleaved cells) stays non-contiguous
|
|
if (!seq_rm(ctx.get(), 0)) {
|
|
return case_result::fail();
|
|
}
|
|
|
|
// restore via the io path under test
|
|
if (!set_seq_state(ctx.get(), 1, flags, state_save)) {
|
|
return case_result::fail();
|
|
}
|
|
TLOG_TRC("%s: seq 1 restored via %s, %zu bytes\n", __func__, on_device ? "device" : "host", state_save.size());
|
|
|
|
std::vector<uint8_t> state_after;
|
|
if (!get_seq_state(ctx.get(), 1, LLAMA_STATE_SEQ_FLAGS_NONE, state_after)) {
|
|
return case_result::fail();
|
|
}
|
|
|
|
// the blob is serialized in sequence cell order, so identical bytes iff the restore wrote the same KV
|
|
CTX_REQUIRE(blobs_equal(state_before, state_after), "the restored KV state is not byte-identical to the saved state");
|
|
|
|
return case_result::pass();
|
|
}
|
|
|
|
private:
|
|
bool on_device;
|
|
};
|
|
|
|
//
|
|
// sessions
|
|
//
|
|
|
|
// - save a sequence, erase it, restore it, save it again
|
|
// - the two blobs must be identical
|
|
// compares blobs rather than generated text: a partially restored cell still decodes to plausible tokens
|
|
struct case_state_roundtrip : ctx_test {
|
|
const char * name() const override { return "rt"; }
|
|
|
|
const char * description() const override {
|
|
return "save, erase, restore, save again - the two blobs must be identical";
|
|
}
|
|
|
|
case_result run(model_run & mr) const override {
|
|
auto ctx = mr.ctx();
|
|
if (!ctx) {
|
|
return case_result::fail();
|
|
}
|
|
|
|
common_batch batch = common_batch_get_one(ctx.get(), mr.tokens);
|
|
if (llama_process(ctx.get(), LLAMA_PROCESS_TYPE_DECODE, batch.get())) {
|
|
TLOG_ERR("%s: failed to decode the prompt\n", __func__);
|
|
return case_result::fail();
|
|
}
|
|
|
|
std::vector<uint8_t> blob_a;
|
|
if (!get_seq_state(ctx.get(), 0, LLAMA_STATE_SEQ_FLAGS_NONE, blob_a)) {
|
|
return case_result::fail();
|
|
}
|
|
|
|
if (!seq_rm(ctx.get(), 0)) {
|
|
return case_result::fail();
|
|
}
|
|
|
|
if (!set_seq_state(ctx.get(), 0, LLAMA_STATE_SEQ_FLAGS_NONE, blob_a)) {
|
|
TLOG_ERR("%s: failed to restore seq 0\n", __func__);
|
|
return case_result::fail();
|
|
}
|
|
|
|
std::vector<uint8_t> blob_b;
|
|
if (!get_seq_state(ctx.get(), 0, LLAMA_STATE_SEQ_FLAGS_NONE, blob_b)) {
|
|
return case_result::fail();
|
|
}
|
|
|
|
CTX_REQUIRE(blobs_equal(blob_a, blob_b), "the state changed across a save / erase / restore cycle");
|
|
|
|
return case_result::pass();
|
|
}
|
|
};
|
|
|
|
// the partial state holds what seq_rm cannot roll back: recurrent states and SWA windows
|
|
struct case_state_partial_size : ctx_test {
|
|
const char * name() const override { return "ps"; }
|
|
|
|
const char * description() const override {
|
|
return "the partial state is not larger than the full state";
|
|
}
|
|
|
|
case_result run(model_run & mr) const override {
|
|
auto ctx = mr.ctx();
|
|
if (!ctx) {
|
|
return case_result::fail();
|
|
}
|
|
|
|
const llama_pos n_tokens = mr.tokens.size();
|
|
const llama_pos n_half = n_tokens / 2;
|
|
|
|
if (!decode_seq(ctx.get(), mr.tokens, 0, n_half, 0, true)) {
|
|
return case_result::fail();
|
|
}
|
|
|
|
const size_t partial_a = llama_state_seq_get_size_ext(ctx.get(), 0, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
|
|
const size_t full_a = llama_state_seq_get_size_ext(ctx.get(), 0, LLAMA_STATE_SEQ_FLAGS_NONE);
|
|
|
|
if (!decode_seq(ctx.get(), mr.tokens, n_half, n_tokens, 0, true)) {
|
|
return case_result::fail();
|
|
}
|
|
|
|
const size_t partial_b = llama_state_seq_get_size_ext(ctx.get(), 0, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
|
|
const size_t full_b = llama_state_seq_get_size_ext(ctx.get(), 0, LLAMA_STATE_SEQ_FLAGS_NONE);
|
|
|
|
TLOG_INF("%s: partial %zu -> %zu, full %zu -> %zu\n", __func__, partial_a, partial_b, full_a, full_b);
|
|
|
|
CTX_REQUIRE(partial_a <= full_a && partial_b <= full_b, "the partial state is larger than the full state");
|
|
|
|
return case_result::pass();
|
|
}
|
|
};
|
|
|
|
// a failed restore must leave the sequence empty and must not change the logits of other sequences
|
|
struct case_state_restore_failure : ctx_test {
|
|
const char * name() const override { return "rf"; }
|
|
|
|
const char * description() const override {
|
|
return "restoring a corrupted state fails and leaves the other sequences untouched";
|
|
}
|
|
|
|
case_result run(model_run & mr) const override {
|
|
ctx_opts opts;
|
|
opts.n_ctx = 256;
|
|
opts.n_seq_max = 4;
|
|
opts.kv_unified = 1;
|
|
|
|
// without flash attention, corrupted data left behind by the restore shows up as NaN logits on the other sequences
|
|
opts.fa = LLAMA_FLASH_ATTN_TYPE_DISABLED;
|
|
|
|
auto ctx = mr.ctx(opts);
|
|
if (!ctx) {
|
|
return case_result::fail();
|
|
}
|
|
|
|
llama_memory_t mem = llama_get_memory(ctx.get());
|
|
if (mem == nullptr) {
|
|
TLOG_INF("no memory to test\n");
|
|
return case_result::pass();
|
|
}
|
|
|
|
const auto decode = [&](const llama_tokens & inp, llama_seq_id seq_id, std::vector<float> * logits_out) {
|
|
if (!decode_tokens(ctx.get(), inp, seq_id)) {
|
|
TLOG_ERR("%s: failed to decode on sequence %d\n", __func__, seq_id);
|
|
return false;
|
|
}
|
|
|
|
if (logits_out && !get_current_logits(ctx.get(), *logits_out)) {
|
|
TLOG_ERR("%s: failed to get logits\n", __func__);
|
|
return false;
|
|
}
|
|
|
|
return true;
|
|
};
|
|
|
|
const llama_tokens tokens_save (mr.tokens.begin(), mr.tokens.begin() + std::min<size_t>(24, mr.tokens.size()));
|
|
const llama_tokens tokens_verify(mr.tokens.end() - std::min<size_t>(8, mr.tokens.size()), mr.tokens.end());
|
|
|
|
// the cases share a working directory, so the state file is named after the model
|
|
const std::string path = "state-restore-failure." + mr.name + ".tmp.bin";
|
|
|
|
llama_memory_clear(mem, true);
|
|
|
|
std::vector<float> baseline;
|
|
if (!decode(tokens_verify, 1, &baseline)) {
|
|
return case_result::fail();
|
|
}
|
|
|
|
const std::vector<std::pair<const char *, std::function<bool()>>> sub_cases = {
|
|
{ "buffer", [&]() {
|
|
std::vector<uint8_t> state(llama_state_seq_get_size(ctx.get(), 0));
|
|
GGML_ASSERT(llama_state_seq_get_data(ctx.get(), state.data(), state.size(), 0) == state.size());
|
|
llama_memory_seq_rm(mem, 0, -1, -1);
|
|
|
|
if (!corrupt_state(state)) {
|
|
return false;
|
|
}
|
|
|
|
return llama_state_seq_set_data(ctx.get(), state.data(), state.size(), 0) == 0;
|
|
}},
|
|
{ "file", [&]() {
|
|
GGML_ASSERT(llama_state_seq_save_file(ctx.get(), path.c_str(), 0, tokens_save.data(), tokens_save.size()) > 0);
|
|
llama_memory_seq_rm(mem, 0, -1, -1);
|
|
|
|
std::vector<uint8_t> data;
|
|
{
|
|
std::ifstream f(path, std::ios::binary);
|
|
data.assign(std::istreambuf_iterator<char>(f), std::istreambuf_iterator<char>());
|
|
}
|
|
|
|
if (!corrupt_state(data)) {
|
|
std::remove(path.c_str());
|
|
return false;
|
|
}
|
|
|
|
{
|
|
std::ofstream f(path, std::ios::binary);
|
|
f.write((const char *) data.data(), data.size());
|
|
}
|
|
|
|
llama_tokens tokens_out(tokens_save.size());
|
|
size_t n_token_count = 0;
|
|
const size_t nread = llama_state_seq_load_file(ctx.get(), path.c_str(), 0, tokens_out.data(), tokens_out.size(), &n_token_count);
|
|
std::remove(path.c_str());
|
|
|
|
return nread == 0;
|
|
}},
|
|
};
|
|
|
|
for (const auto & [name, restore_failed] : sub_cases) {
|
|
llama_memory_clear(mem, true);
|
|
|
|
if (!decode(tokens_save, 0, nullptr)) {
|
|
return case_result::fail();
|
|
}
|
|
|
|
CTX_REQUIRE(restore_failed(), std::string(name) + ": restoring a corrupted state did not fail");
|
|
|
|
CTX_REQUIRE(llama_memory_seq_pos_max(mem, 0) == -1,
|
|
std::string(name) + ": the sequence is not empty after a failed restore");
|
|
|
|
std::vector<float> logits;
|
|
if (!decode(tokens_verify, 1, &logits)) {
|
|
return case_result::fail();
|
|
}
|
|
|
|
const logits_diff diff = compare_logits(baseline.data(), logits.data(), logits.size());
|
|
|
|
CTX_REQUIRE(diff.n_nan == 0 && diff.diff_max <= 1e-6f,
|
|
std::string(name) + ": the logits changed after a failed restore (max diff = " +
|
|
std::to_string(diff.diff_max) + ", nan = " + std::to_string(diff.n_nan) + ")");
|
|
|
|
TLOG_TRC("%s: %s: logits match (max diff = %g)\n", __func__, name, (double) diff.diff_max);
|
|
}
|
|
|
|
return case_result::pass();
|
|
}
|
|
};
|
|
|
|
// a KV state saved with attention rotation enabled must restore only into a context with the
|
|
// same setting
|
|
// note: rotation is only active for quantized KV caches with a head size that is a multiple of
|
|
// 64, for other models the restore into the rotation-disabled context is valid and the
|
|
// test passes vacuously
|
|
struct case_state_rotation : ctx_test {
|
|
const char * name() const override { return "rot"; }
|
|
|
|
const char * description() const override {
|
|
return "a state saved with attention rotation restores only into a matching context";
|
|
}
|
|
|
|
case_result run(model_run & mr) const override {
|
|
const auto make_context = [&](ggml_type type_k, ggml_type type_v, bool disable_rotation) {
|
|
// a setting of the model, not of the process - a sweep that runs several models at a
|
|
// time must not let this leak into the contexts of another model
|
|
llama_model_set_attn_rot(mr.model, !disable_rotation);
|
|
|
|
ctx_opts opts;
|
|
opts.n_ctx = 32;
|
|
opts.n_batch = 1;
|
|
opts.n_ubatch = 1;
|
|
opts.type_k = type_k;
|
|
opts.type_v = type_v;
|
|
opts.fa = LLAMA_FLASH_ATTN_TYPE_ENABLED;
|
|
|
|
// probing whether this KV type is supported is expected to fail for some models
|
|
opts.quiet = true;
|
|
|
|
return mr.ctx(opts);
|
|
};
|
|
|
|
std::vector<std::pair<ggml_type, ggml_type>> type_pairs;
|
|
for (const auto & types : { std::pair{GGML_TYPE_Q8_0, GGML_TYPE_Q8_0} }) {
|
|
if (make_context(types.first, types.second, false)) {
|
|
type_pairs.push_back(types);
|
|
}
|
|
}
|
|
if (type_pairs.empty()) {
|
|
return case_result::skip("no supported quantized KV cache type combination");
|
|
}
|
|
|
|
bool success = true;
|
|
for (const auto & types : type_pairs) {
|
|
auto src = make_context(types.first, types.second, false);
|
|
if (!src) {
|
|
TLOG_ERR("%s: failed to create the source context\n", __func__);
|
|
success = false;
|
|
break;
|
|
}
|
|
|
|
llama_token token = 0;
|
|
if (llama_decode(src.get(), llama_batch_get_one(&token, 1))) {
|
|
TLOG_ERR("%s: failed to decode a token\n", __func__);
|
|
success = false;
|
|
break;
|
|
}
|
|
|
|
std::vector<uint8_t> state;
|
|
if (!get_seq_state(src.get(), 0, LLAMA_STATE_SEQ_FLAGS_NONE, state)) {
|
|
if (state.empty()) {
|
|
continue; // no KV state to test
|
|
}
|
|
success = false;
|
|
break;
|
|
}
|
|
|
|
auto matching = make_context(types.first, types.second, false);
|
|
if (!matching || !set_seq_state(matching.get(), 0, LLAMA_STATE_SEQ_FLAGS_NONE, state)) {
|
|
TLOG_ERR("%s: failed to restore a matching rotation\n", __func__);
|
|
success = false;
|
|
break;
|
|
}
|
|
|
|
auto mismatched = make_context(types.first, types.second, true);
|
|
if (!mismatched) {
|
|
TLOG_ERR("%s: failed to create a mismatched rotation context\n", __func__);
|
|
success = false;
|
|
break;
|
|
}
|
|
// a model that does not use attention rotation accepts the mismatched restore
|
|
if (set_seq_state(mismatched.get(), 0, LLAMA_STATE_SEQ_FLAGS_NONE, state)) {
|
|
TLOG_TRC("%s: state restored into rotation-disabled context, model does not use attention rotation\n", __func__);
|
|
}
|
|
}
|
|
|
|
llama_model_set_attn_rot(mr.model, true); // the default, the cases that follow rely on it
|
|
|
|
if (!success) {
|
|
return case_result::fail();
|
|
}
|
|
|
|
return case_result::pass();
|
|
}
|
|
};
|
|
|
|
// restore a sequence into a fragmented cache: the slots freed by the removed sequence are
|
|
// scattered between live cells of the other sequences, so the restore cannot claim one
|
|
// contiguous block
|
|
// the restored state must come back byte for byte, and the sequence must continue to behave
|
|
// exactly as the same sequence in a context where it was never removed
|
|
// ref: https://github.com/ggml-org/llama.cpp/issues/17527
|
|
struct case_state_restore_fragmented : ctx_test {
|
|
const char * name() const override { return "frag"; }
|
|
|
|
const char * description() const override {
|
|
return "restore a sequence into a cache fragmented by the removal of another sequence";
|
|
}
|
|
|
|
std::string skip_reason(model_run & mr) const override {
|
|
if (arch_reserves_single_seq(mr.model)) {
|
|
return "the interleaved batch is a multi-seq graph";
|
|
}
|
|
|
|
return ctx_test::skip_reason(mr);
|
|
}
|
|
|
|
case_result run(model_run & mr) const override {
|
|
const int n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(mr.model));
|
|
|
|
ctx_opts opts;
|
|
opts.n_ctx = 256;
|
|
opts.n_seq_max = 3;
|
|
opts.kv_unified = 1;
|
|
opts.n_batch = 3*70;
|
|
opts.n_ubatch = 3*70; // the interleaved prompt is a single multi-seq batch
|
|
|
|
// ctx removes and restores seq 1, ctx_keep keeps it untouched as the reference
|
|
llama_context_ptr ctx = mr.ctx(opts);
|
|
llama_context_ptr ctx_keep = mr.ctx(opts);
|
|
if (!ctx || !ctx_keep) {
|
|
return case_result::fail();
|
|
}
|
|
|
|
const llama_pos n_tokens = (llama_pos) std::min<size_t>(70, mr.tokens.size());
|
|
|
|
// interleave the 3 sequences: 01201230123... - seq 1 owns every third cell
|
|
const auto decode_interleaved = [&](llama_context * cur) {
|
|
common_batch batch(cur);
|
|
for (llama_pos i = 0; i < n_tokens; i++) {
|
|
for (llama_seq_id s = 0; s < 3; ++s) {
|
|
batch.add(mr.tokens[i], i, s, false);
|
|
}
|
|
}
|
|
batch.set_output(batch.size() - 1, true);
|
|
|
|
return llama_process(cur, LLAMA_PROCESS_TYPE_DECODE, batch.get()) == 0;
|
|
};
|
|
|
|
if (!decode_interleaved(ctx.get()) || !decode_interleaved(ctx_keep.get())) {
|
|
TLOG_ERR("%s: failed to decode the interleaved prompt\n", __func__);
|
|
return case_result::fail();
|
|
}
|
|
|
|
TLOG_INF("%s: processed the prompt on seq 0, 1, 2 (%d tokens each)\n", __func__, n_tokens);
|
|
|
|
std::vector<uint8_t> state_before;
|
|
if (!get_seq_state(ctx.get(), 1, LLAMA_STATE_SEQ_FLAGS_NONE, state_before)) {
|
|
return case_result::fail();
|
|
}
|
|
|
|
// clearing seq 1 leaves holes where its cells were - no contiguous block large enough
|
|
// for the seq 1 state is left in the cache
|
|
llama_memory_seq_rm(llama_get_memory(ctx.get()), 1, -1, -1);
|
|
|
|
if (!set_seq_state(ctx.get(), 1, LLAMA_STATE_SEQ_FLAGS_NONE, state_before)) {
|
|
TLOG_ERR("%s: failed to restore the seq state into a fragmented cache\n", __func__);
|
|
return case_result::fail();
|
|
}
|
|
|
|
TLOG_INF("%s: restored seq 1 into a fragmented cache, %zu bytes\n", __func__, state_before.size());
|
|
|
|
// the restore has to be lossless - the cells the sequence ends up in are irrelevant, its
|
|
// state has to come back exactly as it was saved
|
|
std::vector<uint8_t> state_after;
|
|
if (!get_seq_state(ctx.get(), 1, LLAMA_STATE_SEQ_FLAGS_NONE, state_after)) {
|
|
return case_result::fail();
|
|
}
|
|
|
|
CTX_REQUIRE(blobs_equal(state_before, state_after), "restoring into a fragmented cache changed the state of seq 1");
|
|
|
|
// the restored sequence must continue exactly as the untouched one - decode the same token
|
|
// on seq 1 in both contexts and compare the logits they produce
|
|
const llama_token token = gen_token(n_vocab, 1, n_tokens);
|
|
|
|
if (!decode_one(ctx.get(), token, n_tokens, 1) || !decode_one(ctx_keep.get(), token, n_tokens, 1)) {
|
|
TLOG_ERR("%s: failed to decode with the restored state\n", __func__);
|
|
return case_result::fail();
|
|
}
|
|
|
|
const float * logits = llama_get_logits(ctx.get());
|
|
const float * logits_keep = llama_get_logits(ctx_keep.get());
|
|
if (!logits || !logits_keep) {
|
|
TLOG_ERR("%s: missing logits after the restore\n", __func__);
|
|
return case_result::fail();
|
|
}
|
|
|
|
const logits_diff diff = compare_logits(logits, logits_keep, n_vocab);
|
|
|
|
CTX_REQUIRE(diff.nmse() <= NMSE_THRESHOLD,
|
|
"the logits of the restored seq 1 differ from the reference: nmse " + std::to_string(diff.nmse()));
|
|
|
|
TLOG_INF("%s: restored seq 1 continued identically to the reference, nmse %g\n", __func__, diff.nmse());
|
|
|
|
return case_result::pass();
|
|
}
|
|
};
|
|
|
|
//
|
|
// no_alloc
|
|
//
|
|
|
|
// - load a second instance of the model with no_alloc, so the tensor data stays on disk
|
|
// - create a context from it, which reserves the compute graphs over the dummy buffers
|
|
// - the memory use must still be reported, and nothing real may be allocated for the memory
|
|
struct case_noalloc : ctx_test {
|
|
const char * name() const override { return "noalloc"; }
|
|
|
|
const char * description() const override {
|
|
return "with no_alloc the memory use is simulated: non-zero breakdown, null buffer bases";
|
|
}
|
|
|
|
case_result run(model_run & mr) const override {
|
|
const std::string & path = mr.params.model.path;
|
|
CTX_REQUIRE(!path.empty(), "the model path is not known");
|
|
|
|
llama_model_params mparams = llama_model_default_params();
|
|
mparams.no_alloc = true;
|
|
mparams.load_mode = LLAMA_LOAD_MODE_NONE;
|
|
|
|
llama_model_ptr model(llama_model_load_from_file(path.c_str(), mparams));
|
|
CTX_REQUIRE(model, "failed to load the model with no_alloc");
|
|
|
|
auto ctx = init_ctx(model.get(), common_context_params_to_llama(mr.params), 0);
|
|
CTX_REQUIRE(ctx, "failed to create the context with no_alloc");
|
|
|
|
size_t n_ctx_mem = 0;
|
|
for (const auto & [_, data] : llama_get_memory_breakdown(ctx.get())) {
|
|
n_ctx_mem += data.context;
|
|
}
|
|
|
|
// models without a memory module (BERT, diffusion) have a null memory
|
|
if (llama_get_memory(ctx.get()) != nullptr) {
|
|
CTX_REQUIRE(n_ctx_mem > 0, "the simulated memory breakdown is zero");
|
|
} else {
|
|
CTX_REQUIRE(n_ctx_mem == 0, "a model without memory reported context memory");
|
|
}
|
|
|
|
for (ggml_backend_buffer_t buf : llama_memory_get_buffers(llama_get_memory(ctx.get()))) {
|
|
CTX_REQUIRE(ggml_backend_buffer_get_base(buf) == nullptr, "a memory buffer was really allocated");
|
|
}
|
|
|
|
return case_result::pass();
|
|
}
|
|
};
|
|
|
|
std::vector<ctx_test_ptr> make_kv_cases() {
|
|
std::vector<ctx_test_ptr> cases;
|
|
|
|
cases.push_back(std::make_unique<case_baseline>());
|
|
cases.push_back(std::make_unique<case_seq_rm_isolated>());
|
|
cases.push_back(std::make_unique<case_state_load>());
|
|
cases.push_back(std::make_unique<case_seq_cp>(false));
|
|
cases.push_back(std::make_unique<case_seq_cp>(true));
|
|
cases.push_back(std::make_unique<case_seq_cp_scatter>(false));
|
|
cases.push_back(std::make_unique<case_seq_cp_scatter>(true));
|
|
cases.push_back(std::make_unique<case_state_roundtrip>());
|
|
cases.push_back(std::make_unique<case_state_partial_size>());
|
|
cases.push_back(std::make_unique<case_state_restore_failure>());
|
|
cases.push_back(std::make_unique<case_state_rotation>());
|
|
cases.push_back(std::make_unique<case_state_restore_fragmented>());
|
|
cases.push_back(std::make_unique<case_noalloc>());
|
|
|
|
return cases;
|
|
}
|