hrx: Q4_K batches off the 256 grid can read the packed layout, so a weight keeps one device copy (engine #2137, opt-in)

ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat.cpp

@@ -2,6 +2,7 @@
 #include "dispatch-i4-prefill-opt-in.h"
 
 #include "dispatch-mul-mat-common.h"
+#include "dispatch-q4k-packed-tail.h"
 #include "ggml.h"
 #include "graph/graph-matcher.h"
 #include "kernel-corpus/kernel-corpus-catalog-verify.h"
@@ -726,6 +727,13 @@ static bool match_mul_mat_swiglu_dispatch(const DispatchMatchContext & context,
     if (use_wmma && match.token_count < 2) {
         return false;
     }
+    // engine #2137 (dispatch-q4k-packed-tail.h): the WMMA body reading both weights in the packed row layout
+    const bool packed_tail = use_wmma && match.gate_format == CommonMulMatWeightFormat::Q4K &&
+                             match.up_format == CommonMulMatWeightFormat::Q4K &&
+                             q4k_packed_tail_weight_ok(*match.gate_weight, match.input_size, match.output_size) &&
+                             q4k_packed_tail_weight_ok(*match.up_weight, match.input_size, match.output_size) &&
+                             q4k_packed_tail_enabled();
+    const bool row_layout  = pack_q4 || packed_tail;
     dispatch.kernel = make_kernel_specialization(publish_q8 ? kMulMatSwiGLUQ4Q8OutputKernel :
         use_q8 ? kMulMatSwiGLUQ4Q8LowTokenDotKernel :
         (use_direct_dot ? kMulMatSwiGLUF32F32LowTokenDotKernel : kMulMatSwiGLUF32F32WmmaKernel));
@@ -738,16 +746,16 @@ static bool match_mul_mat_swiglu_dispatch(const DispatchMatchContext & context,
                                                common_to_config_value(match.output_size));
     dispatch.kernel.compile_parameters.emplace(
         "ggml.mul_mat_swiglu.gate_weight_format",
-        common_to_config_value(common_mul_mat_format_config_value(pack_q4 ? CommonMulMatWeightFormat::Q4KRow64 :
+        common_to_config_value(common_mul_mat_format_config_value(row_layout ? CommonMulMatWeightFormat::Q4KRow64 :
                                                                            match.gate_format)));
     dispatch.kernel.compile_parameters.emplace(
         "ggml.mul_mat_swiglu.up_weight_format",
-        common_to_config_value(common_mul_mat_format_config_value(pack_q4 ? CommonMulMatWeightFormat::Q4KRow64 :
+        common_to_config_value(common_mul_mat_format_config_value(row_layout ? CommonMulMatWeightFormat::Q4KRow64 :
                                                                            match.up_format)));
     dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_swiglu.op",
                                                common_to_config_value(static_cast<int64_t>(binary_kind_config_value(match.op))));
     dispatch.bindings.push_back(activation);
-    if (pack_q4) {
+    if (row_layout) {
         for (const Value * weight : { match.gate_weight, match.up_weight }) {
             dispatch.bindings.push_back({ weight->id, 0, weight->byte_count, kQ4KPackedK256Row64Layout, weight->type,
                                           match.input_size, match.output_size, weight->byte_count });

ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat.cpp

@@ -7,6 +7,7 @@
 #include <cstdlib>
 
 #include "dispatch-mul-mat-common.h"
+#include "dispatch-q4k-packed-tail.h"
 #include "dispatch-token-chunks.h"
 #include "dispatch-mul-mat-iq3-xxs.h"
 #include "dispatch-mul-mat-tail.h"
@@ -175,6 +176,19 @@ static bool weight_has_native_layout_reader(const Graph & graph, const Value & v
     return false;
 }
 
+// engine #2137 (dispatch-q4k-packed-tail.h): the tiled path (more than 5 tokens, f32 activation) reading the packed
+// row layout, when no other node of the graph reads the stored one (weight_has_native_layout_reader).
+static bool q4k_packed_tail_applies(const DispatchMatchContext & context,
+                                    const Value &                weight,
+                                    CommonMulMatWeightFormat     format,
+                                    int64_t                      token_count,
+                                    int64_t                      input_size,
+                                    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();
+}
+
 static bool has_normalized_rope_consumer(const Graph & graph, const Value & value) {
     if (!graph.has_index()) {
         return false;
@@ -1151,6 +1165,12 @@ static bool build_mul_mat_dispatch(const DispatchMatchContext & context,
             return false;
         }
     }
+    // engine #2137: not with the split-k variant, which has not been run on the packed layout
+    if (!pack_q4 && !pack_q6 && !split_k &&
+        q4k_packed_tail_applies(context, *match.weight, match.weight_format, match.token_count, match.input_size,
+                                match.output_size)) {
+        pack_q4 = true;
+    }
     const CommonMulMatWeightFormat format = pack_q4 ? CommonMulMatWeightFormat::Q4KRow64 :
                                            pack_q6 ? CommonMulMatWeightFormat::Q6KRow64 : match.weight_format;
     dispatch.kernel.compile_parameters.emplace(
@@ -1270,6 +1290,11 @@ static bool match_mul_mat_postops_dispatch(const DispatchMatchContext & context,
     if (match.requires_q8_activation && !pack_q4 && !pack_q6) {
         return false;
     }
+    if (!pack_q4 && !pack_q6 &&
+        q4k_packed_tail_applies(context, *match.weight, match.weight_format, match.token_count, match.input_size,
+                                match.output_size)) {
+        pack_q4 = true;  // engine #2137
+    }
     const CommonMulMatWeightFormat format = pack_q4 ? CommonMulMatWeightFormat::Q4KRow64 :
                                            pack_q6 ? CommonMulMatWeightFormat::Q6KRow64 : match.weight_format;
     dispatch.kernel.compile_parameters.emplace(

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

@@ -0,0 +1,67 @@
+// Copyright 2026 bong-water-water-bong
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+// engine #2137: one device copy of a Q4_K weight for every batch size.
+//
+// Decode and the 256-aligned prefill kernels read a Q4_K weight in HRX's packed row layout. The generic tiled kernels
+// (plain matmul, the post-op chains, the fused gate/up pair) take every other batch size, a prompt's last, unaligned
+// batch among them, and read the stored layout. Since llama.cpp #304 the weights stay in the loader's file mapping and
+// 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.
+
+#include "ggml-backend.h"
+#include "ggml-impl.h"
+#include "ggml.h"
+#include "graph/value-map.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;
+}
+
+// The layout's own preconditions, and a weight that may be packed at all. A packed copy is made of a weight once; a
+// tensor in an ordinary buffer can be rewritten between computes, and the executor refuses to read one in a packed
+// layout ("its buffer is not flagged as weights"), which would fail a graph the stored layout runs (measured on
+// test-backend-ops' q4_K large-activation cases at 22 tokens).
+inline bool q4k_packed_tail_weight_ok(const Value & weight, int64_t input_size, int64_t output_size) {
+    if (weight.type != GGML_TYPE_Q4_K || input_size % 256 != 0 || output_size % 64 != 0 ||
+        weight.alias_source.value >= 0) {
+        return false;
+    }
+    const ggml_tensor *   tensor = weight.tensor;
+    ggml_backend_buffer_t buffer =
+        tensor == nullptr ? nullptr : tensor->view_src != nullptr ? tensor->view_src->buffer : tensor->buffer;
+    return buffer != nullptr && ggml_backend_buffer_get_usage(buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS;
+}
+
+}  // namespace ggml::hrx

ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_f32_f32_wmma_core.loom

@@ -123,7 +123,19 @@ func.def inline @ggml_mul_mat_f32_f32_wmma_tiled(%weight_format: index, %token_c
         %valid_weight = scalar.andi %valid_channel, %valid_k : i1
         %weight_values = scf.if %valid_weight -> (vector<4xf16>) {
           %row_byte_base = index.scale %channel, %weight_row_bytes : index, offset -> offset
-          %decoded = func.call @ggml_dequant_f16_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %weight_noalias, %row_byte_base, %bounded_input_size, %quant_block, %quant_group, %load_packet, %weight_k) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf16>)
+          // engine #2137: format 44 is HRX's packed row layout of Q4_K (rows in groups of 64, the layout decode and the
+          // 256-aligned prefill kernel read). Reading it here lets a batch of any size share that one device copy
+          // instead of bringing in a second copy of the weight in the stored layout.
+          %packed_row_format = index.constant 44 : index
+          %packed_rows = index.cmp eq, %weight_format, %packed_row_format : index
+          %decoded = scf.if %packed_rows -> (vector<4xf16>) {
+            %packed_f32 = func.call @ggml_q4k_native_row64_f32_vector4(%weight_noalias, %channel, %bounded_input_size, %quant_block, %quant_group, %load_packet) : (buffer, index, index, index, index, index) -> (vector<4xf32>)
+            %packed_f16 = vector.fptrunc %packed_f32 : vector<4xf32> to vector<4xf16>
+            scf.yield %packed_f16 : vector<4xf16>
+          } else {
+            %native = func.call @ggml_dequant_f16_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %weight_noalias, %row_byte_base, %bounded_input_size, %quant_block, %quant_group, %load_packet, %weight_k) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf16>)
+            scf.yield %native : vector<4xf16>
+          }
           scf.yield %decoded : vector<4xf16>
         } else {
           scf.yield %c0_f16x4 : vector<4xf16>