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

DeepseekV4 MTP + DSpark (#25784)

Browse Source
This commit is contained in:
Aman Gupta
2026-08-02 20:55:34 +08:00
committed by GitHub
parent f5919bf458
commit 596a5795bd
14 changed files with 1534 additions and 225 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/llama-kv-cache-dsv4.cpp
+282 -65
View File
@@ -252,7 +252,8 @@ static void dsv4_state_write_tensor_streams(
uint32_t tensor_rows,
uint32_t n_rows,
uint32_t s0,
uint32_t ns) {
uint32_t ns,
const std::vector<uint32_t> * stream_ids = nullptr) {
const int32_t type_i = (int32_t) tensor->type;
const uint64_t ne0 = tensor->ne[0];
const uint64_t rows = n_rows;
@@ -273,8 +274,16 @@ static void dsv4_state_write_tensor_streams(
return;
}
if (stream_ids && stream_ids->size() != ns) {
throw std::runtime_error("DSV4 state tensor stream map size mismatch");
}
for (uint32_t s = 0; s < ns; ++s) {
const size_t offset = (size_t) (s0 + s)*stream_stride;
const uint32_t stream = stream_ids ? (*stream_ids)[s] : s0 + s;
if ((int64_t) stream >= tensor->ne[2]) {
throw std::runtime_error("DSV4 state tensor stream out of range");
}
const size_t offset = (size_t) stream*stream_stride;
io.write_tensor(tensor, offset, size);
}
}
@@ -421,7 +430,9 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan(
bool overlap,
uint32_t state_size,
uint32_t kv_size,
uint32_t n_stream) {
uint32_t n_stream,
uint32_t n_rs_seq,
const std::vector<uint32_t> & rs_idx) {
llama_kv_cache_dsv4_context::comp_plan plan;
plan.n_visible.resize(ubatch.n_tokens);
plan.n_stream = dsv4_comp_graph_n_stream(ubatch, n_stream);
@@ -451,6 +462,7 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan(
std::vector<int32_t> overlap_cur_reads;
std::map<std::pair<llama_seq_id, llama_pos>, int64_t> curr_token_idx_map;
std::map<llama_seq_id, uint32_t> state_write_counts;
for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
for (int32_t s = 0; s < ubatch.n_seq_id[i]; ++s) {
@@ -513,6 +525,7 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan(
plan.state_write_idxs.push_back(cache_off + pos/ratio);
plan.state_write_pos.push_back((int32_t) source_start);
++state_write_counts[seq_id];
if (overlap) {
const llama_pos prev_start = source_start - ratio;
@@ -531,33 +544,57 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan(
}
}
if (ratio == DSV4_CSA_RATIO && plan.state_write_idxs.empty() && !plan.state_pos.empty()) {
// Non-boundary CSA steps still need a write op so their graph matches
// boundary steps. Use a padded scratch row that is masked from attention.
if (ratio == DSV4_CSA_RATIO && !plan.state_pos.empty()) {
assert(kv_size > 0);
uint32_t i = 0;
while (i < ubatch.n_tokens && ubatch.pos[i] < 0) {
++i;
}
assert(i < ubatch.n_tokens);
// Pad each stream to the reserve plan's block count.
const auto append_dummy_block = [&](llama_seq_id seq_id, uint32_t i) {
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]);
const llama_pos pos = ubatch.pos[i];
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, pos);
plan.state_write_idxs.push_back(cache_off + kv_size - 1);
plan.state_write_pos .push_back(0);
plan.state_write_idxs.push_back(cache_off + kv_size - 1);
plan.state_write_pos .push_back(0);
if (overlap) {
for (uint32_t j = 0; j < ratio; ++j) {
overlap_prev_reads.push_back(source_idx);
overlap_cur_reads .push_back(source_idx);
}
} else {
for (uint32_t j = 0; j < ratio; ++j) {
plan.state_read_idxs.push_back(source_idx);
}
}
};
if (overlap) {
for (uint32_t j = 0; j < ratio; ++j) {
overlap_prev_reads.push_back(source_idx);
overlap_cur_reads .push_back(source_idx);
if (dsv4_ubatch_has_coupled(ubatch)) {
if (plan.state_write_idxs.empty()) {
uint32_t i = 0;
while (i < ubatch.n_tokens && ubatch.pos[i] < 0) {
++i;
}
assert(i < ubatch.n_tokens);
append_dummy_block(ubatch.seq_id[i][0], i);
}
} else {
for (uint32_t j = 0; j < ratio; ++j) {
plan.state_read_idxs.push_back(source_idx);
const uint32_t n_blocks = (std::max<uint32_t>(1, ubatch.n_seq_tokens) + ratio - 1)/ratio;
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 n_writes = state_write_counts[seq_id];
if (n_writes >= n_blocks) {
continue;
}
if (n_writes + 1 != n_blocks) {
throw std::runtime_error("DSV4 CSA sequence positions are not contiguous");
}
uint32_t i = 0;
while (i < ubatch.n_tokens && (ubatch.pos[i] < 0 || !dsv4_token_has_seq(ubatch, i, seq_id))) {
++i;
}
assert(i < ubatch.n_tokens);
append_dummy_block(seq_id, i);
}
}
}
@@ -583,6 +620,63 @@ 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;
}
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;
// 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) {
plan.state_restore_src_idxs.push_back((int32_t) (src_plane + stream_off + r));
plan.state_restore_dst_idxs.push_back((int32_t) (stream_off + r));
}
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 (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;
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);
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);
}
plan.state_snapshot_src_idxs.push_back(src);
plan.state_snapshot_dst_idxs.push_back((int32_t) (dst_plane + stream_off + r));
}
}
}
}
static const bool debug = []() {
const char * env = getenv("LLAMA_DSV4_COMPRESS_DEBUG");
return env && atoi(env) > 0;
@@ -604,12 +698,14 @@ static std::vector<llama_kv_cache_dsv4_context::comp_plan> dsv4_build_comp_plans
bool overlap,
uint32_t state_size,
uint32_t kv_size,
uint32_t n_stream) {
uint32_t n_stream,
uint32_t n_rs_seq,
const std::vector<uint32_t> & rs_idx) {
std::vector<llama_kv_cache_dsv4_context::comp_plan> plans;
plans.reserve(ubatches.size());
for (const llama_ubatch & ubatch : ubatches) {
plans.push_back(dsv4_build_comp_plan(ubatch, ratio, overlap, state_size, kv_size, n_stream));
plans.push_back(dsv4_build_comp_plan(ubatch, ratio, overlap, state_size, kv_size, n_stream, n_rs_seq, rs_idx));
}
return plans;
@@ -696,7 +792,8 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_reserve_comp_plan(
bool overlap,
uint32_t state_size,
uint32_t kv_size,
uint32_t n_stream) {
uint32_t n_stream,
uint32_t n_rs_seq) {
llama_kv_cache_dsv4_context::comp_plan plan;
plan.n_visible.resize(ubatch.n_tokens);
plan.n_stream = dsv4_comp_graph_n_stream(ubatch, n_stream);
@@ -714,10 +811,16 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_reserve_comp_plan(
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);
plan.state_pos .resize(ubatch.n_tokens);
plan.state_persist_src_idxs.resize(n_persist);
plan.state_persist_dst_idxs.resize(n_persist);
plan.state_restore_src_idxs.resize(n_restore);
plan.state_restore_dst_idxs.resize(n_restore);
plan.state_snapshot_src_idxs.resize(n_snapshot);
plan.state_snapshot_dst_idxs.resize(n_snapshot);
plan.state_read_idxs .resize((overlap ? 2u : 1u)*ratio*n_blocks);
plan.state_write_idxs.resize(n_blocks);
plan.state_write_pos .resize(n_blocks);
@@ -743,12 +846,14 @@ llama_dsv4_comp_state::llama_dsv4_comp_state(
uint32_t ratio,
uint32_t state_size,
uint32_t n_embd_state,
uint32_t n_rs_seq,
const char * name,
const llama_memory_i::layer_filter_cb & filter) :
ratio(ratio),
state_size(state_size),
n_embd_state(n_embd_state),
n_stream(unified ? 1 : n_seq_max) {
n_stream(unified ? 1 : n_seq_max),
n_rs_seq(n_rs_seq) {
const llama_hparams & hparams = model.hparams;
struct ggml_backend_buft_comparator {
@@ -804,8 +909,9 @@ llama_dsv4_comp_state::llama_dsv4_comp_state(
throw std::runtime_error("failed to create ggml context for DSV4 compressor state");
}
ggml_tensor * kv = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_state, state_size, n_stream);
ggml_tensor * score = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_state, state_size, n_stream);
const uint32_t n_planes = n_stream*(1 + n_rs_seq);
ggml_tensor * kv = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_state, state_size, n_planes);
ggml_tensor * score = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_state, state_size, n_planes);
ggml_format_name(kv, "dsv4_%s_state_kv_l%d", name, il);
ggml_format_name(score, "dsv4_%s_state_score_l%d", name, il);
@@ -837,8 +943,8 @@ llama_dsv4_comp_state::llama_dsv4_comp_state(
ctxs_bufs.emplace_back(std::move(ctx), buf);
}
LLAMA_LOG_INFO("%s: %s ratio = %u, state = %u x %u, streams = %u, layers = %zu, size = %7.2f MiB\n",
__func__, name, ratio, state_size, n_embd_state, n_stream, layers.size(), total_size()/1024.0/1024.0);
LLAMA_LOG_INFO("%s: %s ratio = %u, state = %u x %u, streams = %u, rs_seq = %u, layers = %zu, size = %7.2f MiB\n",
__func__, name, ratio, state_size, n_embd_state, n_stream, n_rs_seq, layers.size(), total_size()/1024.0/1024.0);
}
void llama_dsv4_comp_state::clear(llama_seq_id seq_id, bool data) {
@@ -848,9 +954,13 @@ void llama_dsv4_comp_state::clear(llama_seq_id seq_id, bool data) {
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);
for (uint32_t d = 0; d <= n_rs_seq; ++d) {
const uint32_t stream = d*n_stream + (uint32_t) seq_id;
dsv4_clear_tensor_stream(layer.kv, stream);
dsv4_clear_tensor_stream(layer.score, stream);
}
}
return;
}
@@ -868,6 +978,8 @@ void llama_dsv4_comp_state::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_
return;
}
clear(seq_id_dst, true);
sc_info.ssrc.push_back((uint32_t) seq_id_src);
sc_info.sdst.push_back((uint32_t) seq_id_dst);
}
@@ -896,6 +1008,14 @@ uint32_t llama_dsv4_comp_state::get_n_stream() const {
return n_stream;
}
uint32_t llama_dsv4_comp_state::get_n_rs_seq() const {
return n_rs_seq;
}
uint32_t llama_dsv4_comp_state::get_n_rows() const {
return state_size*n_stream;
}
std::map<ggml_backend_buffer_type_t, size_t> llama_dsv4_comp_state::memory_breakdown() const {
std::map<ggml_backend_buffer_type_t, size_t> ret;
for (const auto & [_, buf] : ctxs_bufs) {
@@ -905,13 +1025,26 @@ std::map<ggml_backend_buffer_type_t, size_t> llama_dsv4_comp_state::memory_break
return ret;
}
void llama_dsv4_comp_state::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const {
void llama_dsv4_comp_state::state_write(
llama_io_write_i & io,
llama_seq_id seq_id,
llama_state_seq_flags flags,
const std::vector<uint32_t> & rs_idx) const {
GGML_UNUSED(flags);
uint32_t s0;
uint32_t ns;
dsv4_state_src_stream_range(n_stream, seq_id, s0, ns);
std::vector<uint32_t> stream_ids(ns);
for (uint32_t s = 0; s < ns; ++s) {
const uint32_t seq = seq_id >= 0 ? (uint32_t) seq_id : s0 + s;
if (seq >= rs_idx.size() || rs_idx[seq] > n_rs_seq) {
throw std::runtime_error("DSV4 recurrent state rollback index out of range");
}
stream_ids[s] = rs_idx[seq]*n_stream + s0 + s;
}
const uint32_t version = DSV4_COMP_STATE_VER;
const uint32_t n_layer = layers.size();
@@ -925,8 +1058,8 @@ void llama_dsv4_comp_state::state_write(llama_io_write_i & io, llama_seq_id seq_
for (const auto & layer : layers) {
io.write(&layer.il, sizeof(layer.il));
dsv4_state_write_tensor_streams(io, layer.kv, state_size, state_size, s0, ns);
dsv4_state_write_tensor_streams(io, layer.score, state_size, state_size, s0, ns);
dsv4_state_write_tensor_streams(io, layer.kv, state_size, state_size, s0, ns, &stream_ids);
dsv4_state_write_tensor_streams(io, layer.score, state_size, state_size, s0, ns, &stream_ids);
}
}
@@ -972,28 +1105,40 @@ void llama_dsv4_comp_state::state_read(llama_io_read_i & io, llama_seq_id seq_id
}
}
ggml_tensor * llama_dsv4_comp_state::get_kv(ggml_context * ctx, int32_t il) const {
ggml_tensor * llama_dsv4_comp_state::get_kv_all(ggml_context * ctx, int32_t il) const {
const int32_t ids = map_layer_ids.at(il);
ggml_tensor * state = layers[ids].kv;
return ggml_reshape_2d(ctx, state, state->ne[0], state->ne[1]*state->ne[2]);
return ggml_view_2d(ctx, state, state->ne[0], get_n_rows()*(1 + n_rs_seq), state->nb[1], 0);
}
ggml_tensor * llama_dsv4_comp_state::get_score_all(ggml_context * ctx, int32_t il) const {
const int32_t ids = map_layer_ids.at(il);
ggml_tensor * state = layers[ids].score;
return ggml_view_2d(ctx, state, state->ne[0], get_n_rows()*(1 + n_rs_seq), state->nb[1], 0);
}
ggml_tensor * llama_dsv4_comp_state::get_kv(ggml_context * ctx, int32_t il) const {
ggml_tensor * state = get_kv_all(ctx, il);
const size_t row_size = ggml_row_size(state->type, state->ne[0]);
return ggml_view_2d(ctx, state, state->ne[0], get_n_rows(), state->nb[1], 0*row_size);
}
ggml_tensor * llama_dsv4_comp_state::get_score(ggml_context * ctx, int32_t il) const {
const int32_t ids = map_layer_ids.at(il);
ggml_tensor * state = get_score_all(ctx, il);
const size_t row_size = ggml_row_size(state->type, state->ne[0]);
ggml_tensor * state = layers[ids].score;
return ggml_reshape_2d(ctx, state, state->ne[0], state->ne[1]*state->ne[2]);
return ggml_view_2d(ctx, state, state->ne[0], get_n_rows(), state->nb[1], 0*row_size);
}
ggml_tensor * llama_dsv4_comp_state::cpy_kv(ggml_context * ctx, ggml_tensor * cur, ggml_tensor * idxs, int32_t il) const {
return ggml_set_rows(ctx, get_kv(ctx, il), cur, idxs);
return ggml_set_rows(ctx, get_kv_all(ctx, il), cur, idxs);
}
ggml_tensor * llama_dsv4_comp_state::cpy_score(ggml_context * ctx, ggml_tensor * cur, ggml_tensor * idxs, int32_t il) const {
return ggml_set_rows(ctx, get_score(ctx, il), cur, idxs);
return ggml_set_rows(ctx, get_score_all(ctx, il), cur, idxs);
}
size_t llama_dsv4_comp_state::total_size() const {
@@ -1022,13 +1167,16 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4(
uint32_t n_seq_max,
uint32_t n_ubatch,
uint32_t n_pad,
uint32_t n_rs_seq,
const layer_filter_cb & filter,
const layer_reuse_cb & reuse) :
hparams_raw(model.hparams),
hparams_csa(model.hparams),
hparams_hca(model.hparams),
hparams_lid(model.hparams),
n_seq_max(n_seq_max) {
n_seq_max(n_seq_max),
n_rs_seq(n_rs_seq),
rs_idx(n_seq_max, 0) {
const layer_filter_cb filter_raw = [&](int32_t il) {
if (filter && !filter(il)) {
@@ -1043,6 +1191,11 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4(
// Keep DSV4 KV/state streams per sequence even when public KV mode is unified.
const bool unified_raw = false;
hparams_raw.n_layer_nextn = 0;
hparams_csa.n_layer_nextn = 0;
hparams_hca.n_layer_nextn = 0;
hparams_lid.n_layer_nextn = 0;
LLAMA_LOG_INFO("%s: creating DSV4 raw KV cache\n", __func__);
dsv4_make_k_only(hparams_raw);
@@ -1109,19 +1262,19 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4(
csa_state = std::make_unique<llama_dsv4_comp_state>(
model, offload, unified_compressed, n_seq_max, DSV4_CSA_RATIO, 2*DSV4_CSA_RATIO,
2*model.hparams.n_embd_head_k(), "csa", filter_csa);
2*model.hparams.n_embd_head_k(), n_rs_seq, "csa", filter_csa);
LLAMA_LOG_INFO("%s: creating DSV4 HCA compressor state\n", __func__);
hca_state = std::make_unique<llama_dsv4_comp_state>(
model, offload, unified_compressed, n_seq_max, DSV4_HCA_RATIO, DSV4_HCA_RATIO,
model.hparams.n_embd_head_k(), "hca", filter_hca);
model.hparams.n_embd_head_k(), n_rs_seq, "hca", filter_hca);
LLAMA_LOG_INFO("%s: creating DSV4 lightning-indexer compressor state\n", __func__);
lid_state = std::make_unique<llama_dsv4_comp_state>(
model, offload, unified_compressed, n_seq_max, DSV4_CSA_RATIO, 2*DSV4_CSA_RATIO,
2*model.hparams.indexer_head_size, "lid", filter_csa);
2*model.hparams.indexer_head_size, n_rs_seq, "lid", filter_csa);
// DSV4 attention reads compressed-K / compressor-state rows that the current
// graph does not necessarily overwrite; uninitialized buffer contents would
@@ -1255,17 +1408,35 @@ bool llama_kv_cache_dsv4::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1
}
if (p0 > 0) {
if (seq_id < 0 || (uint32_t) seq_id >= n_seq_max ||
p0 <= kv_raw->seq_pos_max(seq_id)) {
if (seq_id < 0 || (uint32_t) seq_id >= n_seq_max) {
return false;
}
bool res = true;
const llama_pos pos_max = kv_raw->seq_pos_max(seq_id);
if (p0 > pos_max) {
bool res = true;
res = res & kv_raw->seq_rm(seq_id, p0, -1);
res = res & kv_csa->seq_rm(seq_id, p0/DSV4_CSA_RATIO, -1);
res = res & kv_hca->seq_rm(seq_id, p0/DSV4_HCA_RATIO, -1);
res = res & kv_lid->seq_rm(seq_id, p0/DSV4_CSA_RATIO, -1);
res = res & kv_raw->seq_rm(seq_id, p0, -1);
res = res & kv_csa->seq_rm(seq_id, p0/DSV4_CSA_RATIO, -1);
res = res & kv_hca->seq_rm(seq_id, p0/DSV4_HCA_RATIO, -1);
res = res & kv_lid->seq_rm(seq_id, p0/DSV4_CSA_RATIO, -1);
return res;
}
if (n_rs_seq == 0) {
return false;
}
const llama_pos rollback = pos_max - (p0 - 1);
if (rollback < 1 || rollback > (llama_pos) n_rs_seq) {
return false;
}
const bool res = kv_raw->seq_rm(seq_id, p0, p1);
if (res) {
rs_idx[seq_id] = (uint32_t) rollback;
}
return res;
}
@@ -1290,6 +1461,10 @@ void llama_kv_cache_dsv4::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_ds
csa_state->seq_cp(seq_id_src, seq_id_dst);
hca_state->seq_cp(seq_id_src, seq_id_dst);
lid_state->seq_cp(seq_id_src, seq_id_dst);
if (seq_id_src != seq_id_dst) {
rs_idx[seq_id_dst] = 0;
}
}
void llama_kv_cache_dsv4::seq_keep(llama_seq_id seq_id) {
@@ -1386,9 +1561,9 @@ void llama_kv_cache_dsv4::state_write(llama_io_write_i & io, llama_seq_id seq_id
dsv4_state_write_k_cache(io, kv_lid.get(), seq_id, flags, n_rows_lid);
}
csa_state->state_write(io, seq_id, flags);
hca_state->state_write(io, seq_id, flags);
lid_state->state_write(io, seq_id, flags);
csa_state->state_write(io, seq_id, flags, rs_idx);
hca_state->state_write(io, seq_id, flags, rs_idx);
lid_state->state_write(io, seq_id, flags, rs_idx);
}
void llama_kv_cache_dsv4::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) {
@@ -1432,6 +1607,12 @@ void llama_kv_cache_dsv4::state_read(llama_io_read_i & io, llama_seq_id seq_id,
hca_state->state_read(io, seq_id, flags);
lid_state->state_read(io, seq_id, flags);
if (seq_id >= 0) {
GGML_ASSERT((uint32_t) seq_id < n_seq_max);
rs_idx[seq_id] = 0;
} else {
std::fill(rs_idx.begin(), rs_idx.end(), 0);
}
}
llama_kv_cache_iswa * llama_kv_cache_dsv4::get_raw() const {
@@ -1462,6 +1643,31 @@ llama_dsv4_comp_state * llama_kv_cache_dsv4::get_lid_state() const {
return lid_state.get();
}
uint32_t llama_kv_cache_dsv4::get_n_rs_seq() const {
return n_rs_seq;
}
const std::vector<uint32_t> & llama_kv_cache_dsv4::get_rs_idx() const {
return rs_idx;
}
void llama_kv_cache_dsv4::reset_rs_idx_for_ubatches(const std::vector<llama_ubatch> & ubatches) {
if (n_rs_seq == 0) {
return;
}
for (const llama_ubatch & ubatch : ubatches) {
for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
for (int32_t s = 0; s < ubatch.n_seq_id[i]; ++s) {
const llama_seq_id seq_id = ubatch.seq_id[i][s];
if (seq_id >= 0 && (uint32_t) seq_id < n_seq_max) {
rs_idx[seq_id] = 0;
}
}
}
}
}
void llama_kv_cache_dsv4::clear_compressed(llama_seq_id seq_id, bool data) {
if (seq_id < 0) {
kv_csa->clear(data);
@@ -1488,6 +1694,12 @@ void llama_kv_cache_dsv4::clear_compressed(llama_seq_id seq_id, bool data) {
csa_state->clear(seq_id, data);
hca_state->clear(seq_id, data);
lid_state->clear(seq_id, data);
if (seq_id >= 0) {
rs_idx[seq_id] = 0;
} else {
std::fill(rs_idx.begin(), rs_idx.end(), 0);
}
}
//
@@ -1779,10 +1991,14 @@ llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context(
std::vector<llama_ubatch> ubatches_raw) :
ubatches(std::move(ubatches)),
plans_csa(dsv4_build_comp_plans(this->ubatches, DSV4_CSA_RATIO, true,
kv->get_csa_state()->get_state_size(), kv->get_csa()->get_size(), kv->get_csa_state()->get_n_stream())),
kv->get_csa_state()->get_state_size(), kv->get_csa()->get_size(), kv->get_csa_state()->get_n_stream(),
kv->get_n_rs_seq(), kv->get_rs_idx())),
plans_hca(dsv4_build_comp_plans(this->ubatches, DSV4_HCA_RATIO, false,
kv->get_hca_state()->get_state_size(), kv->get_hca()->get_size(), kv->get_hca_state()->get_n_stream())),
plans_lid(plans_csa),
kv->get_hca_state()->get_state_size(), kv->get_hca()->get_size(), kv->get_hca_state()->get_n_stream(),
kv->get_n_rs_seq(), kv->get_rs_idx())),
plans_lid(dsv4_build_comp_plans(this->ubatches, DSV4_CSA_RATIO, true,
kv->get_lid_state()->get_state_size(), kv->get_lid()->get_size(), kv->get_lid_state()->get_n_stream(),
kv->get_n_rs_seq(), kv->get_rs_idx())),
ctx_raw(std::make_unique<llama_kv_cache_dsv4_raw_context>(
kv->get_raw(),
std::move(sinfos_raw_base_write),
@@ -1809,6 +2025,7 @@ llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context(
hca_state(kv->get_hca_state()),
lid_state(kv->get_lid_state()),
status(ctx_raw->get_status()) {
kv->reset_rs_idx_for_ubatches(this->ubatches);
}
llama_kv_cache_dsv4_context::~llama_kv_cache_dsv4_context() = default;
@@ -1944,7 +2161,7 @@ const llama_kv_cache_dsv4_context::comp_plan & llama_kv_cache_dsv4_context::get_
reserve_plan_csa = dsv4_build_reserve_comp_plan(
ubatch, DSV4_CSA_RATIO, true,
csa_state->get_state_size(), get_csa()->get_n_kv(), csa_state->get_n_stream());
csa_state->get_state_size(), get_csa()->get_n_kv(), csa_state->get_n_stream(), csa_state->get_n_rs_seq());
return reserve_plan_csa;
}
@@ -1958,7 +2175,7 @@ const llama_kv_cache_dsv4_context::comp_plan & llama_kv_cache_dsv4_context::get_
reserve_plan_hca = dsv4_build_reserve_comp_plan(
ubatch, DSV4_HCA_RATIO, false,
hca_state->get_state_size(), get_hca()->get_n_kv(), hca_state->get_n_stream());
hca_state->get_state_size(), get_hca()->get_n_kv(), hca_state->get_n_stream(), hca_state->get_n_rs_seq());
return reserve_plan_hca;
}
@@ -1972,7 +2189,7 @@ const llama_kv_cache_dsv4_context::comp_plan & llama_kv_cache_dsv4_context::get_
reserve_plan_lid = dsv4_build_reserve_comp_plan(
ubatch, DSV4_CSA_RATIO, true,
lid_state->get_state_size(), get_lid()->get_n_kv(), lid_state->get_n_stream());
lid_state->get_state_size(), get_lid()->get_n_kv(), lid_state->get_n_stream(), lid_state->get_n_rs_seq());
return reserve_plan_lid;
}