788 lines
29 KiB
C++
788 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 and, without SWA, does not grow with the tokens";
|
|
}
|
|
|
|
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");
|
|
|
|
if (llama_model_n_swa(mr.model) == 0) {
|
|
CTX_REQUIRE(partial_a == partial_b, "the partial state grows with the tokens");
|
|
}
|
|
|
|
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;
|
|
}
|