From 3ec6ec16850c3cdcaebad3365f654d04bf8a196a Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Wed, 11 Nov 2026 16:06:18 -0400 Subject: [PATCH 1/4] hrx: prefill keeps the stored layout for a weight that already has a stored device copy (engine #2137) On a mixed-quant dense model the one-token decode goes through the kquant family, which reads the GGUF block layout, so a prefill kernel that reads the same weight in the packed rows makes a second device copy. HostWeightCache now records the host address of each stored-layout copy it makes, and the per-weight default keeps the stored layout for such a weight; the 6-to-16-token Q6_K i8 kernel steps aside for it too. GGML_HRX_PREFILL_FOLLOWS_DECODE=0 turns the lookup off. Co-Authored-By: Claude Opus 5.5 --- .../common/dispatch-mul-mat.cpp | 4 +- .../common/dispatch-q4k-packed-tail.h | 45 +++++++++++++++++-- ggml/src/ggml-hrx/model-buffer-from-mmap.cpp | 29 ++++++++++++ ggml/src/ggml-hrx/model-buffer-from-mmap.h | 7 +++ ggml/src/ggml-hrx/runtime/host-memory.cpp | 8 ++++ 5 files changed, 87 insertions(+), 6 deletions(-) diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat.cpp index 3c9dbe3b0242..da5611e1226a 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat.cpp @@ -659,8 +659,8 @@ static bool match_q6_k_i8_prepacked_dispatch(const DispatchMatchContext & contex // decode and the other prefill batches read as packed rows. Where the packed rows are the weight's layout (a weight // in the file mapping by default, see dispatch-q4k-packed-tail.h), leave the 6-to-16-token batch to the tiled // kernel, which reads them too. - if (q6k_packed_rows_weight_ok(*match.weight, match.input_size, match.output_size) && - !weight_has_native_layout_reader(context.graph, *match.weight) && q6k_packed_rows_enabled(*match.weight)) { + if (!weight_has_native_layout_reader(context.graph, *match.weight) && + q6k_i8_layout_is_extra_copy(*match.weight, match.input_size, match.output_size)) { return false; } diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-q4k-packed-tail.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-q4k-packed-tail.h index 40d7f8e5e4a2..785e032ca62c 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-q4k-packed-tail.h +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-q4k-packed-tail.h @@ -56,16 +56,43 @@ inline PackedLayoutSwitch packed_layout_switch(const char * name, const char * w return value; } -// The weight's own buffer decides the default (see the header comment). -inline bool packed_layout_wanted(PackedLayoutSwitch value, const Value & weight) { - if (value != PackedLayoutSwitch::FollowWeight) { - return value == PackedLayoutSwitch::Packed; +// engine #2137: a weight that already has a device copy in the stored layout. On a mixed-quant dense model +// (Qwen3.8-27B UD-Q4_K_XL) the one-token decode goes through the kquant family, which reads the GGUF block layout, so +// a prefill kernel that reads the same weight in the packed rows makes a second device copy of it. Where the prefill +// kernel has a stored form, it keeps the stored layout for such a weight. The copy is recorded when it is made +// (model-buffer-from-mmap.h), so this follows what ran, not what a matcher was asked about. A weight is known once a +// decode has run, which the tools' warm-up does before the first prompt. +// GGML_HRX_PREFILL_FOLLOWS_DECODE=0 turns the lookup off. +inline bool prefill_follows_decode_enabled() { + static const bool enabled = [] { + const char * env = std::getenv("GGML_HRX_PREFILL_FOLLOWS_DECODE"); + return env == nullptr || env[0] != '0'; + }(); + return enabled; +} + +inline bool decode_reads_stored(const Value & weight) { + const ggml_tensor * tensor = weight.tensor; + if (tensor == nullptr || !prefill_follows_decode_enabled()) { + return false; } + return has_stored_device_copy(tensor->view_src != nullptr ? tensor->view_src->data : tensor->data); +} + +inline bool weight_in_file_mapping(const Value & weight) { const ggml_tensor * tensor = weight.tensor; return tensor != nullptr && buffer_is_file_mapped(tensor->view_src != nullptr ? tensor->view_src->buffer : tensor->buffer); } +// The weight's own buffer decides the default (see the header comment), unless decode reads the weight stored. +inline bool packed_layout_wanted(PackedLayoutSwitch value, const Value & weight) { + if (value != PackedLayoutSwitch::FollowWeight) { + return value == PackedLayoutSwitch::Packed; + } + return weight_in_file_mapping(weight) && !decode_reads_stored(weight); +} + inline bool q4k_packed_tail_enabled(const Value & weight) { static const PackedLayoutSwitch value = packed_layout_switch("GGML_HRX_Q4K_PACKED_TAIL", "q4_k batches off the 256 grid"); @@ -104,4 +131,14 @@ inline bool q6k_packed_rows_weight_ok(const Value & weight, int64_t input_size, return packed_tail_weight_ok(weight, GGML_TYPE_Q6_K, input_size, output_size); } +// The 6-to-16-token Q6_K kernel reads an i8 layout of its own. That is an extra device copy of a weight in the file +// mapping whichever of the other two layouts the weight is read in elsewhere, so the kernel is left out for such a +// weight unless the stored layout is forced for every weight (the previous behaviour). +inline bool q6k_i8_layout_is_extra_copy(const Value & weight, int64_t input_size, int64_t output_size) { + if (!q6k_packed_rows_weight_ok(weight, input_size, output_size)) { + return false; + } + return q6k_packed_rows_enabled(weight) || (weight_in_file_mapping(weight) && decode_reads_stored(weight)); +} + } // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/model-buffer-from-mmap.cpp b/ggml/src/ggml-hrx/model-buffer-from-mmap.cpp index 3a7426824ef6..159ff5608c64 100644 --- a/ggml/src/ggml-hrx/model-buffer-from-mmap.cpp +++ b/ggml/src/ggml-hrx/model-buffer-from-mmap.cpp @@ -19,6 +19,8 @@ #include "ggml-backend-impl.h" #include "runtime/host-memory.h" +#include +#include #include #include @@ -62,6 +64,33 @@ bool model_buffer_from_mmap_enabled() { return HostWeightCache::residency_limit() == 0; } +namespace { +std::mutex & stored_copy_mutex() { + static std::mutex mutex; + return mutex; +} +std::unordered_set & stored_copy_addresses() { + static std::unordered_set addresses; + return addresses; +} +} // namespace + +void note_stored_device_copy(const void * host_address) { + if (host_address == nullptr) { + return; + } + const std::lock_guard lock(stored_copy_mutex()); + stored_copy_addresses().insert(host_address); +} + +bool has_stored_device_copy(const void * host_address) { + if (host_address == nullptr) { + return false; + } + const std::lock_guard lock(stored_copy_mutex()); + return stored_copy_addresses().count(host_address) != 0; +} + bool buffer_is_file_mapped(ggml_backend_buffer_t buffer) { return buffer != nullptr && buffer->buft != nullptr && buffer->buft->iface.get_name == mapped_type_name; } diff --git a/ggml/src/ggml-hrx/model-buffer-from-mmap.h b/ggml/src/ggml-hrx/model-buffer-from-mmap.h index a9f92802c233..79cb1386eba6 100644 --- a/ggml/src/ggml-hrx/model-buffer-from-mmap.h +++ b/ggml/src/ggml-hrx/model-buffer-from-mmap.h @@ -32,6 +32,13 @@ bool model_buffer_from_mmap_enabled(); // weight in such a buffer costs a device copy for every layout the kernels read it in (engine #2137). bool buffer_is_file_mapped(ggml_backend_buffer_t buffer); +// engine #2137: host addresses of weights that have a device copy in the stored (GGUF) layout, recorded when the copy +// is made. A matcher that can read a weight either stored or packed asks this before choosing, so it does not add a +// packed copy next to a stored one that a decode kernel already made. Keyed by the weight's host address, which is +// where both the matcher (tensor->data) and the copy (buffer base + offset) can see it. +void note_stored_device_copy(const void * host_address); +bool has_stored_device_copy(const void * host_address); + // The buffer type of these buffers ("_Mapped"). It is host memory; only this device claims it. void model_buffer_from_mmap_init_type(ggml_backend_hrx_device_context * device); diff --git a/ggml/src/ggml-hrx/runtime/host-memory.cpp b/ggml/src/ggml-hrx/runtime/host-memory.cpp index 200eefa06b8b..f608795a7909 100644 --- a/ggml/src/ggml-hrx/runtime/host-memory.cpp +++ b/ggml/src/ggml-hrx/runtime/host-memory.cpp @@ -1,5 +1,7 @@ #include "host-memory.h" +#include "model-buffer-from-mmap.h" + #include "ggml-impl.h" #include "ggml-quants.h" #include "hrx-interop-utils.h" @@ -1929,6 +1931,12 @@ HostWeightAcquireResult HostWeightCache::acquire(hrx_device_t device } } + // engine #2137: a device copy in the stored layout is about to be made of a host weight; matchers that can read + // the weight either way ask for this before adding a packed copy (model-buffer-from-mmap.h). + if (source.host_data != nullptr && source.layout == kNativeWeightLayout) { + note_stored_device_copy(static_cast(source.host_data) + source.offset); + } + // GGML_HRX_HOST_WEIGHT_TRACE=2: one line per device copy made (a cache miss), so the copies of a run can be counted // per weight and per layout (engine #2137: a weight read in two layouts shows up as two lines with one id/off). // At ERROR because the CLI tools drop INFO. From 53c999c2d31cdf56a833140ec7b22ed398f44237 Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Wed, 11 Nov 2026 16:24:49 -0400 Subject: [PATCH 2/4] tests: a file-mapped weight with a stored device copy is read stored by prefill (engine #2137) hrx-backend-test packed_layout_follows_weight.{q4_k,q6_k}.file_mapped.stored_copy and q6_k.file_mapped.stored_copy.tokens12, device-free. Co-Authored-By: Claude Opus 5.5 --- tests/hrx-backend-test.cpp | 20 +++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/tests/hrx-backend-test.cpp b/tests/hrx-backend-test.cpp index 7aba5f604392..7976a15ffecb 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -3396,7 +3396,8 @@ static bool matmul_graph_is_supported(ggml_context * ctx, ggml_tensor * output) // stored layout on the device, so it stays stored (4 / 6) and the packed copy is not added. // GGML_HRX_Q4K_PACKED_TAIL / GGML_HRX_Q6K_PACKED_ROWS force one or the other; the row is about the default, so it does // nothing when either is set. -static void schedule_packed_layout_follows_weight(ggml_type weight_type, bool file_mapped, int64_t token_count = 100) { +static void schedule_packed_layout_follows_weight(ggml_type weight_type, bool file_mapped, int64_t token_count = 100, + bool stored_copy_exists = false) { if (std::getenv("GGML_HRX_Q4K_PACKED_TAIL") != nullptr || std::getenv("GGML_HRX_Q6K_PACKED_ROWS") != nullptr) { std::printf(" skipped: a packed-layout switch is set in the environment\n"); return; @@ -3433,6 +3434,11 @@ static void schedule_packed_layout_follows_weight(ggml_type weight_type, bool fi ggml_backend_buffer_set_usage(buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); ggml_tallocr allocator = ggml_tallocr_new(buffer); REQUIRE(ggml_tallocr_alloc(&allocator, weight) == GGML_STATUS_SUCCESS); + if (stored_copy_exists) { + // engine #2137: a decode kernel already made a stored-layout device copy of this weight + ggml::hrx::note_stored_device_copy(weight->data); + REQUIRE(ggml::hrx::has_stored_device_copy(weight->data)); + } ggml_tensor * output = ggml_mul_mat(ctx, weight, input); REQUIRE(output != nullptr); @@ -3447,7 +3453,7 @@ static void schedule_packed_layout_follows_weight(ggml_type weight_type, bool fi REQUIRE(scheduler.plan().valid()); const int stored = weight_type == GGML_TYPE_Q4_K ? 4 : 6; - const std::string expected = std::to_string(file_mapped ? stored + 40 : stored); + const std::string expected = std::to_string(file_mapped && !stored_copy_exists ? stored + 40 : stored); size_t found = 0; for (const ggml::hrx::Dispatch & dispatch : scheduler.plan().dispatches) { const auto format = dispatch.kernel.compile_parameters.find("ggml.mul_mat.weight_format"); @@ -3462,7 +3468,7 @@ static void schedule_packed_layout_follows_weight(ggml_type weight_type, bool fi } REQUIRE(format->second == expected); } - if (weight_type == GGML_TYPE_Q6_K && !file_mapped && token_count >= 6 && token_count <= 16) { + if (weight_type == GGML_TYPE_Q6_K && !file_mapped && !stored_copy_exists && token_count >= 6 && token_count <= 16) { // an uploaded Q6_K weight keeps the i8 kernel at this batch size; it carries no weight_format REQUIRE(found == 0); REQUIRE(scheduler.plan().dispatches.size() == 1); @@ -11954,6 +11960,14 @@ int main(int argc, char ** argv) { [] { schedule_packed_layout_follows_weight(GGML_TYPE_Q6_K, true, 12); } }, { "packed_layout_follows_weight.q6_k.uploaded.tokens12", [] { schedule_packed_layout_follows_weight(GGML_TYPE_Q6_K, false, 12); } }, + // engine #2137: a file-mapped weight that already has a stored device copy (a decode kernel read it stored) is + // read stored by prefill too, so it is not copied a second time; and not in the i8 layout at 6 to 16 tokens + { "packed_layout_follows_weight.q4_k.file_mapped.stored_copy", + [] { schedule_packed_layout_follows_weight(GGML_TYPE_Q4_K, true, 100, true); } }, + { "packed_layout_follows_weight.q6_k.file_mapped.stored_copy", + [] { schedule_packed_layout_follows_weight(GGML_TYPE_Q6_K, true, 100, true); } }, + { "packed_layout_follows_weight.q6_k.file_mapped.stored_copy.tokens12", + [] { schedule_packed_layout_follows_weight(GGML_TYPE_Q6_K, true, 12, true); } }, { "q6k_token1_shortlist_group_count.256", [] { schedule_q6k_token1_shortlist_group_count(256, 65536, 8); } }, { "q6k_token1_shortlist_group_count.2560", [] { schedule_q6k_token1_shortlist_group_count(2560, 248320, 80); } }, From ce27788b2387f0de77707e356fe4861193bda69d Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Wed, 11 Nov 2026 16:36:14 -0400 Subject: [PATCH 3/4] hrx: the 256-aligned Q4_K prefill reads a weight stored when decode already holds it stored (engine #2137) The f16 prefill kernel reads Q4_K only in the packed rows. The 8-bit-block kernel (ggml_mul_mat_q5_k_iq4_xs_q8_1_x4_wmma_token256) reads Q4_K stored as well (weight format 4). For a Q4_K weight that already has a stored device copy the aligned prefill now takes that kernel over the stored layout, so it shares decode's copy instead of adding a packed one. GGML_HRX_Q4K_PACKED_TAIL=0 selects it for every Q4_K weight. Test (device-free): hrx-backend-test q4k_aligned_prefill_layout.{packed,stored_copy}. Co-Authored-By: Claude Opus 5.5 --- .../common/dispatch-mul-mat.cpp | 8 ++- .../common/dispatch-q4k-packed-tail.h | 23 ++++++- tests/hrx-backend-test.cpp | 64 +++++++++++++++++++ 3 files changed, 91 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat.cpp index da5611e1226a..58499e701f68 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat.cpp @@ -926,7 +926,11 @@ static bool match_packed_q8_1_x4_prefill_dispatch(const DispatchMatchContext & c } } - const bool use_f16 = !split_tail && match.weight->type == GGML_TYPE_Q4_K && + // engine #2137 (dispatch-q4k-packed-tail.h): a Q4_K weight that already has a stored copy is read stored, by the + // 8-bit-block kernel, instead of being copied into the packed rows for the f16 kernel + const bool stored_q4 = match.weight_format == CommonMulMatWeightFormat::Q4K && + q4k_aligned_prefill_reads_stored(*match.weight); + const bool use_f16 = !split_tail && !stored_q4 && match.weight->type == GGML_TYPE_Q4_K && match.weight->alias_source.value < 0 && match.output_size >= match.input_size / 4; DispatchBinding activation; int activation_scale_log2 = 0; @@ -964,7 +968,7 @@ static bool match_packed_q8_1_x4_prefill_dispatch(const DispatchMatchContext & c "ggml.mul_mat_q8_1_x4.token_capacity", common_to_config_value(q8_tokens)); const bool pack_q4 = match.weight_format == CommonMulMatWeightFormat::Q4K && - match.weight->alias_source.value < 0; + match.weight->alias_source.value < 0 && !stored_q4; if (use_f16) { dispatch.kernel.compile_parameters.emplace("ggml.mul_mat.f16_input_layout", packed_input ? "1" : "0"); common_set_f16_prefill_activation_scale(dispatch, activation_scale_log2); diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-q4k-packed-tail.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-q4k-packed-tail.h index 785e032ca62c..990b9eb79ad8 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-q4k-packed-tail.h +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-q4k-packed-tail.h @@ -93,10 +93,14 @@ inline bool packed_layout_wanted(PackedLayoutSwitch value, const Value & weight) return weight_in_file_mapping(weight) && !decode_reads_stored(weight); } -inline bool q4k_packed_tail_enabled(const Value & weight) { +inline PackedLayoutSwitch q4k_packed_tail_switch() { static const PackedLayoutSwitch value = packed_layout_switch("GGML_HRX_Q4K_PACKED_TAIL", "q4_k batches off the 256 grid"); - return packed_layout_wanted(value, weight); + return value; +} + +inline bool q4k_packed_tail_enabled(const Value & weight) { + return packed_layout_wanted(q4k_packed_tail_switch(), weight); } // The layout's own preconditions, and a weight that may be packed at all. A packed copy is made of a weight once; a @@ -131,6 +135,21 @@ inline bool q6k_packed_rows_weight_ok(const Value & weight, int64_t input_size, return packed_tail_weight_ok(weight, GGML_TYPE_Q6_K, input_size, output_size); } +// engine #2137: the 256-aligned Q4_K prefill. Its f16 kernel reads only the packed rows; the 8-bit-block kernel +// (ggml_mul_mat_q5_k_iq4_xs_q8_1_x4_wmma_token256) reads Q4_K stored as well (weight format 4). For a weight that +// already has a stored device copy, that kernel shares it instead of adding a packed one. GGML_HRX_Q4K_PACKED_TAIL=0 +// (stored for every weight) selects it for every Q4_K weight, which is also how the path is reached by a test. +inline bool q4k_aligned_prefill_reads_stored(const Value & weight) { + if (weight.type != GGML_TYPE_Q4_K || weight.alias_source.value >= 0) { + return false; + } + const PackedLayoutSwitch value = q4k_packed_tail_switch(); + if (value == PackedLayoutSwitch::Stored) { + return true; + } + return value == PackedLayoutSwitch::FollowWeight && weight_in_file_mapping(weight) && decode_reads_stored(weight); +} + // The 6-to-16-token Q6_K kernel reads an i8 layout of its own. That is an extra device copy of a weight in the file // mapping whichever of the other two layouts the weight is read in elsewhere, so the kernel is left out for such a // weight unless the stored layout is forced for every weight (the previous behaviour). diff --git a/tests/hrx-backend-test.cpp b/tests/hrx-backend-test.cpp index 7976a15ffecb..7922dbe98294 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -3396,6 +3396,68 @@ static bool matmul_graph_is_supported(ggml_context * ctx, ggml_tensor * output) // stored layout on the device, so it stays stored (4 / 6) and the packed copy is not added. // GGML_HRX_Q4K_PACKED_TAIL / GGML_HRX_Q6K_PACKED_ROWS force one or the other; the row is about the default, so it does // nothing when either is set. +// engine #2137: the 256-aligned Q4_K prefill. For a file-mapped weight with no stored device copy it is the f16 kernel +// over the packed rows; for one that already has a stored copy it is the 8-bit-block kernel over the stored layout +// (weight format 4), so prefill shares decode's copy. +static void schedule_q4k_aligned_prefill_layout(bool stored_copy_exists) { + if (std::getenv("GGML_HRX_Q4K_PACKED_TAIL") != nullptr) { + std::printf(" skipped: a packed-layout switch is set in the environment\n"); + return; + } + constexpr int64_t input_size = 512; + constexpr int64_t output_size = 256; + constexpr int64_t token_count = 512; + ggml_backend_hrx_device_context device_context = {}; + device_context.name = "HRX0"; + ggml::hrx::model_buffer_from_mmap_init_type(&device_context); + ggml_init_params params = {}; + params.mem_size = 4 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, input_size, output_size); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, input_size, token_count); + const size_t byte_count = GGML_PAD(ggml_nbytes(weight), 64) + 64; + std::vector storage(byte_count + 64); + uint8_t * bytes = storage.data() + (64 - reinterpret_cast(storage.data()) % 64) % 64; + ggml_backend_buffer_t buffer = ggml::hrx::model_buffer_from_mmap(&device_context, bytes, byte_count); + REQUIRE(buffer != nullptr); + ggml_backend_buffer_set_usage(buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + ggml_tallocr allocator = ggml_tallocr_new(buffer); + REQUIRE(ggml_tallocr_alloc(&allocator, weight) == GGML_STATUS_SUCCESS); + if (stored_copy_exists) { + ggml::hrx::note_stored_device_copy(weight->data); + } + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, output); + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + size_t found = 0; + for (const ggml::hrx::Dispatch & dispatch : scheduler.plan().dispatches) { + const std::string name = kernel_name_for_id(dispatch.kernel.kernel_id); + for (const auto & binding : dispatch.bindings) { + if (binding.value != imported.graph.nodes().front().inputs[0]) { + continue; + } + ++found; + if (stored_copy_exists) { + REQUIRE(name == "loom_libs:ggml_mul_mat_q5_k_iq4_xs_q8_1_x4_wmma_token256"); + require_compile_parameter(dispatch, "ggml.mul_mat_q8_1_x4.weight_format", "4"); + REQUIRE(binding.layout == ggml::hrx::kNativeWeightLayout); + } else { + REQUIRE(name == "loom_libs:ggml_mul_mat_q4_k_f16_wmma_prefill_wave32"); + REQUIRE(binding.layout != ggml::hrx::kNativeWeightLayout); + } + } + } + REQUIRE(found == 1); + ggml_backend_buffer_free(buffer); + ggml_free(ctx); +} + static void schedule_packed_layout_follows_weight(ggml_type weight_type, bool file_mapped, int64_t token_count = 100, bool stored_copy_exists = false) { if (std::getenv("GGML_HRX_Q4K_PACKED_TAIL") != nullptr || std::getenv("GGML_HRX_Q6K_PACKED_ROWS") != nullptr) { @@ -11962,6 +12024,8 @@ int main(int argc, char ** argv) { [] { schedule_packed_layout_follows_weight(GGML_TYPE_Q6_K, false, 12); } }, // engine #2137: a file-mapped weight that already has a stored device copy (a decode kernel read it stored) is // read stored by prefill too, so it is not copied a second time; and not in the i8 layout at 6 to 16 tokens + { "q4k_aligned_prefill_layout.packed", [] { schedule_q4k_aligned_prefill_layout(false); } }, + { "q4k_aligned_prefill_layout.stored_copy", [] { schedule_q4k_aligned_prefill_layout(true); } }, { "packed_layout_follows_weight.q4_k.file_mapped.stored_copy", [] { schedule_packed_layout_follows_weight(GGML_TYPE_Q4_K, true, 100, true); } }, { "packed_layout_follows_weight.q6_k.file_mapped.stored_copy", From c100cf535a15f8c18ce637f3e3ad8972ebd84e6a Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Wed, 11 Nov 2026 16:52:09 -0400 Subject: [PATCH 4/4] hrx: one key for the stored-copy record, and a new mapping forgets the entries inside it (review of #320, engine #2137) Both sides now compute the key with stored_copy_key(): the bound tensor's host address, which is the binding's base plus offset on the copy side. The reader used view_src->data for a view, which could differ. model_buffer_from_mmap() forgets every entry inside the range it wraps, so a model mapped where an earlier one was starts with none. File-mapped weights are not evicted (the mapping is only used without a host-weight limit), so within one mapping the entries stay true. Test: hrx-backend-test stored_copy_entries_forgotten_on_remap. Co-Authored-By: Claude Opus 5.5 --- .../common/dispatch-q4k-packed-tail.h | 2 +- ggml/src/ggml-hrx/model-buffer-from-mmap.cpp | 21 ++++++++++++ ggml/src/ggml-hrx/model-buffer-from-mmap.h | 21 +++++++++--- tests/hrx-backend-test.cpp | 33 +++++++++++++++++-- 4 files changed, 68 insertions(+), 9 deletions(-) diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-q4k-packed-tail.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-q4k-packed-tail.h index 990b9eb79ad8..583031c18862 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-q4k-packed-tail.h +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-q4k-packed-tail.h @@ -76,7 +76,7 @@ inline bool decode_reads_stored(const Value & weight) { if (tensor == nullptr || !prefill_follows_decode_enabled()) { return false; } - return has_stored_device_copy(tensor->view_src != nullptr ? tensor->view_src->data : tensor->data); + return has_stored_device_copy(stored_copy_key(tensor)); } inline bool weight_in_file_mapping(const Value & weight) { diff --git a/ggml/src/ggml-hrx/model-buffer-from-mmap.cpp b/ggml/src/ggml-hrx/model-buffer-from-mmap.cpp index 159ff5608c64..33a8c1728c6d 100644 --- a/ggml/src/ggml-hrx/model-buffer-from-mmap.cpp +++ b/ggml/src/ggml-hrx/model-buffer-from-mmap.cpp @@ -19,6 +19,7 @@ #include "ggml-backend-impl.h" #include "runtime/host-memory.h" +#include #include #include #include @@ -83,6 +84,24 @@ void note_stored_device_copy(const void * host_address) { stored_copy_addresses().insert(host_address); } +const void * stored_copy_key(const ggml_tensor * tensor) { + return tensor == nullptr ? nullptr : tensor->data; +} + +void forget_stored_device_copies(const void * begin, size_t size) { + if (begin == nullptr || size == 0) { + return; + } + const uint8_t * first = static_cast(begin); + const uint8_t * last = first + size; + const std::lock_guard lock(stored_copy_mutex()); + auto & addresses = stored_copy_addresses(); + for (auto it = addresses.begin(); it != addresses.end();) { + const uint8_t * address = static_cast(*it); + it = address >= first && address < last ? addresses.erase(it) : std::next(it); + } +} + bool has_stored_device_copy(const void * host_address) { if (host_address == nullptr) { return false; @@ -110,6 +129,8 @@ ggml_backend_buffer_t model_buffer_from_mmap(ggml_backend_hrx_device_context * d ggml_backend_buffer_t buffer = ggml_backend_cpu_buffer_from_ptr(ptr, size); if (buffer != nullptr) { buffer->buft = &device->mapped_buft; + // engine #2137: entries from a mapping that was here before describe other bytes + forget_stored_device_copies(ptr, size); } return buffer; } diff --git a/ggml/src/ggml-hrx/model-buffer-from-mmap.h b/ggml/src/ggml-hrx/model-buffer-from-mmap.h index 79cb1386eba6..20c65654a000 100644 --- a/ggml/src/ggml-hrx/model-buffer-from-mmap.h +++ b/ggml/src/ggml-hrx/model-buffer-from-mmap.h @@ -32,12 +32,23 @@ bool model_buffer_from_mmap_enabled(); // weight in such a buffer costs a device copy for every layout the kernels read it in (engine #2137). bool buffer_is_file_mapped(ggml_backend_buffer_t buffer); -// engine #2137: host addresses of weights that have a device copy in the stored (GGUF) layout, recorded when the copy +// engine #2137: weights in the file mapping that have a device copy in the stored (GGUF) layout, recorded when the copy // is made. A matcher that can read a weight either stored or packed asks this before choosing, so it does not add a -// packed copy next to a stored one that a decode kernel already made. Keyed by the weight's host address, which is -// where both the matcher (tensor->data) and the copy (buffer base + offset) can see it. -void note_stored_device_copy(const void * host_address); -bool has_stored_device_copy(const void * host_address); +// packed copy next to a stored one that a decode kernel already made. +// +// The key is the host address of the bound tensor's bytes: tensor->data on the matcher side, and the binding's buffer +// base plus its offset on the copy side, which is the same address (backend-buffer-binding.cpp computes the offset as +// tensor->data - base). stored_copy_key() is that expression for a tensor; both sides use it. +// +// Lifetime: an entry describes bytes of one mapped model buffer. model_buffer_from_mmap() forgets every entry inside +// the range it wraps, so a model mapped where an earlier one was starts with none. Within one mapping the entries stay +// true: a file-mapped weight is never evicted (the mapping is only used when there is no host-weight limit, see +// model_buffer_from_mmap_enabled()). A stale entry would only make a matcher choose the stored layout, whose copy is +// then made on first use; it cannot make a kernel read bytes that were not written. +const void * stored_copy_key(const ggml_tensor * tensor); +void note_stored_device_copy(const void * host_address); +bool has_stored_device_copy(const void * host_address); +void forget_stored_device_copies(const void * begin, size_t size); // The buffer type of these buffers ("_Mapped"). It is host memory; only this device claims it. void model_buffer_from_mmap_init_type(ggml_backend_hrx_device_context * device); diff --git a/tests/hrx-backend-test.cpp b/tests/hrx-backend-test.cpp index 7922dbe98294..cc86a7135fe0 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -3396,6 +3396,32 @@ static bool matmul_graph_is_supported(ggml_context * ctx, ggml_tensor * output) // stored layout on the device, so it stays stored (4 / 6) and the packed copy is not added. // GGML_HRX_Q4K_PACKED_TAIL / GGML_HRX_Q6K_PACKED_ROWS force one or the other; the row is about the default, so it does // nothing when either is set. +// engine #2137: a stored-copy entry belongs to one mapping. A model mapped over the same addresses later (a reload in +// the same process) starts with none, so a weight there is not taken to have a copy it never had. +static void run_stored_copy_entries_forgotten_on_remap() { + ggml_backend_hrx_device_context device_context = {}; + device_context.name = "HRX0"; + ggml::hrx::model_buffer_from_mmap_init_type(&device_context); + std::vector storage(64 * 1024 + 64); + uint8_t * bytes = storage.data() + (64 - reinterpret_cast(storage.data()) % 64) % 64; + const size_t size = 64 * 1024; + + ggml_backend_buffer_t first = ggml::hrx::model_buffer_from_mmap(&device_context, bytes, size); + REQUIRE(first != nullptr); + ggml::hrx::note_stored_device_copy(bytes + 4096); + ggml::hrx::note_stored_device_copy(bytes + size + 0); // outside the mapping: must survive the remap + REQUIRE(ggml::hrx::has_stored_device_copy(bytes + 4096)); + ggml_backend_buffer_free(first); + + ggml_backend_buffer_t second = ggml::hrx::model_buffer_from_mmap(&device_context, bytes, size); + REQUIRE(second != nullptr); + REQUIRE(!ggml::hrx::has_stored_device_copy(bytes + 4096)); + REQUIRE(ggml::hrx::has_stored_device_copy(bytes + size)); + ggml::hrx::forget_stored_device_copies(bytes + size, 1); + REQUIRE(!ggml::hrx::has_stored_device_copy(bytes + size)); + ggml_backend_buffer_free(second); +} + // engine #2137: the 256-aligned Q4_K prefill. For a file-mapped weight with no stored device copy it is the f16 kernel // over the packed rows; for one that already has a stored copy it is the 8-bit-block kernel over the stored layout // (weight format 4), so prefill shares decode's copy. @@ -3426,7 +3452,7 @@ static void schedule_q4k_aligned_prefill_layout(bool stored_copy_exists) { ggml_tallocr allocator = ggml_tallocr_new(buffer); REQUIRE(ggml_tallocr_alloc(&allocator, weight) == GGML_STATUS_SUCCESS); if (stored_copy_exists) { - ggml::hrx::note_stored_device_copy(weight->data); + ggml::hrx::note_stored_device_copy(ggml::hrx::stored_copy_key(weight)); } ggml_tensor * output = ggml_mul_mat(ctx, weight, input); ggml_cgraph * graph = ggml_new_graph(ctx); @@ -3498,8 +3524,8 @@ static void schedule_packed_layout_follows_weight(ggml_type weight_type, bool fi REQUIRE(ggml_tallocr_alloc(&allocator, weight) == GGML_STATUS_SUCCESS); if (stored_copy_exists) { // engine #2137: a decode kernel already made a stored-layout device copy of this weight - ggml::hrx::note_stored_device_copy(weight->data); - REQUIRE(ggml::hrx::has_stored_device_copy(weight->data)); + ggml::hrx::note_stored_device_copy(ggml::hrx::stored_copy_key(weight)); + REQUIRE(ggml::hrx::has_stored_device_copy(ggml::hrx::stored_copy_key(weight))); } ggml_tensor * output = ggml_mul_mat(ctx, weight, input); @@ -12024,6 +12050,7 @@ int main(int argc, char ** argv) { [] { schedule_packed_layout_follows_weight(GGML_TYPE_Q6_K, false, 12); } }, // engine #2137: a file-mapped weight that already has a stored device copy (a decode kernel read it stored) is // read stored by prefill too, so it is not copied a second time; and not in the i8 layout at 6 to 16 tokens + { "stored_copy_entries_forgotten_on_remap", [] { run_stored_copy_entries_forgotten_on_remap(); } }, { "q4k_aligned_prefill_layout.packed", [] { schedule_q4k_aligned_prefill_layout(false); } }, { "q4k_aligned_prefill_layout.stored_copy", [] { schedule_q4k_aligned_prefill_layout(true); } }, { "packed_layout_follows_weight.q4_k.file_mapped.stored_copy",