Skip to content

Commit 7f575c3

Browse files
am17anam17an
authored
DeepseekV4: fix seq_rm (ggml-org#25588)
* DeepseekV4: fix seq_rm * implement proper seq_cp * create actual update context
1 parent 7cbd610 commit 7f575c3

2 files changed

Lines changed: 106 additions & 19 deletions

File tree

‎src/llama-kv-cache-dsv4.cpp‎

Lines changed: 85 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -720,7 +720,7 @@ llama_dsv4_comp_state::llama_dsv4_comp_state(
720720
auto it = ctx_map.find(buft);
721721
if (it == ctx_map.end()) {
722722
ggml_init_params params = {
723-
/*.mem_size =*/ size_t(2u*hparams.n_layer()*ggml_tensor_overhead()),
723+
/*.mem_size =*/ size_t(2u*(1 + n_stream)*hparams.n_layer()*ggml_tensor_overhead()),
724724
/*.mem_buffer =*/ NULL,
725725
/*.no_alloc =*/ true,
726726
};
@@ -767,9 +767,17 @@ llama_dsv4_comp_state::llama_dsv4_comp_state(
767767
ggml_format_name(kv, "dsv4_%s_state_kv_l%d", name, il);
768768
ggml_format_name(score, "dsv4_%s_state_score_l%d", name, il);
769769

770+
std::vector<ggml_tensor *> kv_stream;
771+
std::vector<ggml_tensor *> score_stream;
772+
773+
for (uint32_t s = 0; s < n_stream; ++s) {
774+
kv_stream.push_back(ggml_view_2d(ctx, kv, n_embd_state, state_size, kv->nb[1], s*kv->nb[2]));
775+
score_stream.push_back(ggml_view_2d(ctx, score, n_embd_state, state_size, score->nb[1], s*score->nb[2]));
776+
}
777+
770778
map_layer_ids[il] = layers.size();
771779

772-
layers.push_back({ il, kv, score });
780+
layers.push_back({ il, kv, score, std::move(kv_stream), std::move(score_stream) });
773781
}
774782

775783
for (auto & [buft, ctx] : ctx_map) {
@@ -809,6 +817,30 @@ void llama_dsv4_comp_state::clear(llama_seq_id seq_id, bool data) {
809817
}
810818
}
811819

820+
void llama_dsv4_comp_state::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst) {
821+
GGML_ASSERT(seq_id_src >= 0 && (uint32_t) seq_id_src < n_stream);
822+
GGML_ASSERT(seq_id_dst >= 0 && (uint32_t) seq_id_dst < n_stream);
823+
824+
if (seq_id_src == seq_id_dst) {
825+
return;
826+
}
827+
828+
sc_info.ssrc.push_back((uint32_t) seq_id_src);
829+
sc_info.sdst.push_back((uint32_t) seq_id_dst);
830+
}
831+
832+
void llama_dsv4_comp_state::apply_copies(const stream_copy_info & sc_info) const {
833+
for (size_t i = 0; i < sc_info.ssrc.size(); ++i) {
834+
const uint32_t ssrc = sc_info.ssrc[i];
835+
const uint32_t sdst = sc_info.sdst[i];
836+
837+
for (const auto & layer : layers) {
838+
ggml_backend_tensor_copy(layer.kv_stream[ssrc], layer.kv_stream[sdst]);
839+
ggml_backend_tensor_copy(layer.score_stream[ssrc], layer.score_stream[sdst]);
840+
}
841+
}
842+
}
843+
812844
uint32_t llama_dsv4_comp_state::get_ratio() const {
813845
return ratio;
814846
}
@@ -1154,7 +1186,13 @@ llama_memory_context_ptr llama_kv_cache_dsv4::init_full() {
11541186
}
11551187

11561188
llama_memory_context_ptr llama_kv_cache_dsv4::init_update(llama_context * lctx, bool optimize) {
1157-
return std::make_unique<llama_kv_cache_dsv4_context>(this, lctx, optimize);
1189+
return std::make_unique<llama_kv_cache_dsv4_context>(
1190+
this,
1191+
lctx,
1192+
optimize,
1193+
std::move(csa_state->sc_info),
1194+
std::move(hca_state->sc_info),
1195+
std::move(lid_state->sc_info));
11581196
}
11591197

11601198
bool llama_kv_cache_dsv4::get_can_shift() const {
@@ -1174,14 +1212,19 @@ bool llama_kv_cache_dsv4::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1
11741212
}
11751213

11761214
if (p0 > 0) {
1177-
// DSV4 compressed cache rows are derived from running compressor state,
1178-
// so arbitrary rollback is not reconstructible from the raw cache alone.
1179-
// Allow the common prompt-cache cleanup no-op: remove [end, infinity).
1180-
if (seq_id >= 0 && p0 > kv_raw->seq_pos_max(seq_id)) {
1181-
return true;
1215+
if (seq_id < 0 || (uint32_t) seq_id >= n_seq_max ||
1216+
p0 <= kv_raw->seq_pos_max(seq_id)) {
1217+
return false;
11821218
}
11831219

1184-
return false;
1220+
bool res = true;
1221+
1222+
res = res & kv_raw->seq_rm(seq_id, p0, -1);
1223+
res = res & kv_csa->seq_rm(seq_id, p0/DSV4_CSA_RATIO, -1);
1224+
res = res & kv_hca->seq_rm(seq_id, p0/DSV4_HCA_RATIO, -1);
1225+
res = res & kv_lid->seq_rm(seq_id, p0/DSV4_CSA_RATIO, -1);
1226+
1227+
return res;
11851228
}
11861229

11871230
const bool res = kv_raw->seq_rm(seq_id, p0, p1);
@@ -1194,7 +1237,16 @@ bool llama_kv_cache_dsv4::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1
11941237
}
11951238

11961239
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) {
1240+
GGML_ASSERT(p0 <= 0 && p1 < 0 && "DSV4 only supports full sequence copies");
1241+
11971242
kv_raw->seq_cp(seq_id_src, seq_id_dst, p0, p1);
1243+
kv_csa->seq_cp(seq_id_src, seq_id_dst, -1, -1);
1244+
kv_hca->seq_cp(seq_id_src, seq_id_dst, -1, -1);
1245+
kv_lid->seq_cp(seq_id_src, seq_id_dst, -1, -1);
1246+
1247+
csa_state->seq_cp(seq_id_src, seq_id_dst);
1248+
hca_state->seq_cp(seq_id_src, seq_id_dst);
1249+
lid_state->seq_cp(seq_id_src, seq_id_dst);
11981250
}
11991251

12001252
void llama_kv_cache_dsv4::seq_keep(llama_seq_id seq_id) {
@@ -1639,20 +1691,26 @@ llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context(
16391691
llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context(
16401692
llama_kv_cache_dsv4 * kv,
16411693
llama_context * lctx,
1642-
bool optimize) :
1694+
bool optimize,
1695+
stream_copy_info sc_info_csa,
1696+
stream_copy_info sc_info_hca,
1697+
stream_copy_info sc_info_lid) :
16431698
ctx_raw(std::make_unique<llama_kv_cache_dsv4_raw_context>(kv->get_raw(), lctx, optimize)),
16441699
ctx_csa_mem(kv->get_csa()->init_update(lctx, optimize)),
16451700
ctx_hca_mem(kv->get_hca()->init_update(lctx, optimize)),
16461701
ctx_lid_mem(kv->get_lid()->init_update(lctx, optimize)),
1647-
ctx_csa(std::make_unique<llama_kv_cache_dsv4_comp_context>(kv->get_csa())),
1648-
ctx_hca(std::make_unique<llama_kv_cache_dsv4_comp_context>(kv->get_hca())),
1649-
ctx_lid(std::make_unique<llama_kv_cache_dsv4_comp_context>(kv->get_lid())),
16501702
csa_state(kv->get_csa_state()),
16511703
hca_state(kv->get_hca_state()),
16521704
lid_state(kv->get_lid_state()),
1705+
sc_info_csa(std::move(sc_info_csa)),
1706+
sc_info_hca(std::move(sc_info_hca)),
1707+
sc_info_lid(std::move(sc_info_lid)),
16531708
status(llama_memory_status_combine(
1654-
llama_memory_status_combine(ctx_raw->get_status(), ctx_csa_mem->get_status()),
1655-
llama_memory_status_combine(ctx_hca_mem->get_status(), ctx_lid_mem->get_status()))) {
1709+
llama_memory_status_combine(
1710+
llama_memory_status_combine(ctx_raw->get_status(), ctx_csa_mem->get_status()),
1711+
llama_memory_status_combine(ctx_hca_mem->get_status(), ctx_lid_mem->get_status())),
1712+
this->sc_info_csa.empty() && this->sc_info_hca.empty() && this->sc_info_lid.empty() ?
1713+
LLAMA_MEMORY_STATUS_NO_UPDATE : LLAMA_MEMORY_STATUS_SUCCESS)) {
16561714
}
16571715

16581716
llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context(
@@ -1720,6 +1778,18 @@ bool llama_kv_cache_dsv4_context::apply() {
17201778

17211779
res = res & ctx_raw->apply();
17221780

1781+
if (ctx_csa_mem) {
1782+
res = res & ctx_csa_mem->apply();
1783+
res = res & ctx_hca_mem->apply();
1784+
res = res & ctx_lid_mem->apply();
1785+
}
1786+
1787+
if (ubatches.empty()) {
1788+
csa_state->apply_copies(sc_info_csa);
1789+
hca_state->apply_copies(sc_info_hca);
1790+
lid_state->apply_copies(sc_info_lid);
1791+
}
1792+
17231793
return res;
17241794
}
17251795

‎src/llama-kv-cache-dsv4.h‎

Lines changed: 21 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,10 @@
1010

1111
class llama_dsv4_comp_state {
1212
public:
13+
using stream_copy_info = llama_kv_cache::stream_copy_info;
14+
15+
stream_copy_info sc_info;
16+
1317
llama_dsv4_comp_state(
1418
const llama_model & model,
1519
bool offload,
@@ -22,6 +26,8 @@ class llama_dsv4_comp_state {
2226
const llama_memory_i::layer_filter_cb & filter);
2327

2428
void clear(llama_seq_id seq_id, bool data);
29+
void seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst);
30+
void apply_copies(const stream_copy_info & sc_info) const;
2531

2632
uint32_t get_ratio() const;
2733
uint32_t get_state_size() const;
@@ -44,6 +50,9 @@ class llama_dsv4_comp_state {
4450

4551
ggml_tensor * kv;
4652
ggml_tensor * score;
53+
54+
std::vector<ggml_tensor *> kv_stream;
55+
std::vector<ggml_tensor *> score_stream;
4756
};
4857

4958
const uint32_t ratio;
@@ -245,6 +254,7 @@ class llama_kv_cache_dsv4_comp_context {
245254
class llama_kv_cache_dsv4_context : public llama_memory_context_i {
246255
public:
247256
using slot_info_vec_t = llama_kv_cache::slot_info_vec_t;
257+
using stream_copy_info = llama_kv_cache::stream_copy_info;
248258

249259
struct comp_plan {
250260
// Per-ubatch recipe for updating compressor state, committing completed
@@ -291,7 +301,10 @@ class llama_kv_cache_dsv4_context : public llama_memory_context_i {
291301
llama_kv_cache_dsv4_context(
292302
llama_kv_cache_dsv4 * kv,
293303
llama_context * lctx,
294-
bool optimize);
304+
bool optimize,
305+
stream_copy_info sc_info_csa,
306+
stream_copy_info sc_info_hca,
307+
stream_copy_info sc_info_lid);
295308

296309
llama_kv_cache_dsv4_context(
297310
llama_kv_cache_dsv4 * kv,
@@ -351,9 +364,13 @@ class llama_kv_cache_dsv4_context : public llama_memory_context_i {
351364
const std::unique_ptr<llama_kv_cache_dsv4_comp_context> ctx_hca;
352365
const std::unique_ptr<llama_kv_cache_dsv4_comp_context> ctx_lid;
353366

354-
const llama_dsv4_comp_state * csa_state = nullptr;
355-
const llama_dsv4_comp_state * hca_state = nullptr;
356-
const llama_dsv4_comp_state * lid_state = nullptr;
367+
llama_dsv4_comp_state * csa_state = nullptr;
368+
llama_dsv4_comp_state * hca_state = nullptr;
369+
llama_dsv4_comp_state * lid_state = nullptr;
370+
371+
stream_copy_info sc_info_csa;
372+
stream_copy_info sc_info_hca;
373+
stream_copy_info sc_info_lid;
357374

358375
bool reserve_plans = false;
359376
mutable comp_plan reserve_plan_csa;

0 commit comments

Comments
 (0)