Superminaren/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-09-04 02:37:27 +02:00
Code Issues Packages Projects Releases Wiki Activity

kv-cache : optimize restoring non-contiguous cells (#27991)

Browse Source 
* kv cache : batch state restore scatter reads per contiguous run

When restoring state into non-contiguous destination cells (e.g. a
prompt-cache snapshot into a fragmented ring), state_read_data issued
one small copy per KV cell - ~1.4M copies of a few KiB each for a
40k+ token restore, taking 25-63 s on the CUDA backend.

The snapshot stores cell rows in cell order, so a maximal run of
consecutive destination indices maps to one contiguous block and can
be restored with a single copy. Precompute the runs once and use them
in all three scatter loops (K, V, transposed V). Byte-identical.

The on-device reader copies with a byte cursor when the read and
write chunking differs, so the batched reads are safe for it as well.
Batching makes equal tensor counts with a different split reachable
(save ranges [2,1] vs restore runs [1,2]); the next commit teaches the
reader's 1:1 path to fall back to the byte cursor in that case.

Verified in a production setup: 1,363,616 copies / 25-63 s -> 224
copies / 221-424 ms for the same restores (42,603 cells, 4 runs).

Assisted-by: Claude Code (unsloth/qwen3.8-27b)

* context : fall back to the byte cursor when read and write chunking differ

the on-device reader copies saved state back with a 1:1 copy by tensor
index whenever the write and read sides recorded the same number of
tensors, guarded by a per-tensor size assert.

equal tensor counts do not imply equal chunking: a state restore may
batch its reads per contiguous run of destination cells while the save
used per-range reads, so both sides can record two tensors that split
the same data differently, and the assert aborts in all builds.

compare the per-tensor sizes and only take the 1:1 path when the
chunking actually matches, otherwise fall through to the existing
byte-cursor copy. both sides enumerate the same logical data in the
same order, so the cursor copy is well-defined across tensor
boundaries.

Assisted-by: Claude Code (unsloth/qwen3.8-27b)

* tests : cover state restore scatter reads on host and on-device paths

decode the same prefix on two sequences, interleaving the seq 0 cells
between the seq 1 cells, so the seq 1 cells are isolated from each
other in the kv cache (three cells, two saved ranges). save the seq 1
state, free the interleaved seq 0 cells, and restore: the destination
is then non-contiguous (two runs), and the restore-side chunking has
the same tensor count as the save-side with a different split, so the
scatter path is batched per contiguous run and the on-device reader's
byte-cursor fallback is exercised.

the restored state is saved again on the host and compared byte for
byte with the first save: the blob is serialized in sequence cell
order, so the two saves are identical if and only if the scatter
restore wrote exactly the same KV content. this documents the
byte-identical guarantee of the run-batched scatter reads.

one test per io backend: the host (CPU) path and the on-device path.

Assisted-by: Claude Code (unsloth/qwen3.8-27b)
This commit is contained in:
itsnotoger
2026-08-31 19:49:58 +03:00
committed by GitHub
parent 010be9683a
commit 2d8d612e4c
3 changed files with 147 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/llama-context.cpp
+16 -5
View File
@@ -2907,17 +2907,28 @@ public:
}
if (mbuf_cur.n_tensors == mbuf.n_tensors) {
// same chunking: copy 1:1 by index
// an equal tensor count does not imply the same chunking, e.g. save ranges [2,1] vs restore runs [1,2]
bool same_chunking = true;
for (size_t i = 0; i < mbuf_cur.org.size(); ++i) {
GGML_ASSERT(ggml_nbytes(mbuf_cur.cpy[i]) == ggml_nbytes(mbuf.org[i]));
ggml_backend_tensor_copy(mbuf_cur.cpy[i], mbuf.org[i]);
if (ggml_nbytes(mbuf_cur.cpy[i]) != ggml_nbytes(mbuf.org[i])) {
same_chunking = false;
break;
}
}
if (same_chunking) {
// same chunking: copy 1:1 by index
for (size_t i = 0; i < mbuf_cur.org.size(); ++i) {
ggml_backend_tensor_copy(mbuf_cur.cpy[i], mbuf.org[i]);
}
continue;
}
continue;
}
// different chunking: copy the write-side data (mbuf_cur.cpy) into the read-side targets (mbuf.org)
// with a byte cursor. Write and read enumerate the same logical data in the same order but may chunk
// it differently, so copy across tensor boundaries rather than 1:1 by index.
// it differently (even with an equal number of tensors), so copy across tensor boundaries rather than
// 1:1 by index.
const size_t total = mbuf_cur.total_size;
ggml_init_params params_scratch = {
src/llama-kv-cache.cpp
+26 -38
View File
@@ -2533,6 +2533,24 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32
bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32_t cell_count, const slot_info & sinfo) {
auto & cells = v_cells[strm];
// batch the scatter reads per contiguous run of destination indices
// from inclusive, to exclusive - same convention as cell_ranges_t
// contiguous cells yield a single run covering the whole block
struct cell_run { uint32_t from; uint32_t to; };
std::vector<cell_run> runs;
if (cell_count > 0) {
const auto & idxs = sinfo.idxs[0];
uint32_t i0 = 0;
while (i0 < cell_count) {
uint32_t i1 = i0 + 1;
while (i1 < cell_count && idxs[i1] == idxs[i1 - 1] + 1) {
++i1;
}
runs.push_back({idxs[i0], idxs[i1 - 1] + 1});
i0 = i1;
}
}
uint32_t v_trans;
uint32_t n_layer;
@@ -2580,17 +2598,8 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32
return false;
}
if (cell_count) {
if (sinfo.is_contiguous()) {
// Fast path: contiguous cells, single memcpy
io.read_tensor(k, sinfo.head() * k_size_row, cell_count * k_size_row);
} else {
// Slow path: scatter to non-contiguous positions
for (uint32_t i = 0; i < cell_count; ++i) {
const size_t dst_offset = sinfo.idxs[0][i] * k_size_row;
io.read_tensor(k, dst_offset, k_size_row);
}
}
for (const auto & r : runs) {
io.read_tensor(k, (size_t) r.from * k_size_row, (size_t) (r.to - r.from) * k_size_row);
}
}
@@ -2623,17 +2632,8 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32
return false;
}
if (cell_count) {
if (sinfo.is_contiguous()) {
// Fast path: contiguous cells, single memcpy
io.read_tensor(v, sinfo.head() * v_size_row, cell_count * v_size_row);
} else {
// Slow path: scatter to non-contiguous positions
for (uint32_t i = 0; i < cell_count; ++i) {
const size_t dst_offset = sinfo.idxs[0][i] * v_size_row;
io.read_tensor(v, dst_offset, v_size_row);
}
}
for (const auto & r : runs) {
io.read_tensor(v, (size_t) r.from * v_size_row, (size_t) (r.to - r.from) * v_size_row);
}
}
} else {
@@ -2674,22 +2674,10 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32
return false;
}
if (cell_count) {
if (sinfo.is_contiguous()) {
// Fast path: contiguous cells
const uint32_t h = sinfo.head();
for (uint32_t j = 0; j < n_embd_v_gqa; ++j) {
const size_t dst_offset = (h + j * cells.size()) * v_size_el;
io.read_tensor(v, dst_offset, cell_count * v_size_el);
}
} else {
// Slow path: scatter to non-contiguous positions
for (uint32_t j = 0; j < n_embd_v_gqa; ++j) {
for (uint32_t i = 0; i < cell_count; ++i) {
const size_t dst_offset = (sinfo.idxs[0][i] + j * cells.size()) * v_size_el;
io.read_tensor(v, dst_offset, v_size_el);
}
}
for (uint32_t j = 0; j < n_embd_v_gqa; ++j) {
for (const auto & r : runs) {
const size_t dst_offset = ((size_t) r.from + j * cells.size()) * v_size_el;
io.read_tensor(v, dst_offset, (size_t) (r.to - r.from) * v_size_el);
}
}
}
tests/test-save-load-state.cpp
+105 -1
View File
@@ -355,7 +355,101 @@ static bool test_seq_cp_device(struct llama_model * model, const struct common_p
}
// Run the full save/load test suite (tests 1-5) for a single model.
// Test 6/7: seq copy (scatter)
// - 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
static bool test_seq_cp_scatter(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens, int test_num, bool on_device) {
auto params_ctx = common_context_params_to_llama(params);
params_ctx.n_ctx = 256;
params_ctx.n_seq_max = 2;
params_ctx.kv_unified = true;
auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)};
LOG("\n=== Test %d: seq copy (%s, scatter) ===\n", test_num, on_device ? "device" : "host");
const uint32_t flags = on_device ? LLAMA_STATE_SEQ_FLAGS_ON_DEVICE : LLAMA_STATE_SEQ_FLAGS_NONE;
auto decode_one = [&](llama_token tok, int pos, llama_seq_id seq) {
llama_batch_ptr batch(1, 0, 1);
common_batch_add(batch.get(), tok, pos, { seq }, false);
return llama_decode(ctx.get(), batch.get()) == 0;
};
// seq 0 cells 0,1,4 interleave the seq 1 cells 2,3,5
if (!decode_one(tokens[0], 0, 0) ||
!decode_one(tokens[1], 1, 0) ||
!decode_one(tokens[0], 0, 1) ||
!decode_one(tokens[1], 1, 1) ||
!decode_one(tokens[2], 2, 0) ||
!decode_one(tokens[2], 2, 1)) {
LOG_ERR("%s: failed to build interleaved state\n", __func__);
return false;
}
const auto get_seq_state = [&](llama_seq_id seq_id, uint32_t fl, std::vector<uint8_t> & state) {
const size_t state_size = llama_state_seq_get_size_ext(ctx.get(), seq_id, fl);
if (state_size == 0) {
LOG_ERR("%s: sequence state is empty\n", __func__);
return false;
}
state.resize(state_size);
const size_t ncopy = llama_state_seq_get_data_ext(ctx.get(), state.data(), state.size(), seq_id, fl);
if (ncopy != state.size()) {
LOG_ERR("%s: sequence state length %zu does not match expected length %zu\n",
__func__, ncopy, state.size());
return false;
}
return true;
};
// host blob: contains the KV data, used for the byte-for-byte comparison
std::vector<uint8_t> state_before;
if (!get_seq_state(1, LLAMA_STATE_SEQ_FLAGS_NONE, state_before)) {
return false;
}
// save via the io path under test
std::vector<uint8_t> state_save;
if (!get_seq_state(1, flags, state_save)) {
return false;
}
LOG_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 (!llama_memory_seq_rm(llama_get_memory(ctx.get()), 0, -1, -1)) {
LOG_ERR("%s: failed to remove sequence 0\n", __func__);
return false;
}
// restore via the io path under test
const size_t nset = llama_state_seq_set_data_ext(ctx.get(), state_save.data(), state_save.size(), 1, flags);
if (nset != state_save.size()) {
LOG_ERR("%s: seq set data length %zu does not match expected length %zu\n", __func__, nset, state_save.size());
return false;
}
LOG_TRC("%s: seq 1 restored via %s, %zu bytes\n", __func__, on_device ? "device" : "host", nset);
std::vector<uint8_t> state_after;
if (!get_seq_state(1, LLAMA_STATE_SEQ_FLAGS_NONE, state_after)) {
return false;
}
// the blob is serialized in sequence cell order, so identical bytes iff the restore wrote the same KV
if (state_before.size() != state_after.size() || memcmp(state_before.data(), state_after.data(), state_before.size()) != 0) {
LOG_ERR("\n%s: error: restored KV state is not byte-identical to the saved state\n", __func__);
return false;
}
LOG("\nPASS\n");
return true;
}
// Run the full save/load test suite (tests 1-7) for a single model.
// Returns true if all tests pass, false otherwise.
static bool run_save_load_tests_for_model(const std::string & model_path, const struct common_params & base_params) {
struct common_params params = base_params;
@@ -422,6 +516,16 @@ static bool run_save_load_tests_for_model(const std::string & model_path, const
return false;
}
// Test 6: seq copy (host, scatter)
if (!test_seq_cp_scatter(model, params, tokens, 6, false)) {
return false;
}
// Test 7: seq copy (device, scatter)
if (!test_seq_cp_scatter(model, params, tokens, 7, true)) {
return false;
}
LOG("\nAll tests passed.\n");
return true;