Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
378 lines
16 KiB
C++
378 lines
16 KiB
C++
// Unit tests for TransformerKVCache ring_mode: slot mapping, wrap advance,
|
|
// chronological export, wrap-placement import, and non-ring regression, plus
|
|
// the bucketed step-graph capacity a static-cache AR loop sizes itself from.
|
|
#include "engine/framework/runtime/kv_cache.h"
|
|
|
|
#include "engine/framework/core/backend.h"
|
|
|
|
#include "ggml-backend.h"
|
|
#include "ggml.h"
|
|
|
|
#include <algorithm>
|
|
#include <cstdint>
|
|
#include <cstring>
|
|
#include <iostream>
|
|
#include <limits>
|
|
#include <memory>
|
|
#include <stdexcept>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
namespace {
|
|
|
|
namespace core = engine::core;
|
|
namespace runtime = engine::runtime;
|
|
|
|
int failures = 0;
|
|
|
|
void check(bool condition, const std::string & message) {
|
|
if (!condition) {
|
|
++failures;
|
|
std::cerr << "FAIL: " << message << "\n";
|
|
}
|
|
}
|
|
|
|
struct GgmlContextDeleter {
|
|
void operator()(ggml_context * ctx) const noexcept {
|
|
if (ctx != nullptr) {
|
|
ggml_free(ctx);
|
|
}
|
|
}
|
|
};
|
|
|
|
struct BackendDeleter {
|
|
void operator()(ggml_backend_t backend) const noexcept {
|
|
if (backend != nullptr) {
|
|
ggml_backend_free(backend);
|
|
}
|
|
}
|
|
};
|
|
|
|
// One layer of key/value tensors, F32 by default, with (steps * elems) elements.
|
|
struct TestTensors {
|
|
std::unique_ptr<ggml_context, GgmlContextDeleter> ctx;
|
|
std::unique_ptr<std::remove_pointer_t<ggml_backend_buffer_t>, void (*)(ggml_backend_buffer_t)> buffer{
|
|
nullptr, ggml_backend_buffer_free};
|
|
core::TensorValue keys;
|
|
core::TensorValue values;
|
|
};
|
|
|
|
ggml_backend_t cpu_backend() {
|
|
static std::unique_ptr<std::remove_pointer_t<ggml_backend_t>, BackendDeleter> backend(
|
|
core::init_backend(core::BackendConfig{}), BackendDeleter{});
|
|
if (backend == nullptr) {
|
|
throw std::runtime_error("test CPU backend init failed");
|
|
}
|
|
return backend.get();
|
|
}
|
|
|
|
TestTensors make_tensors(int64_t steps, int64_t elems, ggml_type type = GGML_TYPE_F32) {
|
|
TestTensors out;
|
|
ggml_init_params params{1024ull * 1024ull, nullptr, true};
|
|
out.ctx.reset(ggml_init(params));
|
|
if (out.ctx == nullptr) {
|
|
throw std::runtime_error("test context init failed");
|
|
}
|
|
const int64_t count = steps * elems;
|
|
ggml_tensor * k = ggml_new_tensor_1d(out.ctx.get(), type, count);
|
|
ggml_tensor * v = ggml_new_tensor_1d(out.ctx.get(), type, count);
|
|
out.buffer.reset(ggml_backend_alloc_ctx_tensors(out.ctx.get(), cpu_backend()));
|
|
if (out.buffer == nullptr) {
|
|
throw std::runtime_error("test buffer alloc failed");
|
|
}
|
|
out.keys = core::TensorValue{k, core::TensorShape::from_dims({count}), type};
|
|
out.values = core::TensorValue{v, core::TensorShape::from_dims({count}), type};
|
|
return out;
|
|
}
|
|
|
|
runtime::TransformerKVCache make_cache(
|
|
TestTensors & tensors, int64_t steps, int64_t elems, runtime::TransformerKVCacheOptions options) {
|
|
return runtime::TransformerKVCache(
|
|
steps, elems, {tensors.keys}, {tensors.values}, options);
|
|
}
|
|
|
|
// Row r filled with float(position): positions are directly readable.
|
|
std::vector<float> tagged_row(int64_t position, int64_t elems) {
|
|
return std::vector<float>(static_cast<size_t>(elems), static_cast<float>(position));
|
|
}
|
|
|
|
runtime::KVLayerState tagged_layer(const std::vector<int64_t> & positions, int64_t elems) {
|
|
runtime::KVLayerState layer;
|
|
layer.valid_steps = static_cast<int64_t>(positions.size());
|
|
for (const int64_t pos : positions) {
|
|
const auto row = tagged_row(pos, elems);
|
|
layer.key.insert(layer.key.end(), row.begin(), row.end());
|
|
layer.value.insert(layer.value.end(), row.begin(), row.end());
|
|
}
|
|
return layer;
|
|
}
|
|
|
|
int64_t row_position(const std::vector<float> & flat, int64_t row, int64_t elems) {
|
|
return static_cast<int64_t>(flat[static_cast<size_t>(row * elems)]);
|
|
}
|
|
|
|
void test_slot_mapping() {
|
|
runtime::TransformerKVCacheOptions off;
|
|
{
|
|
TestTensors t = make_tensors(8, 4);
|
|
runtime::TransformerKVCache cache = make_cache(t, 8, 4, off);
|
|
check(cache.slot_for_position(0) == 0, "identity slot 0 when ring off");
|
|
check(cache.slot_for_position(100) == 100, "identity slot 100 when ring off");
|
|
}
|
|
runtime::TransformerKVCacheOptions ring;
|
|
ring.ring_mode = true;
|
|
ring.ring_pinned_steps = 2;
|
|
{
|
|
TestTensors t = make_tensors(8, 4);
|
|
runtime::TransformerKVCache cache = make_cache(t, 8, 4, ring);
|
|
check(cache.slot_for_position(0) == 0, "pinned slot 0");
|
|
check(cache.slot_for_position(1) == 1, "pinned slot 1");
|
|
check(cache.slot_for_position(2) == 2, "first ring slot");
|
|
check(cache.slot_for_position(7) == 7, "last ring slot first lap");
|
|
check(cache.slot_for_position(8) == 2, "wrap to slot 2");
|
|
check(cache.slot_for_position(9) == 3, "wrap to slot 3");
|
|
check(cache.slot_for_position(14) == 2, "second wrap to slot 2");
|
|
bool threw = false;
|
|
try {
|
|
cache.slot_for_position(-1);
|
|
} catch (const std::runtime_error &) {
|
|
threw = true;
|
|
}
|
|
check(threw, "negative position throws");
|
|
}
|
|
std::cout << "slot mapping done\n";
|
|
}
|
|
|
|
void test_option_validation() {
|
|
auto expect_throw = [](runtime::TransformerKVCacheOptions options, const std::string & what) {
|
|
bool threw = false;
|
|
try {
|
|
TestTensors t = make_tensors(8, 4);
|
|
runtime::TransformerKVCache cache = make_cache(t, 8, 4, options);
|
|
(void)cache;
|
|
} catch (const std::runtime_error &) {
|
|
threw = true;
|
|
}
|
|
check(threw, what);
|
|
};
|
|
runtime::TransformerKVCacheOptions zero_cap;
|
|
zero_cap.ring_mode = true;
|
|
// NOTE: cache_steps is clamped to >= 0 in the ctor; ring requires positive.
|
|
{
|
|
bool threw = false;
|
|
try {
|
|
TestTensors t = make_tensors(0, 4);
|
|
runtime::TransformerKVCache cache = make_cache(t, 0, 4, zero_cap);
|
|
(void)cache;
|
|
} catch (const std::runtime_error &) {
|
|
threw = true;
|
|
}
|
|
check(threw, "ring with zero capacity throws");
|
|
}
|
|
runtime::TransformerKVCacheOptions bad_pin;
|
|
bad_pin.ring_mode = true;
|
|
bad_pin.ring_pinned_steps = 8;
|
|
expect_throw(bad_pin, "pinned == capacity throws");
|
|
bad_pin.ring_pinned_steps = -1;
|
|
expect_throw(bad_pin, "negative pinned throws");
|
|
std::cout << "option validation done\n";
|
|
}
|
|
|
|
void test_wrap_round_trip() {
|
|
constexpr int64_t steps = 8;
|
|
constexpr int64_t elems = 4;
|
|
runtime::TransformerKVCacheOptions ring;
|
|
ring.ring_mode = true;
|
|
ring.ring_pinned_steps = 2;
|
|
TestTensors t = make_tensors(steps, elems);
|
|
runtime::TransformerKVCache cache = make_cache(t, steps, elems, ring);
|
|
// Chronological rows for positions [0,1] + [6..11] (2,3,4,5 evicted).
|
|
runtime::TransformerKVState state;
|
|
state.current_end = 12;
|
|
state.layers.push_back(tagged_layer({0, 1, 6, 7, 8, 9, 10, 11}, elems));
|
|
cache.import_state(state);
|
|
check(cache.valid_steps() == 8, "import keeps valid count");
|
|
check(cache.current_end() == 12, "import keeps absolute end");
|
|
// Physical placement: prompt at 0,1; tail 6..11 wraps into 4,5,6,7,2,3.
|
|
const std::vector<float> raw = core::read_tensor_f32(t.keys.tensor);
|
|
const int64_t expected_slots[8] = {0, 1, 6, 7, 8, 9, 10, 11};
|
|
check(row_position(raw, 0, elems) == 0, "slot 0 holds pinned 0");
|
|
check(row_position(raw, 1, elems) == 1, "slot 1 holds pinned 1");
|
|
check(row_position(raw, 2, elems) == 8, "slot 2 holds wrapped 8");
|
|
check(row_position(raw, 3, elems) == 9, "slot 3 holds wrapped 9");
|
|
check(row_position(raw, 4, elems) == 10, "slot 4 holds wrapped 10");
|
|
check(row_position(raw, 5, elems) == 11, "slot 5 holds wrapped 11");
|
|
check(row_position(raw, 6, elems) == 6, "slot 6 holds 6");
|
|
check(row_position(raw, 7, elems) == 7, "slot 7 holds 7");
|
|
// Chronological export restores input order.
|
|
const runtime::TransformerKVState back = cache.export_state();
|
|
check(back.current_end == 12, "export keeps absolute end");
|
|
check(back.layers.front().valid_steps == 8, "export keeps valid count");
|
|
for (int row = 0; row < 8; ++row) {
|
|
check(row_position(back.layers.front().key, row, elems) == expected_slots[row],
|
|
"export row " + std::to_string(row) + " chronological");
|
|
}
|
|
std::cout << "wrap round-trip done\n";
|
|
}
|
|
|
|
void test_wrap_advance() {
|
|
constexpr int64_t steps = 8;
|
|
constexpr int64_t elems = 4;
|
|
runtime::TransformerKVCacheOptions ring;
|
|
ring.ring_mode = true;
|
|
ring.ring_pinned_steps = 2;
|
|
TestTensors t = make_tensors(steps, elems);
|
|
runtime::TransformerKVCache cache = make_cache(t, steps, elems, ring);
|
|
runtime::TransformerKVState state;
|
|
state.current_end = 2;
|
|
state.layers.push_back(tagged_layer({0, 1}, elems));
|
|
cache.import_state(state);
|
|
cache.advance_after_direct_append(6);
|
|
check(cache.valid_steps() == 8, "advance fills to capacity");
|
|
check(cache.current_end() == 8, "advance moves end");
|
|
cache.advance_after_direct_append(4);
|
|
check(cache.valid_steps() == 8, "advance past capacity stays capped");
|
|
check(cache.current_end() == 12, "advance past capacity moves end");
|
|
// Non-ring still throws on overflow.
|
|
runtime::TransformerKVCacheOptions off;
|
|
TestTensors t2 = make_tensors(steps, elems);
|
|
runtime::TransformerKVCache plain = make_cache(t2, steps, elems, off);
|
|
plain.import_state(state);
|
|
bool threw = false;
|
|
try {
|
|
plain.advance_after_direct_append(7);
|
|
} catch (const std::runtime_error &) {
|
|
threw = true;
|
|
}
|
|
check(threw, "non-ring advance overflow still throws");
|
|
std::cout << "wrap advance done\n";
|
|
}
|
|
|
|
void test_non_ring_regression() {
|
|
constexpr int64_t steps = 8;
|
|
constexpr int64_t elems = 4;
|
|
runtime::TransformerKVCacheOptions off;
|
|
TestTensors t = make_tensors(steps, elems);
|
|
runtime::TransformerKVCache cache = make_cache(t, steps, elems, off);
|
|
runtime::TransformerKVState state;
|
|
state.current_end = 5;
|
|
state.layers.push_back(tagged_layer({10, 11, 12, 13, 14}, elems));
|
|
cache.import_state(state);
|
|
const runtime::TransformerKVState back = cache.export_state();
|
|
check(back.current_end == 5, "dense export keeps end");
|
|
for (int row = 0; row < 5; ++row) {
|
|
check(row_position(back.layers.front().key, row, elems) == 10 + row, "dense export row order");
|
|
}
|
|
cache.retain_prefix(3);
|
|
check(cache.valid_steps() == 3, "retain_prefix keeps count");
|
|
check(cache.current_end() == 3, "retain_prefix rewinds end");
|
|
std::cout << "non-ring regression done\n";
|
|
}
|
|
|
|
void test_bucketed_cache_steps() {
|
|
// The shape families size their step graph from: doubling buckets from 128,
|
|
// clamped to prompt + max_tokens, exact once doubling would overshoot.
|
|
check(runtime::bucketed_cache_steps(114, 128) == 242, "one generated window is exact");
|
|
check(runtime::bucketed_cache_steps(100, 200) == 256, "short request rounds to a bucket");
|
|
check(runtime::bucketed_cache_steps(200, 8000) == 512, "generated window stops at 128 steps");
|
|
check(runtime::bucketed_cache_steps(8000, 0) == 8000, "long prompt lands on its own length");
|
|
check(runtime::bucketed_cache_steps(0, 64) == 64, "promptless request clamps to the total");
|
|
check(runtime::bucketed_cache_steps(0, 0) == 0, "empty request is empty");
|
|
check(runtime::bucketed_cache_steps(10, 10, 4, 8) == 16, "window and floor knobs are honoured");
|
|
for (const int64_t prompt : {0, 1, 43, 114, 1000, 8000}) {
|
|
int64_t previous = 0;
|
|
for (const int64_t tokens : {0, 1, 64, 128, 129, 1024, 8192}) {
|
|
const int64_t maximum = prompt + tokens;
|
|
const int64_t bucket = runtime::bucketed_cache_steps(prompt, tokens);
|
|
check(bucket <= maximum, "bucket never exceeds the request total");
|
|
check(bucket >= std::min(prompt + std::min(tokens, int64_t{128}), maximum),
|
|
"bucket never undercuts the generated window it needs");
|
|
check(bucket >= previous, "bucket grows with the request");
|
|
previous = bucket;
|
|
}
|
|
}
|
|
std::cout << "bucketed_cache_steps done\n";
|
|
}
|
|
|
|
void test_direct_import() {
|
|
constexpr int64_t steps = 8, elems = 3;
|
|
auto ordinary = make_tensors(steps, elems);
|
|
auto direct = make_tensors(steps, elems);
|
|
auto reference = make_cache(ordinary, steps, elems, {});
|
|
runtime::TransformerKVCacheOptions options;
|
|
options.direct_import = true;
|
|
auto cache = make_cache(direct, steps, elems, options);
|
|
const std::vector<float> poison(steps * elems, std::numeric_limits<float>::quiet_NaN());
|
|
for (auto * tensor : {ordinary.keys.tensor, ordinary.values.tensor, direct.keys.tensor, direct.values.tensor}) {
|
|
ggml_backend_tensor_set(tensor, poison.data(), 0, poison.size() * sizeof(float));
|
|
}
|
|
for (const int64_t length : {6, 2, 0, 8, 3}) {
|
|
runtime::TransformerKVState state;
|
|
state.current_end = length;
|
|
state.layers.resize(1);
|
|
auto & layer = state.layers.front();
|
|
layer.valid_steps = length;
|
|
for (int64_t i = 0; i < length * elems; ++i) {
|
|
layer.key.push_back(static_cast<float>(i - 4) / 7.0F);
|
|
layer.value.push_back(i == 0 ? -0.0F : static_cast<float>(i + 2) / -11.0F);
|
|
}
|
|
reference.import_state(state);
|
|
cache.import_state(state);
|
|
for (const auto & tensors : {std::pair{ordinary.keys.tensor, direct.keys.tensor},
|
|
std::pair{ordinary.values.tensor, direct.values.tensor}}) {
|
|
std::vector<float> expected(steps * elems), actual(steps * elems);
|
|
ggml_backend_tensor_get(tensors.first, expected.data(), 0, expected.size() * sizeof(float));
|
|
ggml_backend_tensor_get(tensors.second, actual.data(), 0, actual.size() * sizeof(float));
|
|
check(std::memcmp(expected.data(), actual.data(), actual.size() * sizeof(float)) == 0,
|
|
"direct import preserves every prefix and unused-tail byte");
|
|
}
|
|
const auto exported = cache.export_state();
|
|
check(exported.current_end == length && cache.valid_steps() == length, "direct import counters");
|
|
check(exported.layers.front().key == layer.key && exported.layers.front().value == layer.value,
|
|
"direct import/export round trip");
|
|
}
|
|
// Poison rows that a failed or truncated direct append could leave behind.
|
|
ggml_backend_tensor_set(direct.keys.tensor, poison.data(), 0, poison.size() * sizeof(float));
|
|
cache.advance_after_direct_append(2);
|
|
cache.retain_prefix(1);
|
|
runtime::TransformerKVState empty;
|
|
empty.layers.resize(1);
|
|
cache.import_state(empty);
|
|
std::vector<float> cleared(steps * elems);
|
|
ggml_backend_tensor_get(direct.keys.tensor, cleared.data(), 0, cleared.size() * sizeof(float));
|
|
check(cleared == std::vector<float>(steps * elems, 0.0F), "direct reset clears all stale rows after truncation");
|
|
for (const bool ring : {false, true}) {
|
|
auto invalid_options = options;
|
|
invalid_options.ring_mode = ring;
|
|
invalid_options.allow_f16_storage = true;
|
|
auto tensors = make_tensors(steps, elems, ring ? GGML_TYPE_F32 : GGML_TYPE_F16);
|
|
bool rejected = false;
|
|
try { (void)make_cache(tensors, steps, elems, invalid_options); }
|
|
catch (const std::runtime_error &) { rejected = true; }
|
|
check(rejected, "direct import rejects unsupported storage and ring mode");
|
|
}
|
|
std::cout << "direct import done\n";
|
|
}
|
|
|
|
} // namespace
|
|
|
|
int main() {
|
|
try {
|
|
test_slot_mapping();
|
|
test_option_validation();
|
|
test_wrap_round_trip();
|
|
test_wrap_advance();
|
|
test_non_ring_regression();
|
|
test_bucketed_cache_steps();
|
|
test_direct_import();
|
|
} catch (const std::exception & e) {
|
|
std::cerr << "ERROR: " << e.what() << "\n";
|
|
return 1;
|
|
}
|
|
if (failures != 0) {
|
|
std::cerr << failures << " check(s) failed\n";
|
|
return 1;
|
|
}
|
|
std::cout << "transformer_kv_ring: all checks passed\n";
|
|
return 0;
|
|
}
|