// prefix sharing for prompt cache states (tools/server/server-kv-share.cpp): a delta decodes to the original state
// byte for byte, shares only rows whose bytes equal the prefix's, and is refused against any other prefix
#include "../tools/server/server-kv-share.h"
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <string>
#include <vector>
static int n_checks = 0;
static void check(bool ok, const char * what) {
++n_checks;
if (!ok) {
fprintf(stderr, "FAIL: %s\n", what);
exit(1);
}
}
template <typename T> static void put(std::vector<uint8_t> & b, T v) {
const auto * p = reinterpret_cast<const uint8_t *>(&v);
b.insert(b.end(), p, p + sizeof(T));
}
// a row's bytes for a cell at `pos` in layer `l`: a function of both, so equal positions give equal rows
static void row(std::vector<uint8_t> & b, size_t n, int32_t pos, uint32_t l, uint8_t salt) {
for (size_t i = 0; i < n; ++i) {
b.push_back(uint8_t(pos * 31 + l * 7 + i + salt));
}
}
struct cell { int32_t pos; uint8_t salt; };
// a sequence state in llama_kv_cache::state_write's layout, then an opaque recurrent part
static std::vector<uint8_t> state(const std::vector<cell> & cells, bool mrope, bool v_trans, size_t tail, uint8_t tail_seed) {
const uint32_t n_layer = 3;
const size_t k_row = 64, v_row = 48, el = 2;
std::vector<uint8_t> b;
put<uint32_t>(b, 0x4c4c414d); // magic
put<int32_t>(b, 1); // seq id
put<uint32_t>(b, 1); // n_stream
put<uint32_t>(b, (uint32_t) cells.size());
for (const auto & c : cells) {
put<int32_t>(b, c.pos);
put<uint32_t>(b, 1);
if (mrope) {
put<int32_t>(b, c.pos % 5);
put<int32_t>(b, 0);
}
put<int32_t>(b, 1);
}
put<uint32_t>(b, v_trans ? 1 : 0);
put<uint32_t>(b, n_layer);
for (uint32_t l = 0; l < n_layer; ++l) {
put<int32_t>(b, 1);
put<uint64_t>(b, k_row);
for (const auto & c : cells) {
row(b, k_row, c.pos, l, c.salt);
}
}
for (uint32_t l = 0; l < n_layer; ++l) {
put<int32_t>(b, 1);
if (!v_trans) {
put<uint64_t>(b, v_row);
for (const auto & c : cells) {
row(b, v_row, c.pos, l + 100, c.salt);
}
} else {
const uint32_t n_embd = v_row / el;
put<uint32_t>(b, (uint32_t) el);
put<uint32_t>(b, n_embd);
for (uint32_t j = 0; j < n_embd; ++j) {
for (const auto & c : cells) {
row(b, el, c.pos, l + 200 + j, c.salt);
}
}
}
}
for (size_t i = 0; i < tail; ++i) {
b.push_back(uint8_t(i * 13 + tail_seed));
}
return b;
}
static std::vector<cell> range(int32_t from, int32_t to, uint8_t salt = 0) {
std::vector<cell> v;
for (int32_t p = from; p < to; ++p) {
v.push_back({ p, salt });
}
return v;
}
static bool round_trip(const std::vector<uint8_t> & child, const kv_share_parent & parent, const kv_share_layout & layout, size_t & shared) {
std::vector<uint8_t> delta, back;
std::string err;
return kv_share_encode(child, parent, layout, delta, shared) && kv_share_is_delta(delta) &&
kv_share_parent_id(delta) == parent.id && delta.size() < child.size() &&
kv_share_decode(delta, parent, back, err) && back == child;
}
int main() {
const int32_t P = 60000, N = 90000; // ~45 MiB of shared rows: above the 16 MiB floor
for (int mode = 0; mode < 4; ++mode) {
const bool mrope = mode & 1, v_trans = mode & 2;
const kv_share_layout layout { mrope ? 8u : 0u };
const auto parent = state(range(0, P), mrope, v_trans, 4096, 1);
const kv_share_parent pp { 0x1234, &parent };
// the conversation continues the prefix: its rows for the prefix's positions are shared, the rest and the
// recurrent part are its own
size_t shared = 0;
const auto child = state(range(0, N), mrope, v_trans, 8192, 9);
check(round_trip(child, pp, layout, shared), "a continuation decodes to itself");
check(shared == (size_t) P * (64 + 48) * 3, "every prefix row is shared, nothing else");
// cells in another order (a unified KV with two slots interleaves them by ubatch) match by position
auto blocks = range(0, N);
for (size_t i = 0; i + 1024 <= blocks.size(); i += 1024) {
std::swap_ranges(blocks.begin() + i, blocks.begin() + i + 512, blocks.begin() + i + 512);
}
check(round_trip(state(blocks, mrope, v_trans, 100, 3), pp, layout, shared), "a reordered child decodes to itself");
check(shared == (size_t) P * (64 + 48) * 3, "and shares the same rows");
// shared rows scattered one cell at a time are not worth their copy records
auto pairs = range(0, N);
for (size_t i = 0; i + 1 < pairs.size(); i += 2) {
std::swap(pairs[i], pairs[i + 1]);
}
std::vector<uint8_t> scattered;
check(!kv_share_encode(state(pairs, mrope, v_trans, 100, 3), pp, layout, scattered, shared) ||
scattered.size() + (16u << 20) <= state(pairs, mrope, v_trans, 100, 3).size(),
"a delta is only kept when it saves 16 MiB");
// a prefix computed with other bits is stored literally where it differs
auto drifted = range(0, N);
for (int32_t p = 0; p < P; p += 10) {
drifted[p].salt = 7;
}
// (a transposed V splits every cell into one element per row: a delta there costs more than it saves)
if (!v_trans) {
check(round_trip(state(drifted, mrope, v_trans, 100, 3), pp, layout, shared), "a child that differs in places decodes to itself");
check(shared < (size_t) P * (64 + 48) * 3, "and shares only the equal rows");
} else {
std::vector<uint8_t> d;
const auto whole = state(drifted, mrope, v_trans, 100, 3);
check(!kv_share_encode(whole, pp, layout, d, shared) || d.size() + (16u << 20) <= whole.size(), "a scattered transposed delta is kept only if it saves 16 MiB");
}
// too little in common: no delta
std::vector<uint8_t> delta;
check(!kv_share_encode(state(range(0, 100), mrope, v_trans, 10, 3), pp, layout, delta, shared) && delta.empty(), "a short child is kept whole");
check(!kv_share_encode(state(range(P, N), mrope, v_trans, 10, 3), pp, layout, delta, shared), "an unrelated child is kept whole");
}
// a delta only decodes against the prefix it was written against
{
const kv_share_layout layout { 0 };
const auto parent = state(range(0, P), false, false, 64, 1);
const auto child = state(range(0, N), false, false, 64, 2);
std::vector<uint8_t> delta, back;
size_t shared = 0;
std::string err;
check(kv_share_encode(child, { 7, &parent }, layout, delta, shared), "encoded");
check(!kv_share_decode(delta, { 8, &parent }, back, err), "another id is refused");
auto other = parent;
other[other.size() - 1] ^= 1;
check(!kv_share_decode(delta, { 7, &other }, back, err), "other bytes under the same id are refused");
check(!kv_share_decode(child, { 7, &parent }, back, err) && !kv_share_is_delta(child), "a whole state is not a delta");
auto cut = delta;
cut.resize(cut.size() - 5);
check(!kv_share_decode(cut, { 7, &parent }, back, err), "a truncated delta is refused");
check(!kv_share_encode(delta, { 7, &parent }, layout, back, shared), "a delta is not encoded again");
}
// positions twice (image chunks share one): nothing ambiguous is shared
{
const kv_share_layout layout { 0 };
auto cells = range(0, P);
cells.push_back({ 5, 0 });
const auto parent = state(cells, false, false, 64, 1);
std::vector<uint8_t> delta;
size_t shared = 0;
check(!kv_share_encode(state(range(0, N), false, false, 64, 2), { 1, &parent }, layout, delta, shared), "a prefix with a repeated position shares nothing");
}
printf("test-kv-share: %d checks passed\n", n_checks);
return 0;
}