From 2d8d612e4c68d3801e556a1b4a028f55ec33ecbb Mon Sep 17 00:00:00 2001 From: itsnotoger <19309683+itsnotoger@users.noreply.github.com> Date: Mon, 31 Aug 2026 18:49:58 +0200 Subject: [PATCH] kv-cache : optimize restoring non-contiguous cells (#27991) * 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) --- src/llama-context.cpp | 21 +++++-- src/llama-kv-cache.cpp | 64 ++++++++------------ tests/test-save-load-state.cpp | 106 ++++++++++++++++++++++++++++++++- 3 files changed, 147 insertions(+), 44 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index a6d2f264fc..a920c42311 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -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 = { diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 65afbd8c37..0e095f8b6b 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -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 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); } } } diff --git a/tests/test-save-load-state.cpp b/tests/test-save-load-state.cpp index 0ceab7c545..35c7690587 100644 --- a/tests/test-save-load-state.cpp +++ b/tests/test-save-load-state.cpp @@ -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 & 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 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 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 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;