diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat.cpp index 145f7d347508..983b8cb8e068 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat.cpp @@ -734,11 +734,11 @@ static bool match_mul_mat_swiglu_dispatch(const DispatchMatchContext & context, return format; } if (format == CommonMulMatWeightFormat::Q4K && - q4k_packed_tail_weight_ok(weight, match.input_size, match.output_size) && q4k_packed_tail_enabled()) { + q4k_packed_tail_weight_ok(weight, match.input_size, match.output_size) && q4k_packed_tail_enabled(weight)) { return CommonMulMatWeightFormat::Q4KRow64; } if (format == CommonMulMatWeightFormat::Q6K && - q6k_packed_rows_weight_ok(weight, match.input_size, match.output_size) && q6k_packed_rows_enabled()) { + q6k_packed_rows_weight_ok(weight, match.input_size, match.output_size) && q6k_packed_rows_enabled(weight)) { return CommonMulMatWeightFormat::Q6KRow64; } return format; 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 5233efe375da..108553916790 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 @@ -186,7 +186,7 @@ static bool q4k_packed_tail_applies(const DispatchMatchContext & context, int64_t output_size) { return format == CommonMulMatWeightFormat::Q4K && token_count > 5 && token_count <= kMulMatPrefillMaxTokens && q4k_packed_tail_weight_ok(weight, input_size, output_size) && - !weight_has_native_layout_reader(context.graph, weight) && q4k_packed_tail_enabled(); + !weight_has_native_layout_reader(context.graph, weight) && q4k_packed_tail_enabled(weight); } static bool q6k_packed_tail_applies(const DispatchMatchContext & context, @@ -197,7 +197,7 @@ static bool q6k_packed_tail_applies(const DispatchMatchContext & context, int64_t output_size) { return format == CommonMulMatWeightFormat::Q6K && token_count > 5 && token_count <= kMulMatPrefillMaxTokens && q6k_packed_rows_weight_ok(weight, input_size, output_size) && - !weight_has_native_layout_reader(context.graph, weight) && q6k_packed_rows_enabled(); + !weight_has_native_layout_reader(context.graph, weight) && q6k_packed_rows_enabled(weight); } static bool has_normalized_rope_consumer(const Graph & graph, const Value & value) { @@ -751,7 +751,8 @@ static bool match_q6_k_prefill_wave32_dispatch(const DispatchMatchContext & cont const bool from_packed = q6k_packed_prefill_enabled() && match.token_count % 256 == 0 && match.output_size % 128 == 0; const bool packed_rows = 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(); + !weight_has_native_layout_reader(context.graph, *match.weight) && + q6k_packed_rows_enabled(*match.weight); return build_q6_k_prefill_wave32_dispatch(match, activation.value, activation.length, activation_scale_log2, packed_rows ? kMulMatQ6KF16WmmaPrefillWave32Kernel : from_packed ? kMulMatQ6KF16WmmaPrefillFromPackedKernel 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 a26eece88b9c..40d7f8e5e4a2 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 @@ -23,14 +23,17 @@ // each layout HRX reads is a device copy of its own, so a prompt with a remainder left two copies of each Q4_K weight // (Qwen3-4B Q4_K_M, 1,700-token prompt: 4,032.5 MiB of copies against a 2,375.9 MiB model). // -// With GGML_HRX_Q4K_PACKED_TAIL=1 the tiled kernels read the packed layout too (weight format 44). Opt-in until PPL -// and tok/s are measured against the stored path. It only saves memory when the stored layout costs a copy, that is -// with #304's mapping; with GGML_HRX_MODEL_BUFFER_COPY=1 it adds packed copies instead. +// The tiled kernels can read the packed layout too (weight format 44). That saves a copy exactly when the stored +// layout costs one, which is when the weight is in #304's file mapping; for a weight that was uploaded +// (GGML_HRX_MODEL_BUFFER_COPY=1, a host-weight limit, or a load without mmap) the stored layout is already on the +// device and the packed one would be the extra copy. So the default follows the weight: packed for a weight in the +// file mapping, stored otherwise. GGML_HRX_Q4K_PACKED_TAIL=1 forces packed for every weight, =0 forces stored. #include "ggml-backend.h" #include "ggml-impl.h" #include "ggml.h" #include "graph/value-map.h" +#include "model-buffer-from-mmap.h" #include @@ -38,15 +41,35 @@ namespace ggml::hrx { // Says which path it took, once, at ERROR: the CLI's log sink drops WARN and below, and this selection changes a // compile parameter inside existing registrations, so GGML_HRX_LOG_DISPATCH cannot show it. -inline bool q4k_packed_tail_enabled() { - static const bool enabled = [] { - const char * env = std::getenv("GGML_HRX_Q4K_PACKED_TAIL"); - const bool on = env != nullptr && env[0] == '1'; - GGML_LOG_ERROR("%s: engine #2137 q4_k batches off the 256 grid read the %s layout (GGML_HRX_Q4K_PACKED_TAIL=%s)\n", - __func__, on ? "PACKED" : "stored", on ? "1" : "unset"); - return on; - }(); - return enabled; +enum class PackedLayoutSwitch { FollowWeight, Stored, Packed }; + +inline PackedLayoutSwitch packed_layout_switch(const char * name, const char * what) { + const char * env = std::getenv(name); + const PackedLayoutSwitch value = env == nullptr || env[0] == '\0' ? PackedLayoutSwitch::FollowWeight : + env[0] == '0' ? PackedLayoutSwitch::Stored : + PackedLayoutSwitch::Packed; + GGML_LOG_ERROR("%s: engine #2137 %s: %s (%s=%s)\n", __func__, what, + value == PackedLayoutSwitch::Packed ? "PACKED layout for every weight" : + value == PackedLayoutSwitch::Stored ? "stored layout for every weight" : + "PACKED layout for a weight in the file mapping, stored otherwise", + name, env == nullptr || env[0] == '\0' ? "unset" : env); + 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; + } + 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) { + 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); } // The layout's own preconditions, and a weight that may be packed at all. A packed copy is made of a weight once; a @@ -67,19 +90,14 @@ inline bool q4k_packed_tail_weight_ok(const Value & weight, int64_t input_size, return packed_tail_weight_ok(weight, GGML_TYPE_Q4_K, input_size, output_size); } -// Q6_K (GGML_HRX_Q6K_PACKED_ROWS=1): prefill reads the packed row layout decode reads, at every batch size. The tiled -// kernels take it as weight format 46 and the 128-aligned f16 prefill kernel as weight_layout 2. Stored stays the -// default. (GGML_HRX_Q6K_PACKED_PREFILL is an older switch on the aligned kernel alone; it binds the stored layout and -// saves no copy.) -inline bool q6k_packed_rows_enabled() { - static const bool enabled = [] { - const char * env = std::getenv("GGML_HRX_Q6K_PACKED_ROWS"); - const bool on = env != nullptr && env[0] == '1'; - GGML_LOG_ERROR("%s: engine #2137 q6_k prefill reads the %s layout (GGML_HRX_Q6K_PACKED_ROWS=%s)\n", __func__, - on ? "PACKED" : "stored", on ? "1" : "unset"); - return on; - }(); - return enabled; +// Q6_K: prefill reads the packed row layout decode reads, at every batch size. The tiled kernels take it as weight +// format 46 and the 128-aligned f16 prefill kernel as weight_layout 2. Same default as Q4_K above: packed for a weight +// in the file mapping, stored otherwise; GGML_HRX_Q6K_PACKED_ROWS=1 forces packed, =0 forces stored. +// (GGML_HRX_Q6K_PACKED_PREFILL is an older switch on the aligned kernel alone; it binds the stored layout and saves no +// copy.) +inline bool q6k_packed_rows_enabled(const Value & weight) { + static const PackedLayoutSwitch value = packed_layout_switch("GGML_HRX_Q6K_PACKED_ROWS", "q6_k prefill"); + return packed_layout_wanted(value, weight); } inline bool q6k_packed_rows_weight_ok(const Value & weight, int64_t input_size, int64_t output_size) { diff --git a/ggml/src/ggml-hrx/model-buffer-from-mmap.cpp b/ggml/src/ggml-hrx/model-buffer-from-mmap.cpp index b2d535e6fd35..3a7426824ef6 100644 --- a/ggml/src/ggml-hrx/model-buffer-from-mmap.cpp +++ b/ggml/src/ggml-hrx/model-buffer-from-mmap.cpp @@ -62,6 +62,10 @@ bool model_buffer_from_mmap_enabled() { return HostWeightCache::residency_limit() == 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; +} + void model_buffer_from_mmap_init_type(ggml_backend_hrx_device_context * device) { device->mapped_buft_context = { device, device->name + "_Mapped", true }; device->mapped_buft = { diff --git a/ggml/src/ggml-hrx/model-buffer-from-mmap.h b/ggml/src/ggml-hrx/model-buffer-from-mmap.h index 3e075be1e091..a9f92802c233 100644 --- a/ggml/src/ggml-hrx/model-buffer-from-mmap.h +++ b/ggml/src/ggml-hrx/model-buffer-from-mmap.h @@ -28,6 +28,10 @@ namespace ggml::hrx { // restores the old upload, and so does a residency limit (GGML_HRX_HOST_WEIGHT_LIMIT_MB). bool model_buffer_from_mmap_enabled(); +// True for a buffer of the mapped type, i.e. weights that are host pages of the model file and not a device upload. A +// 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); + // 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 dee609dcab95..14d8f00f0eff 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -18,6 +18,7 @@ #include "graph/graph.h" #include "hrx-interop-utils.h" #include "kernel-corpus/kernel-corpus.h" +#include "model-buffer-from-mmap.h" #include "runtime/command-program-executor.h" #include "runtime/graph-executor.h" #include "runtime/graph-program-cache.h" @@ -3389,6 +3390,84 @@ static bool matmul_graph_is_supported(ggml_context * ctx, ggml_tensor * output) return ggml::hrx::DispatchScheduler::can_schedule_graph(imported.graph, test_dispatch_target()); } +// engine #2137: which layout a prefill batch off the kernel's grid reads a Q4_K / Q6_K weight in follows the weight's +// own buffer. A weight in the loader's file mapping (llama.cpp #304) costs a device copy per layout, so it is read in +// the packed rows decode reads (weight format 44 / 46) and holds one copy. A weight that was uploaded already has its +// 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) { + 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; + } + constexpr int64_t input_size = 512; + constexpr int64_t output_size = 256; + constexpr int64_t token_count = 100; // off the 256 (Q4_K) and 128 (Q6_K) grids of the aligned prefill kernels + + 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, weight_type, input_size, output_size); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, input_size, token_count); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + + // the weight's bytes, as the loader would hand them over: host memory, either wrapped as the file mapping or as a + // plain weights buffer (which is what an upload looks like to the matcher: weights, but not the mapped type) + // (ggml_backend_cpu_buffer_from_ptr requires the tensor alignment, which a vector's storage does not promise) + 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 = file_mapped ? ggml::hrx::model_buffer_from_mmap(&device_context, bytes, byte_count) : + ggml_backend_cpu_buffer_from_ptr(bytes, byte_count); + REQUIRE(buffer != nullptr); + REQUIRE(ggml::hrx::buffer_is_file_mapped(buffer) == file_mapped); + 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); + + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + 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())); + 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); + 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"); + if (format == dispatch.kernel.compile_parameters.end()) { + continue; + } + ++found; + if (format->second != expected) { + std::printf(" %s weight, %s buffer: weight_format=%s, expected %s (kernel %s)\n", + ggml_type_name(weight_type), file_mapped ? "file-mapped" : "uploaded", format->second.c_str(), + expected.c_str(), kernel_name_for_id(dispatch.kernel.kernel_id).c_str()); + } + REQUIRE(format->second == expected); + } + // the case must reach a kernel that takes the format, or it asserts nothing + REQUIRE(found == 1); + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); +} + // engine #1962: the draft head's logits step (common.mul_mat.q6_k_token1_shortlist) selects activation groups, and the // scan reads as many group ids as it is told were selected. The selection kernel writes one id per 32-wide group that // exists, so the request must never exceed input_size / 32: at 96 for every width, an input below 3072 left the tail @@ -11827,6 +11906,14 @@ int main(int argc, char ** argv) { { "llama_attention_matmul_dispatch", [] { run_llama_attention_matmul_dispatch_checks(); } }, { "quantized_value_projection_dispatch", [] { run_quantized_value_projection_dispatch_checks(); } }, // widths 256 and 2560 (Qwen3.8-Flash-Next) have fewer than 96 groups; 4096 has more + { "packed_layout_follows_weight.q4_k.file_mapped", + [] { schedule_packed_layout_follows_weight(GGML_TYPE_Q4_K, true); } }, + { "packed_layout_follows_weight.q4_k.uploaded", + [] { schedule_packed_layout_follows_weight(GGML_TYPE_Q4_K, false); } }, + { "packed_layout_follows_weight.q6_k.file_mapped", + [] { schedule_packed_layout_follows_weight(GGML_TYPE_Q6_K, true); } }, + { "packed_layout_follows_weight.q6_k.uploaded", + [] { schedule_packed_layout_follows_weight(GGML_TYPE_Q6_K, false); } }, { "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); } },