From fc349bd785d8bd878235039d8e35a186d872d0cb Mon Sep 17 00:00:00 2001 From: ARC <261443955+arc-uri-el@users.noreply.github.com> Date: Tue, 8 Sep 2026 12:54:30 +0200 Subject: [PATCH] context : size device checkpoint views in quantized elements A quantized type size describes a block rather than one tensor element. Include the block size when converting saved bytes into view dimensions in the device writer, reader and fragmented-copy fallback. Extend the existing save/load regression with fragmented-to-compact restoration. Assisted-by: GPT 6 Astra --- src/llama-context.cpp | 7 +++---- tests/test-save-load-state.cpp | 16 +++++++++++++--- 2 files changed, 16 insertions(+), 7 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index e202c68400f6..31750f071e12 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2788,7 +2788,7 @@ class llama_io_write_device : public llama_io_write_i { for (const auto & winfo : winfos) { auto * buft = ggml_backend_buffer_get_type(winfo.tensor->buffer); - const int64_t n = winfo.size/ggml_element_size(winfo.tensor); + const int64_t n = (winfo.size / ggml_type_size(winfo.tensor->type)) * ggml_blck_size(winfo.tensor->type); auto & mbuf = mbufs_new[buft]; @@ -2919,7 +2919,7 @@ class llama_io_read_device : public llama_io_read_i { for (const auto & rinfo : rinfos) { auto * buft = ggml_backend_buffer_get_type(rinfo.tensor->buffer); - const int64_t n = rinfo.size/ggml_element_size(rinfo.tensor); + const int64_t n = (rinfo.size / ggml_type_size(rinfo.tensor->type)) * ggml_blck_size(rinfo.tensor->type); auto & mbuf = mbufs_new[buft]; @@ -2986,8 +2986,7 @@ class llama_io_read_device : public llama_io_read_i { const size_t n_copy = std::min(src_size - src_off, dst_size - dst_off); - const size_t el = ggml_element_size(src_t); - const int64_t n_el = (int64_t) (n_copy / el); + const int64_t n_el = (n_copy / ggml_type_size(src_t->type)) * ggml_blck_size(src_t->type); auto * src_v = ggml_view_1d(ctx_scratch, src_t, n_el, src_off); ggml_backend_view_init(src_v); diff --git a/tests/test-save-load-state.cpp b/tests/test-save-load-state.cpp index 6179e6c10848..618b6635e988 100644 --- a/tests/test-save-load-state.cpp +++ b/tests/test-save-load-state.cpp @@ -360,14 +360,14 @@ static bool test_seq_cp_device(struct llama_model * model, const struct common_p // - 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) { +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, bool compact = false) { 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"); + LOG("\n=== Test %d: seq copy (%s, %s) ===\n", test_num, on_device ? "device" : "host", compact ? "scatter to compact" : "scatter"); const uint32_t flags = on_device ? LLAMA_STATE_SEQ_FLAGS_ON_DEVICE : LLAMA_STATE_SEQ_FLAGS_NONE; @@ -425,6 +425,11 @@ static bool test_seq_cp_scatter(struct llama_model * model, const struct common_ return false; } + if (compact) { + // Saved ranges [2, 1] must restore into one contiguous range, including quantized KV blocks. + llama_memory_clear(llama_get_memory(ctx.get()), true); + } + // 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()) { @@ -508,7 +513,7 @@ static bool test_state_roundtrip(struct llama_model * model, const struct common } -// Run the full save/load test suite (tests 1-8) for a single model. +// Run the full save/load test suite (tests 1-9) 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; @@ -590,6 +595,11 @@ static bool run_save_load_tests_for_model(const std::string & model_path, const return false; } + // Test 9: on-device restore with different source and destination chunking. + if (!test_seq_cp_scatter(model, params, tokens, 9, true, true)) { + return false; + } + LOG("\nAll tests passed.\n"); return true;