From 01e5e57f1fa22a9db53db4d21d0e0e808865af3a Mon Sep 17 00:00:00 2001 From: Aman Gupta Date: Fri, 10 Jul 2026 16:49:11 +0800 Subject: [PATCH 1/2] DeepseekV4: clear cache only for seq rather than full --- src/llama-kv-cache-dsv4.cpp | 74 +++++++++++++++++++----- src/llama-kv-cache-dsv4.h | 4 +- tests/test-save-load-state.cpp | 100 ++++++++++++++++++++++++++++++--- 3 files changed, 152 insertions(+), 26 deletions(-) diff --git a/src/llama-kv-cache-dsv4.cpp b/src/llama-kv-cache-dsv4.cpp index 9fccf347ed66..ebafac091993 100644 --- a/src/llama-kv-cache-dsv4.cpp +++ b/src/llama-kv-cache-dsv4.cpp @@ -29,6 +29,15 @@ static uint32_t dsv4_comp_size(uint32_t kv_size, uint32_t ratio) { return std::max(1, (kv_size + ratio - 1)/ratio); } +static void dsv4_clear_tensor_stream(ggml_tensor * tensor, uint32_t stream) { + GGML_ASSERT(ggml_is_contiguous(tensor)); + GGML_ASSERT(tensor->ne[3] == 1); + GGML_ASSERT(stream < (uint32_t) tensor->ne[2]); + + const size_t stream_size = tensor->nb[2]; + ggml_backend_tensor_memset(tensor, 0, stream*stream_size, stream_size); +} + static int64_t dsv4_stream_offset(uint32_t n_stream, llama_seq_id seq_id, uint32_t size) { if (n_stream <= 1) { return 0; @@ -781,11 +790,20 @@ llama_dsv4_comp_state::llama_dsv4_comp_state( __func__, name, ratio, state_size, n_embd_state, n_stream, layers.size(), total_size()/1024.0/1024.0); } -void llama_dsv4_comp_state::clear(bool data) { +void llama_dsv4_comp_state::clear(llama_seq_id seq_id, bool data) { if (!data) { return; } + if (seq_id >= 0) { + GGML_ASSERT((uint32_t) seq_id < n_stream); + for (const auto & layer : layers) { + dsv4_clear_tensor_stream(layer.kv, (uint32_t) seq_id); + dsv4_clear_tensor_stream(layer.score, (uint32_t) seq_id); + } + return; + } + for (auto & [_, buf] : ctxs_bufs) { ggml_backend_buffer_clear(buf.get(), 0); } @@ -1034,7 +1052,7 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4( // graph does not necessarily overwrite; uninitialized buffer contents would // otherwise leak in (instance-specific garbage) and corrupt recall. Zero all // compressed buffers up front so reads of un-written rows are deterministic. - clear_compressed(true); + clear_compressed(-1, true); } llama_memory_context_ptr llama_kv_cache_dsv4::init_batch( @@ -1147,7 +1165,7 @@ bool llama_kv_cache_dsv4::get_can_shift() const { void llama_kv_cache_dsv4::clear(bool data) { kv_raw->clear(data); - clear_compressed(true); // DSV4 compressed buffers must never expose stale/uninit rows + clear_compressed(-1, true); // DSV4 compressed buffers must never expose stale/uninit rows } bool llama_kv_cache_dsv4::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) { @@ -1169,7 +1187,7 @@ bool llama_kv_cache_dsv4::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1 const bool res = kv_raw->seq_rm(seq_id, p0, p1); if (res) { - clear_compressed(true); + clear_compressed(seq_id, true); } return res; @@ -1177,22 +1195,29 @@ bool llama_kv_cache_dsv4::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1 void llama_kv_cache_dsv4::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) { kv_raw->seq_cp(seq_id_src, seq_id_dst, p0, p1); - clear_compressed(true); } void llama_kv_cache_dsv4::seq_keep(llama_seq_id seq_id) { + GGML_ASSERT(seq_id >= 0 && (uint32_t) seq_id < n_seq_max); + kv_raw->seq_keep(seq_id); - clear_compressed(true); + + for (llama_seq_id id = 0; id < (llama_seq_id) n_seq_max; ++id) { + if (id == seq_id) { + continue; + } + + kv_raw->seq_rm(id, -1, -1); + clear_compressed(id, true); + } } void llama_kv_cache_dsv4::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) { kv_raw->seq_add(seq_id, p0, p1, shift); - clear_compressed(true); } void llama_kv_cache_dsv4::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) { kv_raw->seq_div(seq_id, p0, p1, d); - clear_compressed(true); } llama_pos llama_kv_cache_dsv4::seq_pos_min(llama_seq_id seq_id) const { @@ -1328,13 +1353,32 @@ llama_dsv4_comp_state * llama_kv_cache_dsv4::get_lid_state() const { return lid_state.get(); } -void llama_kv_cache_dsv4::clear_compressed(bool data) { - kv_csa->clear(data); - kv_hca->clear(data); - kv_lid->clear(data); - csa_state->clear(data); - hca_state->clear(data); - lid_state->clear(data); +void llama_kv_cache_dsv4::clear_compressed(llama_seq_id seq_id, bool data) { + if (seq_id < 0) { + kv_csa->clear(data); + kv_hca->clear(data); + kv_lid->clear(data); + } else { + GGML_ASSERT((uint32_t) seq_id < n_seq_max); + + const auto clear_seq = [seq_id, data](llama_kv_cache * kv) { + kv->seq_rm(seq_id, -1, -1); + + if (data) { + for (uint32_t il : kv->get_layer_ids()) { + dsv4_clear_tensor_stream(kv->get_k_storage(il), (uint32_t) seq_id); + } + } + }; + + clear_seq(kv_csa.get()); + clear_seq(kv_hca.get()); + clear_seq(kv_lid.get()); + } + + csa_state->clear(seq_id, data); + hca_state->clear(seq_id, data); + lid_state->clear(seq_id, data); } // diff --git a/src/llama-kv-cache-dsv4.h b/src/llama-kv-cache-dsv4.h index 772b428cd79f..07c170a74ad9 100644 --- a/src/llama-kv-cache-dsv4.h +++ b/src/llama-kv-cache-dsv4.h @@ -21,7 +21,7 @@ class llama_dsv4_comp_state { const char * name, const llama_memory_i::layer_filter_cb & filter); - void clear(bool data); + void clear(llama_seq_id seq_id, bool data); uint32_t get_ratio() const; uint32_t get_state_size() const; @@ -146,7 +146,7 @@ class llama_kv_cache_dsv4 : public llama_memory_i { std::unique_ptr hca_state; std::unique_ptr lid_state; - void clear_compressed(bool data); + void clear_compressed(llama_seq_id seq_id, bool data); }; // DSV4 raw attention only uses the SWA half of kv_raw. The base half is kept diff --git a/tests/test-save-load-state.cpp b/tests/test-save-load-state.cpp index b097d752ab71..6e93ce6fb8da 100644 --- a/tests/test-save-load-state.cpp +++ b/tests/test-save-load-state.cpp @@ -78,7 +78,84 @@ static llama_tokens test_baseline(struct llama_model * model, const struct commo } -// Test 2: state load +// Test 2: sequence removal isolation +// - decode the same prefix into two sequences +// - remove sequence 0 +// - verify that sequence 1 remains unchanged +static bool test_seq_rm_isolated( + struct llama_model * model, + const struct common_params & params, + const llama_tokens & tokens) { + 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)}; + if (!ctx) { + LOG_ERR("%s: failed to create context\n", __func__); + return false; + } + + LOG("\n=== Test 2: sequence removal isolation ===\n"); + + const size_t n_tokens = tokens.size() < 128 ? tokens.size() : 128; + for (llama_seq_id seq_id = 0; seq_id < 2; ++seq_id) { + llama_batch_ptr batch(n_tokens, 0, 1); + for (size_t i = 0; i < n_tokens; ++i) { + common_batch_add(batch.get(), tokens[i], i, { seq_id }, false); + } + + if (llama_decode(ctx.get(), batch.get())) { + LOG_ERR("%s: failed to decode prompt for sequence %d\n", __func__, seq_id); + return false; + } + } + + const auto get_seq_state = [&](llama_seq_id seq_id, std::vector & state) { + const size_t state_size = llama_state_seq_get_size(ctx.get(), seq_id); + 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(ctx.get(), state.data(), state.size(), seq_id); + 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; + }; + + std::vector state_before; + if (!get_seq_state(1, state_before)) { + return false; + } + + 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; + } + + std::vector state_after; + if (!get_seq_state(1, state_after)) { + return false; + } + + if (state_before != state_after) { + LOG_ERR("%s: removing sequence 0 changed sequence 1\n", __func__); + return false; + } + + LOG("PASS\n"); + return true; +} + + +// Test 3: state load // - create a new context // - load state from file // - replay the last prompt token @@ -90,7 +167,7 @@ static bool test_state_load(struct llama_model * model, const struct common_para auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)}; llama_sampler_chain_add(smpl.get(), llama_sampler_init_dist(params.sampling.seed)); - LOG("\n=== Test 2: state load ===\n"); + LOG("\n=== Test 3: state load ===\n"); // Load state from file llama_tokens unused_sts(tokens.size()); @@ -126,7 +203,7 @@ static bool test_state_load(struct llama_model * model, const struct common_para } -// Test 3: seq copy (host) +// Test 4: seq copy (host) // - create a multi-seq context // - load state from file // - replay the last prompt token @@ -141,7 +218,7 @@ static bool test_seq_cp_host(struct llama_model * model, const struct common_par auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)}; llama_sampler_chain_add(smpl.get(), llama_sampler_init_dist(params.sampling.seed)); - LOG("\n=== Test 3: seq copy (host) ===\n"); + LOG("\n=== Test 4: seq copy (host) ===\n"); // Load state from file llama_tokens unused_sts(tokens.size()); @@ -198,7 +275,7 @@ static bool test_seq_cp_host(struct llama_model * model, const struct common_par } -// Test 4: seq copy (device) +// Test 5: seq copy (device) // - create a multi-seq context // - load state from file // - replay the last prompt token @@ -213,7 +290,7 @@ static bool test_seq_cp_device(struct llama_model * model, const struct common_p auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)}; llama_sampler_chain_add(smpl.get(), llama_sampler_init_dist(params.sampling.seed)); - LOG("\n=== Test 4: seq copy (device) ===\n"); + LOG("\n=== Test 5: seq copy (device) ===\n"); // Load state from file llama_tokens unused_sts(tokens.size()); @@ -337,17 +414,22 @@ int main(int argc, char ** argv) { return 1; } - // Test 2: state load + // Test 2: sequence removal isolation + if (!test_seq_rm_isolated(model, params, tokens)) { + return 1; + } + + // Test 3: state load if (!test_state_load(model, params, tokens, result_baseline)) { return 1; } - // Test 3: seq copy (host) + // Test 4: seq copy (host) if (!test_seq_cp_host(model, params, tokens, result_baseline)) { return 1; } - // Test 4: seq copy (device) + // Test 5: seq copy (device) if (!test_seq_cp_device(model, params, tokens, result_baseline)) { return 1; } From fcf2bb7e54be63603d348b115dd6ec602fa4b481 Mon Sep 17 00:00:00 2001 From: Aman Gupta Date: Fri, 10 Jul 2026 18:50:42 +0800 Subject: [PATCH 2/2] add comment about dsv4 cache limitations --- src/llama-kv-cache-dsv4.h | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/llama-kv-cache-dsv4.h b/src/llama-kv-cache-dsv4.h index 07c170a74ad9..91619ce9e1b0 100644 --- a/src/llama-kv-cache-dsv4.h +++ b/src/llama-kv-cache-dsv4.h @@ -67,6 +67,8 @@ class llama_dsv4_comp_state { // DSV4 uses a normal raw/SWA token cache plus compressed K-only block caches. // The compressed caches are storage only; DSV4-specific visibility and block // planning are handled by llama_kv_cache_dsv4_context / llm_graph_input_dsv4. +// FIXME: currently the cache only supports non-unified mode even if unified flag is passed +// FIXME: we currently conflate token_pos and buffer contents. See https://github.com/ggml-org/llama.cpp/pull/25521#discussion_r3558173819 class llama_kv_cache_dsv4 : public llama_memory_i { public: