Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

DeepseekV4: fix rollback with multi-seq (#26756)

Browse Source 
* DeepseekV4: fix rollback with multi-seq

* fix model loading

* make pending rollback single use

* only clear cache for seq_id for full load

* add assert for compress ratio

* make graph topology static

* pass true instead of flags in clear_compressed

* cont : clean-up + TODOs

---------

Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
This commit is contained in:
Aman Gupta
2026-08-23 13:57:49 +03:00
committed by GitHub
co-authored by Georgi Gerganov
parent d3371929bb
commit b0539c43ed
10 changed files with 345 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/llama-kv-cache-dsv4.cpp
+97 -40
View File
@@ -599,6 +599,33 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan(
}
}
if (ratio == DSV4_HCA_RATIO && !plan.state_pos.empty() && plan.state_write_idxs.empty()) {
assert(kv_size > 0);
// the last slot must not be live, or the dummy write would corrupt it;
// a full stream implies a completed block, which implies real writes
assert(plan.n_kv < (int64_t) kv_size);
// Keep the compress/write ops in the graph when no HCA block completes
// in this ubatch. The dummy block writes to the last cache slot and is
// masked out.
uint32_t i = 0;
while (i < ubatch.n_tokens && ubatch.pos[i] < 0) {
++i;
}
assert(i < ubatch.n_tokens);
const llama_seq_id seq_id = ubatch.seq_id[i][0];
const int64_t cache_off = dsv4_stream_offset(n_stream, seq_id, kv_size);
const int32_t source_idx = state_source_idx(seq_id, ubatch.pos[i]);
plan.state_write_idxs.push_back(cache_off + kv_size - 1);
plan.state_write_pos .push_back(0);
for (uint32_t j = 0; j < ratio; ++j) {
plan.state_read_idxs.push_back(source_idx);
}
}
if (overlap) {
// [ all blocks' prev-window indices | all blocks' cur-window indices ]
plan.state_read_idxs.reserve(overlap_prev_reads.size() + overlap_cur_reads.size());
@@ -608,7 +635,10 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan(
overlap_cur_reads.begin(), overlap_cur_reads.end());
}
plan.n_kv = GGML_PAD(plan.n_kv, 256u);
// Keep the mask (and with it the compressed-attention branch) present even
// before the first block is visible, so the graph topology never changes.
// Padded slots are masked out; comp cache buffers are zero-initialized.
plan.n_kv = std::max<int64_t>(GGML_PAD(plan.n_kv, 256u), 256);
std::sort(persist_rows.begin(), persist_rows.end(),
[](const persist_row & a, const persist_row & b) {
@@ -620,16 +650,26 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan(
plan.state_persist_dst_idxs.push_back(row.dst);
}
if (n_rs_seq > 0) {
for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) {
const llama_seq_id seq_id = ubatch.seq_id_unq[s];
if (seq_id < 0 || (uint32_t) seq_id >= n_stream) {
continue;
// Emit restore/snapshot entries for all layout streams so that the
// graph tensor sizes do not depend on the ubatch's sequence count.
// Streams not present in the ubatch get no-op entries.
for (uint32_t stream = 0; stream < n_stream; ++stream) {
llama_seq_id seq_id = -1;
if (n_stream == 1) {
// a unified stream serves any single sequence
seq_id = ubatch.n_seqs_unq > 0 ? ubatch.seq_id_unq[0] : -1;
} else {
for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) {
if (ubatch.seq_id_unq[s] == (llama_seq_id) stream) {
seq_id = ubatch.seq_id_unq[s];
break;
}
}
}
const int64_t stream_off = dsv4_stream_offset(n_stream, seq_id, state_size);
const uint32_t rollback = (uint32_t) seq_id < rs_idx.size() ? rs_idx[seq_id] : 0;
const int64_t stream_off = (int64_t) stream*state_size;
const uint32_t rollback = seq_id >= 0 && (uint32_t) seq_id < rs_idx.size() ? rs_idx[seq_id] : 0;
// Keep the restore graph fixed-width when no rollback is pending.
const int64_t src_plane = rollback > 0 && rollback <= n_rs_seq ? (int64_t) rollback*state_rows : 0;
for (uint32_t r = 0; r < state_size; ++r) {
@@ -639,35 +679,33 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan(
std::vector<uint32_t> token_idxs;
token_idxs.reserve(ubatch.n_tokens);
for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
if (dsv4_token_has_seq(ubatch, i, seq_id)) {
token_idxs.push_back(i);
if (seq_id >= 0) {
for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
if (dsv4_token_has_seq(ubatch, i, seq_id)) {
token_idxs.push_back(i);
}
}
}
if (token_idxs.empty()) {
continue;
}
const uint32_t n_seq_tokens = (uint32_t) token_idxs.size();
const int64_t scratch_off = (int64_t) state_rows*(1 + n_rs_seq);
for (uint32_t d = 1; d <= n_rs_seq; ++d) {
const int64_t dst_plane = (int64_t) d*state_rows;
const uint32_t prefix = d <= n_seq_tokens ? n_seq_tokens - d : 0;
for (uint32_t r = 0; r < state_size; ++r) {
int32_t src;
if (d <= n_seq_tokens) {
const uint32_t prefix = n_seq_tokens - d;
src = (int32_t) (stream_off + r);
int32_t src = (int32_t) (stream_off + r);
for (uint32_t j = 0; j < prefix; ++j) {
const uint32_t i_tok = token_idxs[j];
if (ubatch.pos[i_tok] >= 0 && (uint32_t) (ubatch.pos[i_tok]%state_size) == r) {
src = (int32_t) (scratch_off + i_tok);
}
for (uint32_t j = 0; j < prefix; ++j) {
const uint32_t i_tok = token_idxs[j];
if (ubatch.pos[i_tok] >= 0 && (uint32_t) (ubatch.pos[i_tok]%state_size) == r) {
src = (int32_t) (scratch_off + i_tok);
}
} else {
const int64_t src_plane = (int64_t) (d - n_seq_tokens)*state_rows;
src = (int32_t) (src_plane + stream_off + r);
}
if (n_seq_tokens == 0) {
// no-op: copy the snapshot plane onto itself
src = (int32_t) (dst_plane + stream_off + r);
}
plan.state_snapshot_src_idxs.push_back(src);
@@ -683,10 +721,16 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan(
}();
if (debug) {
LLAMA_LOG_INFO("%s: ratio=%u, n_tokens=%u, state_persist_dst=%s, state_write_pos=%s\n",
__func__, ratio, ubatch.n_tokens,
LLAMA_LOG_DEBUG("%s: ratio=%u, n_tokens=%u, n_seqs_unq=%u, state_persist_dst=%s, state_write_pos=%s\n",
__func__, ratio, ubatch.n_tokens, ubatch.n_seqs_unq,
dsv4_plan_positions(plan.state_persist_dst_idxs).c_str(),
dsv4_plan_positions(plan.state_write_pos).c_str());
for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) {
const llama_seq_id seq_id = ubatch.seq_id_unq[s];
const uint32_t rollback = seq_id >= 0 && (uint32_t) seq_id < rs_idx.size() ? rs_idx[seq_id] : 0;
LLAMA_LOG_DEBUG("%s: seq %d pos [%d, %d] rollback=%u\n", __func__, seq_id,
ubatch.pos[0], ubatch.pos[ubatch.n_tokens - 1], rollback);
}
}
return plan;
@@ -704,8 +748,17 @@ static std::vector<llama_kv_cache_dsv4_context::comp_plan> dsv4_build_comp_plans
std::vector<llama_kv_cache_dsv4_context::comp_plan> plans;
plans.reserve(ubatches.size());
// the first ubatch touching a seq consumes its rollback restore
std::vector<uint32_t> rs(rs_idx);
for (const llama_ubatch & ubatch : ubatches) {
plans.push_back(dsv4_build_comp_plan(ubatch, ratio, overlap, state_size, kv_size, n_stream, n_rs_seq, rs_idx));
plans.push_back(dsv4_build_comp_plan(ubatch, ratio, overlap, state_size, kv_size, n_stream, n_rs_seq, rs));
for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) {
const llama_seq_id seq_id = ubatch.seq_id_unq[s];
if (seq_id >= 0 && (size_t) seq_id < rs.size()) {
rs[seq_id] = 0;
}
}
}
return plans;
@@ -803,16 +856,15 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_reserve_comp_plan(
return plan;
}
const uint32_t n_seqs = std::max<uint32_t>(1, ubatch.n_seqs);
const uint32_t n_seq_tokens = std::max<uint32_t>(1, ubatch.n_seq_tokens);
const uint64_t n_blocks_u64 = (uint64_t) n_seqs*((n_seq_tokens + ratio - 1)/ratio);
const size_t n_blocks = (size_t) std::max<uint64_t>(1, n_blocks_u64);
GGML_ASSERT((uint64_t) n_blocks == std::max<uint64_t>(1, n_blocks_u64));
// worst case over every seq split: sum of per-seq ceil(tokens/ratio) is at
// most floor(n_tokens/ratio) + n_seqs
const uint32_t n_seqs = std::max<uint32_t>(1, ubatch.n_seqs);
const size_t n_blocks = (size_t) ubatch.n_tokens/ratio + n_seqs;
const uint64_t state_rows = (uint64_t) state_size*n_stream;
const size_t n_persist = (size_t) std::min<uint64_t>(ubatch.n_tokens, state_rows);
const size_t n_restore = n_rs_seq > 0 ? (size_t) state_size*std::max<uint32_t>(1, ubatch.n_seqs_unq) : 0;
const size_t n_snapshot = (size_t) n_rs_seq*state_size*std::max<uint32_t>(1, ubatch.n_seqs_unq);
const size_t n_restore = n_rs_seq > 0 ? (size_t) state_size*n_stream : 0;
const size_t n_snapshot = (size_t) n_rs_seq*state_size*n_stream;
plan.state_pos .resize(ubatch.n_tokens);
plan.state_persist_src_idxs.resize(n_persist);
@@ -1356,7 +1408,9 @@ llama_memory_context_ptr llama_kv_cache_dsv4::init_batch(
if (has_coupled) {
ubatch = balloc.split_seq(n_ubatch);
} else {
ubatch = balloc.split_equal(n_ubatch, raw_per_seq || comp_per_seq, 0);
// [TAG_RECURRENT_ROLLBACK_SPLITS]
// the trailing (1 + n_rs_seq) tokens of each seq must stay in the same ubatch
ubatch = balloc.split_equal(n_ubatch, raw_per_seq || comp_per_seq, n_rs_seq > 0 ? n_rs_seq + 1 : 0);
}
if (ubatch.n_tokens == 0) {
@@ -1433,6 +1487,11 @@ bool llama_kv_cache_dsv4::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1
return false;
}
// pending rollback is single-use: stacked partial removals don't compose
if (rs_idx[seq_id] != 0) {
return false;
}
const bool res = kv_raw->seq_rm(seq_id, p0, p1);
if (res) {
rs_idx[seq_id] = (uint32_t) rollback;
@@ -1594,9 +1653,7 @@ void llama_kv_cache_dsv4::state_read(llama_io_read_i & io, llama_seq_id seq_id,
kv_raw->state_read(io, seq_id, flags);
if (!partial_only) {
kv_csa->clear(true);
kv_hca->clear(true);
kv_lid->clear(true);
clear_compressed(seq_id, true);
dsv4_state_read_k_cache(io, kv_csa.get(), seq_id, flags);
dsv4_state_read_k_cache(io, kv_hca.get(), seq_id, flags);