hrx: prefill keeps the stored layout for a weight that already has a stored device copy (engine #2137)

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);

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

ggml/src/ggml-hrx/model-buffer-from-mmap.cpp

@@ -19,6 +19,9 @@
 #include "ggml-backend-impl.h"
 #include "runtime/host-memory.h"
 
+#include <iterator>
+#include <mutex>
+#include <unordered_set>
 #include <cstdlib>
 #include <cstring>
 
@@ -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<const void *> & stored_copy_addresses() {
+    static std::unordered_set<const void *> addresses;
+    return addresses;
+}
+}  // namespace
+
+void note_stored_device_copy(const void * host_address) {
+    if (host_address == nullptr) {
+        return;
+    }
+    const std::lock_guard<std::mutex> 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<const uint8_t *>(begin);
+    const uint8_t *             last  = first + size;
+    const std::lock_guard<std::mutex> lock(stored_copy_mutex());
+    auto &                      addresses = stored_copy_addresses();
+    for (auto it = addresses.begin(); it != addresses.end();) {
+        const uint8_t * address = static_cast<const uint8_t *>(*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<std::mutex> 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;
 }

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

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<const uint8_t *>(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.