Skip to content

Commit 5ec9c50

Browse files
AnbeeldAnbeeld
committed
feat(kvarn): group-range state serialization for prompt-cache compression
KVarN cache types with >4 bit widths (kvarn5, kvarn5/5, kvarn4/4 etc.) wrote full pre-allocated records tensors on state save, making checkpoints ~5-7x larger than needed. This caused the checkpoint RAM budget to evict all but one checkpoint at typical 18k+ token positions, leaving no viable restore point for prompt-cache reuse. Fix: bump KVAR_N_STATE_VERSION to 3 and serialize only the used portion of the records tensors, computed as n_groups_used = ceil((pos_max+1)/128) groups from metadata seq_pos_max. Stages are still written in full (only 384 tokens per stream per layer, negligible). On read, the unused tail is zero-filled via ggml_backend_tensor_memset (GPU-safe, no host allocation). The wire format change is additive: v3 writes n_groups_used after the stream list, v2 reads skip it. V2 readers reject v3 data via the version check; v3 readers handle v1 and v2 transparently. Other fixes in the same change: - write_kvarn_tensor_slice uses uint64_t for the size field on the wire (was size_t, which would be 4 bytes on 32-bit, breaking portability) - ggml_backend_tensor_memset replaces the host-allocated std::vector zero buffer, eliminating per-layer heap allocations on checkpoint restore Measured impact: kvarn5/5 checkpoint at 18k tokens drops from ~1992 MiB to ~658 MiB (3x reduction). Multiple checkpoints now fit in the default --cache-ram budget, enabling prompt-cache reuse for KVarN >4 types.
1 parent 84d4f1d commit 5ec9c50

2 files changed

Lines changed: 95 additions & 21 deletions

File tree

‎src/llama-kv-cache-kvarn.cpp‎

Lines changed: 87 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
#include "llama-kv-cache-kvarn.h"
22

3+
#include "ggml-backend.h"
34
#include "llama-context.h"
45
#include "llama-hparams.h"
56
#include "llama-impl.h"
@@ -11,13 +12,14 @@
1112
#include <limits>
1213
#include <map>
1314
#include <stdexcept>
15+
#include <vector>
1416

1517
namespace {
1618

1719
constexpr uint32_t KVAR_N_GROUP = 128;
1820
constexpr uint32_t KVAR_N_STAGE_GROUPS = 3;
1921
constexpr uint32_t KVAR_N_STATE_MAGIC = 0x4e52564b; // "KVRN"
20-
constexpr uint32_t KVAR_N_STATE_VERSION = 2;
22+
constexpr uint32_t KVAR_N_STATE_VERSION = 3;
2123

2224
bool kvarn_backend_supports_native_ops(ggml_backend_dev_t dev) {
2325
if (dev == nullptr || ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU) {
@@ -42,15 +44,39 @@ void write_kvarn_tensor(llama_io_write_i & io, ggml_tensor * tensor) {
4244
io.write_tensor(tensor, 0, size);
4345
}
4446

47+
void write_kvarn_tensor_slice(llama_io_write_i & io, ggml_tensor * tensor, size_t offset, size_t size) {
48+
GGML_ASSERT(offset + size <= (size_t) ggml_nbytes(tensor));
49+
const uint64_t size64 = size;
50+
io.write(&size64, sizeof(size64));
51+
io.write_tensor(tensor, offset, size);
52+
}
53+
4554
void read_kvarn_tensor(llama_io_read_i & io, ggml_tensor * tensor) {
4655
uint64_t size;
4756
io.read(&size, sizeof(size));
48-
if (size != ggml_nbytes(tensor)) {
57+
if (size != (uint64_t) ggml_nbytes(tensor)) {
4958
throw std::runtime_error("mismatched KVarN cache tensor size");
5059
}
5160
io.read_tensor(tensor, 0, size);
5261
}
5362

63+
void read_kvarn_tensor_slice(llama_io_read_i & io, ggml_tensor * tensor, size_t offset, size_t size) {
64+
GGML_ASSERT(offset + size <= (size_t) ggml_nbytes(tensor));
65+
uint64_t saved_size;
66+
io.read(&saved_size, sizeof(saved_size));
67+
if (saved_size != size) {
68+
throw std::runtime_error("mismatched KVarN cache tensor slice size");
69+
}
70+
io.read_tensor(tensor, offset, size);
71+
}
72+
73+
void zero_kvarn_tensor_range(ggml_tensor * tensor, size_t offset, size_t size) {
74+
if (size == 0) {
75+
return;
76+
}
77+
ggml_backend_tensor_memset(tensor, 0, offset, size);
78+
}
79+
5480
} // namespace
5581

5682
llama_kv_cache_kvarn_context::llama_kv_cache_kvarn_context(
@@ -662,18 +688,36 @@ void llama_kv_cache_kvarn::state_write(llama_io_write_i & io, llama_seq_id seq_i
662688
io.write(&stream, sizeof(stream));
663689
}
664690

691+
// n_groups_used is single-valued across all saved streams. This is correct
692+
// because when seq_id >= 0, saved_streams has exactly 1 entry (the stream
693+
// for that sequence), so n_groups_used applies to that one stream only.
694+
// When seq_id == -1, n_groups_used equals n_groups_per_stream (no compression).
695+
// If multi-stream partial writes are ever added, n_groups_used must become
696+
// per-stream.
697+
uint32_t n_groups_used = n_groups_per_stream;
698+
if (seq_id >= 0) {
699+
const llama_pos pos_max = metadata->seq_pos_max(seq_id);
700+
if (pos_max >= 0) {
701+
// ceil((pos_max + 1) / KVAR_N_GROUP) — covers all groups that the
702+
// materialize kernel might read. After restore the next store will
703+
// compute live_group ≤ pos_max / KVAR_N_GROUP < n_groups_used, so
704+
// every accessed group stays within the restored range.
705+
n_groups_used = std::min(n_groups_per_stream, (uint32_t) ((pos_max + KVAR_N_GROUP) / KVAR_N_GROUP));
706+
}
707+
}
708+
io.write(&n_groups_used, sizeof(n_groups_used));
709+
665710
for (const auto & layer : layers) {
666711
io.write(&layer.il, sizeof(layer.il));
667712
for (const uint32_t stream : saved_streams) {
668713
io.write(&stream, sizeof(stream));
669-
for (auto * tensor : {
670-
layer.k_records_stream[stream],
671-
layer.v_records_stream[stream],
672-
layer.k_stage_stream[stream],
673-
layer.v_stage_stream[stream],
674-
}) {
675-
write_kvarn_tensor(io, tensor);
676-
}
714+
715+
const size_t k_records_used = n_groups_used * layer.k_records_stream[stream]->nb[2];
716+
const size_t v_records_used = n_groups_used * layer.v_records_stream[stream]->nb[2];
717+
write_kvarn_tensor_slice(io, layer.k_records_stream[stream], 0, k_records_used);
718+
write_kvarn_tensor_slice(io, layer.v_records_stream[stream], 0, v_records_used);
719+
write_kvarn_tensor(io, layer.k_stage_stream[stream]);
720+
write_kvarn_tensor(io, layer.v_stage_stream[stream]);
677721
}
678722
}
679723
}
@@ -689,7 +733,7 @@ void llama_kv_cache_kvarn::state_read(llama_io_read_i & io, llama_seq_id seq_id,
689733
io.read(&version, sizeof(version));
690734
io.read(&type, sizeof(type));
691735
io.read(&n_layers, sizeof(n_layers));
692-
if (magic != KVAR_N_STATE_MAGIC || (version != 1 && version != KVAR_N_STATE_VERSION) ||
736+
if (magic != KVAR_N_STATE_MAGIC || version > KVAR_N_STATE_VERSION ||
693737
type != params.type || n_layers != layers.size()) {
694738
throw std::runtime_error("incompatible KVarN cache state");
695739
}
@@ -723,6 +767,14 @@ void llama_kv_cache_kvarn::state_read(llama_io_read_i & io, llama_seq_id seq_id,
723767
}
724768
}
725769

770+
uint32_t n_groups_used = n_groups_per_stream;
771+
if (version >= 3) {
772+
io.read(&n_groups_used, sizeof(n_groups_used));
773+
if (n_groups_used == 0 || n_groups_used > n_groups_per_stream) {
774+
throw std::runtime_error("invalid KVarN cache group count");
775+
}
776+
}
777+
726778
const uint32_t seq_stream = seq_id == -1 ? 0 : metadata->get_stream_for_seq(seq_id);
727779
if (seq_id != -1 && seq_stream >= n_stream) {
728780
throw std::runtime_error("invalid KVarN sequence stream");
@@ -743,13 +795,30 @@ void llama_kv_cache_kvarn::state_read(llama_io_read_i & io, llama_seq_id seq_id,
743795
}
744796

745797
const uint32_t stream_dst = seq_id == -1 ? stream : seq_stream;
746-
for (auto * tensor : {
747-
layer.k_records_stream[stream_dst],
748-
layer.v_records_stream[stream_dst],
749-
layer.k_stage_stream[stream_dst],
750-
layer.v_stage_stream[stream_dst],
751-
}) {
752-
read_kvarn_tensor(io, tensor);
798+
799+
if (version >= 3) {
800+
const size_t k_records_used = n_groups_used * layer.k_records_stream[stream_dst]->nb[2];
801+
const size_t v_records_used = n_groups_used * layer.v_records_stream[stream_dst]->nb[2];
802+
const size_t k_records_total = n_groups_per_stream * layer.k_records_stream[stream_dst]->nb[2];
803+
const size_t v_records_total = n_groups_per_stream * layer.v_records_stream[stream_dst]->nb[2];
804+
805+
read_kvarn_tensor_slice(io, layer.k_records_stream[stream_dst], 0, k_records_used);
806+
zero_kvarn_tensor_range(layer.k_records_stream[stream_dst], k_records_used, k_records_total - k_records_used);
807+
808+
read_kvarn_tensor_slice(io, layer.v_records_stream[stream_dst], 0, v_records_used);
809+
zero_kvarn_tensor_range(layer.v_records_stream[stream_dst], v_records_used, v_records_total - v_records_used);
810+
811+
read_kvarn_tensor(io, layer.k_stage_stream[stream_dst]);
812+
read_kvarn_tensor(io, layer.v_stage_stream[stream_dst]);
813+
} else {
814+
for (auto * tensor : {
815+
layer.k_records_stream[stream_dst],
816+
layer.v_records_stream[stream_dst],
817+
layer.k_stage_stream[stream_dst],
818+
layer.v_stage_stream[stream_dst],
819+
}) {
820+
read_kvarn_tensor(io, tensor);
821+
}
753822
}
754823
}
755824
}

‎tests/test-dflash-plumbing.cpp‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2255,12 +2255,17 @@ int main(int argc, char ** argv) {
22552255
ok &= expect(kv_cache_kvarn_cpp.find("GGML_ABORT(\"KVarN does not support position shifts\")") != std::string::npos &&
22562256
kv_cache_kvarn_cpp.find("GGML_ABORT(\"KVarN does not support position division\")") != std::string::npos,
22572257
"KVarN seq_add/seq_div must fail fast instead of logging and continuing");
2258-
ok &= expect(kv_cache_kvarn_cpp.find("constexpr uint32_t KVAR_N_STATE_VERSION = 2") != std::string::npos &&
2258+
ok &= expect(kv_cache_kvarn_cpp.find("constexpr uint32_t KVAR_N_STATE_VERSION = 3") != std::string::npos &&
22592259
kv_cache_kvarn_cpp.find("saved_streams") != std::string::npos &&
22602260
kv_cache_kvarn_cpp.find("metadata->get_stream_for_seq(seq_id)") != std::string::npos &&
22612261
kv_cache_kvarn_cpp.find("layer.k_records_stream[stream]") != std::string::npos &&
2262-
kv_cache_kvarn_cpp.find("layer.v_stage_stream[stream]") != std::string::npos,
2263-
"KVarN sequence state must serialize stream-scoped tensors instead of whole-cache tensors");
2262+
kv_cache_kvarn_cpp.find("layer.v_stage_stream[stream]") != std::string::npos &&
2263+
kv_cache_kvarn_cpp.find("n_groups_used") != std::string::npos &&
2264+
kv_cache_kvarn_cpp.find("write_kvarn_tensor_slice") != std::string::npos &&
2265+
kv_cache_kvarn_cpp.find("read_kvarn_tensor_slice") != std::string::npos &&
2266+
kv_cache_kvarn_cpp.find("ggml_backend_tensor_memset") != std::string::npos &&
2267+
kv_cache_kvarn_cpp.find("const uint64_t size64 = size") != std::string::npos,
2268+
"KVarN sequence state must serialize stream-scoped tensors with group-range compression");
22642269
ok &= expect(kv_cache_kvarn_cpp.find("pending_stream_copies") != std::string::npos &&
22652270
kv_cache_kvarn_cpp.find("llama_synchronize(lctx)") != std::string::npos &&
22662271
kv_cache_kvarn_cpp.find("copy_kvarn_stream") != std::string::npos,

0 commit comments

Comments
 (0)