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..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 @@ -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; } @@ -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 40d7f8e5e4a2..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 @@ -56,20 +56,51 @@ 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(stored_copy_key(tensor)); +} + +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); } -inline bool q4k_packed_tail_enabled(const Value & weight) { +// 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 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 @@ -104,4 +135,29 @@ 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). +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..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,9 @@ #include "ggml-backend-impl.h" #include "runtime/host-memory.h" +#include +#include +#include #include #include @@ -62,6 +65,51 @@ 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); +} + +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; + } + 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; } @@ -81,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 a9f92802c233..20c65654a000 100644 --- a/ggml/src/ggml-hrx/model-buffer-from-mmap.h +++ b/ggml/src/ggml-hrx/model-buffer-from-mmap.h @@ -32,6 +32,24 @@ 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: 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. +// +// 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/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. diff --git a/tests/hrx-backend-test.cpp b/tests/hrx-backend-test.cpp index 7aba5f604392..cc86a7135fe0 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -3396,7 +3396,96 @@ 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) { +// 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. +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(ggml::hrx::stored_copy_key(weight)); + } + 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) { std::printf(" skipped: a packed-layout switch is set in the environment\n"); return; @@ -3433,6 +3522,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(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); REQUIRE(output != nullptr); @@ -3447,7 +3541,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 +3556,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 +12048,17 @@ 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 + { "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", + [] { 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); } },