diff --git a/src/llama-context.cpp b/src/llama-context.cpp index a6d2f264fc..a920c42311 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2907,17 +2907,28 @@ public: } if (mbuf_cur.n_tensors == mbuf.n_tensors) { - // same chunking: copy 1:1 by index + // an equal tensor count does not imply the same chunking, e.g. save ranges [2,1] vs restore runs [1,2] + bool same_chunking = true; for (size_t i = 0; i < mbuf_cur.org.size(); ++i) { - GGML_ASSERT(ggml_nbytes(mbuf_cur.cpy[i]) == ggml_nbytes(mbuf.org[i])); - ggml_backend_tensor_copy(mbuf_cur.cpy[i], mbuf.org[i]); + if (ggml_nbytes(mbuf_cur.cpy[i]) != ggml_nbytes(mbuf.org[i])) { + same_chunking = false; + break; + } + } + + if (same_chunking) { + // same chunking: copy 1:1 by index + for (size_t i = 0; i < mbuf_cur.org.size(); ++i) { + ggml_backend_tensor_copy(mbuf_cur.cpy[i], mbuf.org[i]); + } + continue; } - continue; } // different chunking: copy the write-side data (mbuf_cur.cpy) into the read-side targets (mbuf.org) // with a byte cursor. Write and read enumerate the same logical data in the same order but may chunk - // it differently, so copy across tensor boundaries rather than 1:1 by index. + // it differently (even with an equal number of tensors), so copy across tensor boundaries rather than + // 1:1 by index. const size_t total = mbuf_cur.total_size; ggml_init_params params_scratch = { diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 65afbd8c37..0e095f8b6b 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -2533,6 +2533,24 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32 bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32_t cell_count, const slot_info & sinfo) { auto & cells = v_cells[strm]; + // batch the scatter reads per contiguous run of destination indices + // from inclusive, to exclusive - same convention as cell_ranges_t + // contiguous cells yield a single run covering the whole block + struct cell_run { uint32_t from; uint32_t to; }; + std::vector runs; + if (cell_count > 0) { + const auto & idxs = sinfo.idxs[0]; + uint32_t i0 = 0; + while (i0 < cell_count) { + uint32_t i1 = i0 + 1; + while (i1 < cell_count && idxs[i1] == idxs[i1 - 1] + 1) { + ++i1; + } + runs.push_back({idxs[i0], idxs[i1 - 1] + 1}); + i0 = i1; + } + } + uint32_t v_trans; uint32_t n_layer; @@ -2580,17 +2598,8 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32 return false; } - if (cell_count) { - if (sinfo.is_contiguous()) { - // Fast path: contiguous cells, single memcpy - io.read_tensor(k, sinfo.head() * k_size_row, cell_count * k_size_row); - } else { - // Slow path: scatter to non-contiguous positions - for (uint32_t i = 0; i < cell_count; ++i) { - const size_t dst_offset = sinfo.idxs[0][i] * k_size_row; - io.read_tensor(k, dst_offset, k_size_row); - } - } + for (const auto & r : runs) { + io.read_tensor(k, (size_t) r.from * k_size_row, (size_t) (r.to - r.from) * k_size_row); } } @@ -2623,17 +2632,8 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32 return false; } - if (cell_count) { - if (sinfo.is_contiguous()) { - // Fast path: contiguous cells, single memcpy - io.read_tensor(v, sinfo.head() * v_size_row, cell_count * v_size_row); - } else { - // Slow path: scatter to non-contiguous positions - for (uint32_t i = 0; i < cell_count; ++i) { - const size_t dst_offset = sinfo.idxs[0][i] * v_size_row; - io.read_tensor(v, dst_offset, v_size_row); - } - } + for (const auto & r : runs) { + io.read_tensor(v, (size_t) r.from * v_size_row, (size_t) (r.to - r.from) * v_size_row); } } } else { @@ -2674,22 +2674,10 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32 return false; } - if (cell_count) { - if (sinfo.is_contiguous()) { - // Fast path: contiguous cells - const uint32_t h = sinfo.head(); - for (uint32_t j = 0; j < n_embd_v_gqa; ++j) { - const size_t dst_offset = (h + j * cells.size()) * v_size_el; - io.read_tensor(v, dst_offset, cell_count * v_size_el); - } - } else { - // Slow path: scatter to non-contiguous positions - for (uint32_t j = 0; j < n_embd_v_gqa; ++j) { - for (uint32_t i = 0; i < cell_count; ++i) { - const size_t dst_offset = (sinfo.idxs[0][i] + j * cells.size()) * v_size_el; - io.read_tensor(v, dst_offset, v_size_el); - } - } + for (uint32_t j = 0; j < n_embd_v_gqa; ++j) { + for (const auto & r : runs) { + const size_t dst_offset = ((size_t) r.from + j * cells.size()) * v_size_el; + io.read_tensor(v, dst_offset, (size_t) (r.to - r.from) * v_size_el); } } } diff --git a/tests/test-save-load-state.cpp b/tests/test-save-load-state.cpp index 0ceab7c545..35c7690587 100644 --- a/tests/test-save-load-state.cpp +++ b/tests/test-save-load-state.cpp @@ -355,7 +355,101 @@ static bool test_seq_cp_device(struct llama_model * model, const struct common_p } -// Run the full save/load test suite (tests 1-5) for a single model. +// Test 6/7: seq copy (scatter) +// - decode the same prefix on two sequences, interleaving seq 0 cells between the seq 1 cells +// - save the seq 1 state, free the interleaved seq 0 cells, and restore via the given io path +// - the restore destination is non-contiguous: scatter reads are batched per contiguous run +// - save again on the host and compare the two blobs byte for byte +static bool test_seq_cp_scatter(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens, int test_num, bool on_device) { + 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)}; + + LOG("\n=== Test %d: seq copy (%s, scatter) ===\n", test_num, on_device ? "device" : "host"); + + const uint32_t flags = on_device ? LLAMA_STATE_SEQ_FLAGS_ON_DEVICE : LLAMA_STATE_SEQ_FLAGS_NONE; + + auto decode_one = [&](llama_token tok, int pos, llama_seq_id seq) { + llama_batch_ptr batch(1, 0, 1); + common_batch_add(batch.get(), tok, pos, { seq }, false); + return llama_decode(ctx.get(), batch.get()) == 0; + }; + + // seq 0 cells 0,1,4 interleave the seq 1 cells 2,3,5 + if (!decode_one(tokens[0], 0, 0) || + !decode_one(tokens[1], 1, 0) || + !decode_one(tokens[0], 0, 1) || + !decode_one(tokens[1], 1, 1) || + !decode_one(tokens[2], 2, 0) || + !decode_one(tokens[2], 2, 1)) { + LOG_ERR("%s: failed to build interleaved state\n", __func__); + return false; + } + + const auto get_seq_state = [&](llama_seq_id seq_id, uint32_t fl, std::vector & state) { + const size_t state_size = llama_state_seq_get_size_ext(ctx.get(), seq_id, fl); + 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_ext(ctx.get(), state.data(), state.size(), seq_id, fl); + 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; + }; + + // host blob: contains the KV data, used for the byte-for-byte comparison + std::vector state_before; + if (!get_seq_state(1, LLAMA_STATE_SEQ_FLAGS_NONE, state_before)) { + return false; + } + + // save via the io path under test + std::vector state_save; + if (!get_seq_state(1, flags, state_save)) { + return false; + } + LOG_TRC("%s: seq 1 saved via %s, %zu bytes\n", __func__, on_device ? "device" : "host", state_save.size()); + + // free seq 0's cells so the ring is fragmented: the restore destination (seq 1's interleaved cells) stays non-contiguous + 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; + } + + // restore via the io path under test + const size_t nset = llama_state_seq_set_data_ext(ctx.get(), state_save.data(), state_save.size(), 1, flags); + if (nset != state_save.size()) { + LOG_ERR("%s: seq set data length %zu does not match expected length %zu\n", __func__, nset, state_save.size()); + return false; + } + LOG_TRC("%s: seq 1 restored via %s, %zu bytes\n", __func__, on_device ? "device" : "host", nset); + + std::vector state_after; + if (!get_seq_state(1, LLAMA_STATE_SEQ_FLAGS_NONE, state_after)) { + return false; + } + + // the blob is serialized in sequence cell order, so identical bytes iff the restore wrote the same KV + if (state_before.size() != state_after.size() || memcmp(state_before.data(), state_after.data(), state_before.size()) != 0) { + LOG_ERR("\n%s: error: restored KV state is not byte-identical to the saved state\n", __func__); + return false; + } + + LOG("\nPASS\n"); + return true; +} + + +// Run the full save/load test suite (tests 1-7) for a single model. // Returns true if all tests pass, false otherwise. static bool run_save_load_tests_for_model(const std::string & model_path, const struct common_params & base_params) { struct common_params params = base_params; @@ -422,6 +516,16 @@ static bool run_save_load_tests_for_model(const std::string & model_path, const return false; } + // Test 6: seq copy (host, scatter) + if (!test_seq_cp_scatter(model, params, tokens, 6, false)) { + return false; + } + + // Test 7: seq copy (device, scatter) + if (!test_seq_cp_scatter(model, params, tokens, 7, true)) { + return false; + } + LOG("\nAll tests passed.\n"); return true;