hrx: a Q4_K/Q6_K weight in the file mapping is read in the packed layout by default, so it has one device copy (engine #2137)

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;

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

ggml/src/ggml-hrx/dispatch_registration/common/dispatch-q4k-packed-tail.h

@@ -23,30 +23,53 @@
 // 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 <cstdlib>
 
 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) {

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         = {

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 ("<device>_Mapped"). It is host memory; only this device claims it.
 void model_buffer_from_mmap_init_type(ggml_backend_hrx_device_context * device);
 

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<uint8_t>  storage(byte_count + 64);
+    uint8_t *             bytes  = storage.data() + (64 - reinterpret_cast<uintptr_t>(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); } },