Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
am17an merged 2 commits into from
Jul 11, 2026
+154 −26
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
01e5e57
DeepseekV4: clear cache only for seq rather than full
am17an Jul 10, 2026
fcf2bb7
add comment about dsv4 cache limitations
am17an Jul 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
74 changes: 59 additions & 15 deletions src/llama-kv-cache-dsv4.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
}
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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) {
Expand All @@ -1169,30 +1187,37 @@ 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;
}

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 {
Expand Down Expand Up @@ -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);
}

//
Expand Down
6 changes: 4 additions & 2 deletions src/llama-kv-cache-dsv4.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -146,7 +148,7 @@ class llama_kv_cache_dsv4 : public llama_memory_i {
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
Expand Down
100 changes: 91 additions & 9 deletions tests/test-save-load-state.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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());
Expand Down Expand Up @@ -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
Expand All @@ -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());
Expand Down Expand Up @@ -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
Expand All @@ -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());
Expand Down Expand Up @@ -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;
}
Expand Down