Skip to content

Commit 91d2fc3

Browse files
am17anam17an
authored
llama_dsv4: write only used rows in state (ggml-org#25325)
* llama_dsv4: write only used rows in state * add TODO about conflating token pos with kv rows
1 parent 4ee6a9a commit 91d2fc3

1 file changed

Lines changed: 79 additions & 22 deletions

File tree

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

Lines changed: 79 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ static constexpr uint32_t DSV4_STATE_MAGIC = 0x34565344; // DSV4
2222
static constexpr uint32_t DSV4_STATE_VERSION = 1;
2323
static constexpr uint32_t DSV4_STATE_MODE_FULL = 0;
2424
static constexpr uint32_t DSV4_STATE_MODE_PARTIAL = 1;
25-
static constexpr uint32_t DSV4_K_CACHE_STATE_VER = 1;
25+
static constexpr uint32_t DSV4_K_CACHE_STATE_VER = 2;
2626
static constexpr uint32_t DSV4_COMP_STATE_VER = 1;
2727

2828
static uint32_t dsv4_comp_size(uint32_t kv_size, uint32_t ratio) {
@@ -38,6 +38,16 @@ static void dsv4_clear_tensor_stream(ggml_tensor * tensor, uint32_t stream) {
3838
ggml_backend_tensor_memset(tensor, 0, stream*stream_size, stream_size);
3939
}
4040

41+
static uint32_t dsv4_state_n_used_k_rows(llama_pos pos_max, uint32_t ratio, uint32_t kv_size) {
42+
if (pos_max < 0) {
43+
return 0;
44+
}
45+
46+
const uint64_t n_rows = ((uint64_t) pos_max + 1)/ratio;
47+
48+
return (uint32_t) std::min<uint64_t>(kv_size, n_rows);
49+
}
50+
4151
static int64_t dsv4_stream_offset(uint32_t n_stream, llama_seq_id seq_id, uint32_t size) {
4252
if (n_stream <= 1) {
4353
return 0;
@@ -239,6 +249,7 @@ static void dsv4_state_dst_stream_range(
239249
static void dsv4_state_write_tensor_streams(
240250
llama_io_write_i & io,
241251
ggml_tensor * tensor,
252+
uint32_t tensor_rows,
242253
uint32_t n_rows,
243254
uint32_t s0,
244255
uint32_t ns) {
@@ -247,20 +258,31 @@ static void dsv4_state_write_tensor_streams(
247258
const uint64_t rows = n_rows;
248259
const uint64_t row_size = ggml_row_size(tensor->type, tensor->ne[0]);
249260

261+
if (n_rows > tensor_rows) {
262+
throw std::runtime_error("DSV4 state tensor row count exceeds storage");
263+
}
264+
250265
io.write(&type_i, sizeof(type_i));
251266
io.write(&ne0, sizeof(ne0));
252267
io.write(&rows, sizeof(rows));
253268
io.write(&row_size, sizeof(row_size));
254269

255-
const size_t offset = (size_t) s0*n_rows*row_size;
256-
const size_t size = (size_t) ns*n_rows*row_size;
270+
const size_t stream_stride = (size_t) tensor_rows*row_size;
271+
const size_t size = (size_t) n_rows*row_size;
272+
if (size == 0) {
273+
return;
274+
}
257275

258-
io.write_tensor(tensor, offset, size);
276+
for (uint32_t s = 0; s < ns; ++s) {
277+
const size_t offset = (size_t) (s0 + s)*stream_stride;
278+
io.write_tensor(tensor, offset, size);
279+
}
259280
}
260281

261282
static void dsv4_state_read_tensor_streams(
262283
llama_io_read_i & io,
263284
ggml_tensor * tensor,
285+
uint32_t tensor_rows,
264286
uint32_t n_rows,
265287
uint32_t s0,
266288
uint32_t ns) {
@@ -282,18 +304,28 @@ static void dsv4_state_read_tensor_streams(
282304
if (type_i != type_i_ref || ne0 != ne0_ref || rows != rows_ref || row_size != row_size_ref) {
283305
throw std::runtime_error("DSV4 state tensor metadata mismatch");
284306
}
307+
if (n_rows > tensor_rows) {
308+
throw std::runtime_error("DSV4 state tensor row count exceeds storage");
309+
}
285310

286-
const size_t offset = (size_t) s0*n_rows*row_size;
287-
const size_t size = (size_t) ns*n_rows*row_size;
311+
const size_t stream_stride = (size_t) tensor_rows*row_size;
312+
const size_t size = (size_t) n_rows*row_size;
313+
if (size == 0) {
314+
return;
315+
}
288316

289-
io.read_tensor(tensor, offset, size);
317+
for (uint32_t s = 0; s < ns; ++s) {
318+
const size_t offset = (size_t) (s0 + s)*stream_stride;
319+
io.read_tensor(tensor, offset, size);
320+
}
290321
}
291322

292323
static void dsv4_state_write_k_cache(
293324
llama_io_write_i & io,
294325
const llama_kv_cache * kv,
295326
llama_seq_id seq_id,
296-
llama_state_seq_flags flags) {
327+
llama_state_seq_flags flags,
328+
uint32_t n_rows) {
297329
GGML_UNUSED(flags);
298330

299331
uint32_t s0;
@@ -305,14 +337,18 @@ static void dsv4_state_write_k_cache(
305337
const auto layer_ids = kv->get_layer_ids();
306338
const uint32_t n_layer = layer_ids.size();
307339

340+
if (n_rows > kv_size) {
341+
throw std::runtime_error("DSV4 K-cache state row count exceeds cache size");
342+
}
343+
308344
io.write(&version, sizeof(version));
309-
io.write(&kv_size, sizeof(kv_size));
345+
io.write(&n_rows, sizeof(n_rows));
310346
io.write(&ns, sizeof(ns));
311347
io.write(&n_layer, sizeof(n_layer));
312348

313349
for (uint32_t il : layer_ids) {
314350
io.write(&il, sizeof(il));
315-
dsv4_state_write_tensor_streams(io, kv->get_k_storage(il), kv_size, s0, ns);
351+
dsv4_state_write_tensor_streams(io, kv->get_k_storage(il), kv_size, n_rows, s0, ns);
316352
}
317353
}
318354

@@ -324,19 +360,26 @@ static void dsv4_state_read_k_cache(
324360
GGML_UNUSED(flags);
325361

326362
uint32_t version;
327-
uint32_t kv_size_ref;
363+
uint32_t n_rows_ref;
328364
uint32_t ns;
329365
uint32_t n_layer_ref;
330366

331367
io.read(&version, sizeof(version));
332-
io.read(&kv_size_ref, sizeof(kv_size_ref));
368+
io.read(&n_rows_ref, sizeof(n_rows_ref));
333369
io.read(&ns, sizeof(ns));
334370
io.read(&n_layer_ref, sizeof(n_layer_ref));
335371

336-
if (version != DSV4_K_CACHE_STATE_VER) {
372+
if (version != 1 && version != DSV4_K_CACHE_STATE_VER) {
337373
throw std::runtime_error("DSV4 K-cache state version mismatch");
338374
}
339-
if (kv_size_ref != kv->get_size()) {
375+
376+
const uint32_t kv_size = kv->get_size();
377+
if (version == 1 && n_rows_ref != kv_size) {
378+
LLAMA_LOG_INFO("kv size ref %d kv %d\n", n_rows_ref, kv_size);
379+
throw std::runtime_error("DSV4 K-cache state size mismatch");
380+
}
381+
if (n_rows_ref > kv_size) {
382+
LLAMA_LOG_INFO("kv rows ref %d kv %d\n", n_rows_ref, kv_size);
340383
throw std::runtime_error("DSV4 K-cache state size mismatch");
341384
}
342385

@@ -355,7 +398,7 @@ static void dsv4_state_read_k_cache(
355398
throw std::runtime_error("DSV4 K-cache layer id mismatch");
356399
}
357400

358-
dsv4_state_read_tensor_streams(io, kv->get_k_storage(il), kv->get_size(), s0, ns);
401+
dsv4_state_read_tensor_streams(io, kv->get_k_storage(il), kv_size, n_rows_ref, s0, ns);
359402
}
360403
}
361404

@@ -882,8 +925,8 @@ void llama_dsv4_comp_state::state_write(llama_io_write_i & io, llama_seq_id seq_
882925
for (const auto & layer : layers) {
883926
io.write(&layer.il, sizeof(layer.il));
884927

885-
dsv4_state_write_tensor_streams(io, layer.kv, state_size, s0, ns);
886-
dsv4_state_write_tensor_streams(io, layer.score, state_size, s0, ns);
928+
dsv4_state_write_tensor_streams(io, layer.kv, state_size, state_size, s0, ns);
929+
dsv4_state_write_tensor_streams(io, layer.score, state_size, state_size, s0, ns);
887930
}
888931
}
889932

@@ -924,8 +967,8 @@ void llama_dsv4_comp_state::state_read(llama_io_read_i & io, llama_seq_id seq_id
924967
throw std::runtime_error("DSV4 compressor state layer id mismatch");
925968
}
926969

927-
dsv4_state_read_tensor_streams(io, layer.kv, state_size, s0, ns);
928-
dsv4_state_read_tensor_streams(io, layer.score, state_size, s0, ns);
970+
dsv4_state_read_tensor_streams(io, layer.kv, state_size, state_size, s0, ns);
971+
dsv4_state_read_tensor_streams(io, layer.score, state_size, state_size, s0, ns);
929972
}
930973
}
931974

@@ -1328,9 +1371,19 @@ void llama_kv_cache_dsv4::state_write(llama_io_write_i & io, llama_seq_id seq_id
13281371
kv_raw->state_write(io, seq_id, flags);
13291372

13301373
if (!partial_only) {
1331-
dsv4_state_write_k_cache(io, kv_csa.get(), seq_id, flags);
1332-
dsv4_state_write_k_cache(io, kv_hca.get(), seq_id, flags);
1333-
dsv4_state_write_k_cache(io, kv_lid.get(), seq_id, flags);
1374+
const llama_pos pos_max = seq_id >= 0 ? kv_raw->seq_pos_max(seq_id) : -1;
1375+
1376+
//FIXME : note that we conflate token positions with rows, which is not true for multi-modal case.
1377+
const uint32_t n_rows_csa = seq_id >= 0 ?
1378+
dsv4_state_n_used_k_rows(pos_max, DSV4_CSA_RATIO, kv_csa->get_size()) : kv_csa->get_size();
1379+
const uint32_t n_rows_hca = seq_id >= 0 ?
1380+
dsv4_state_n_used_k_rows(pos_max, DSV4_HCA_RATIO, kv_hca->get_size()) : kv_hca->get_size();
1381+
const uint32_t n_rows_lid = seq_id >= 0 ?
1382+
dsv4_state_n_used_k_rows(pos_max, DSV4_CSA_RATIO, kv_lid->get_size()) : kv_lid->get_size();
1383+
1384+
dsv4_state_write_k_cache(io, kv_csa.get(), seq_id, flags, n_rows_csa);
1385+
dsv4_state_write_k_cache(io, kv_hca.get(), seq_id, flags, n_rows_hca);
1386+
dsv4_state_write_k_cache(io, kv_lid.get(), seq_id, flags, n_rows_lid);
13341387
}
13351388

13361389
csa_state->state_write(io, seq_id, flags);
@@ -1366,6 +1419,10 @@ void llama_kv_cache_dsv4::state_read(llama_io_read_i & io, llama_seq_id seq_id,
13661419
kv_raw->state_read(io, seq_id, flags);
13671420

13681421
if (!partial_only) {
1422+
kv_csa->clear(true);
1423+
kv_hca->clear(true);
1424+
kv_lid->clear(true);
1425+
13691426
dsv4_state_read_k_cache(io, kv_csa.get(), seq_id, flags);
13701427
dsv4_state_read_k_cache(io, kv_hca.get(), seq_id, flags);
13711428
dsv4_state_read_k_cache(io, kv_lid.get(), seq_id, flags);

0 commit comments

Comments
 (0)