upstream-archive/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-07-23 11:10:55 -05:00
Code Issues Packages Projects Releases Wiki Activity

DeepseekV4: clear cache only for seq rather than full (#25521)

Browse Source
This commit is contained in:
Aman Gupta
2026-07-11 23:35:45 +08:00
committed by GitHub
parent c92e806d1c
commit 13f2b28b09
3 changed files with 154 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

74
src/llama-kv-cache-dsv4.cpp
View File

@@ -29,6 +29,15 @@ static uint32_t dsv4_comp_size(uint32_t kv_size, uint32_t ratio) {
return std::max<uint32_t>(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);
}
//

6
src/llama-kv-cache-dsv4.h
View File

@@ -21,7 +21,7 @@ public:
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;
@@ -67,6 +67,8 @@ private:
// 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:
@@ -146,7 +148,7 @@ private:
std::unique_ptr<llama_dsv4_comp_state> hca_state;
std::unique_ptr<llama_dsv4_comp_state> 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

100
tests/test-save-load-state.cpp
View File

@@ -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<uint8_t> & 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<uint8_t> 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<uint8_t> 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;
}