From 1c075686c876276339e513391f63f72a27d89ace Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Wed, 11 Nov 2026 09:45:05 -0400 Subject: [PATCH] hrx: a 6-to-16-token batch reads a Q6_K weight in its packed rows too, so a short prompt and its reply share one copy (engine #2137) The i8 kernel for Q6_K at 6 to 16 tokens reads the weight in a layout of its own. With the weight in the file mapping that is a third device copy: after a 12-token prompt and a reply, Qwen3-4B Q4_K_M held 35 Q6_K weights in two layouts (433 device copies, 2,856.2 MiB, for a 2,375.9 MiB model). Where the packed rows are the weight's layout (a weight in the file mapping by default, or GGML_HRX_Q6K_PACKED_ROWS=1), the matcher now declines and the tiled kernel takes the batch from the packed rows. An uploaded weight keeps the i8 kernel. Test (device-free): hrx-backend-test packed_layout_follows_weight.q6_k.{file_mapped,uploaded}.tokens12. Co-Authored-By: Claude Fable 5.2 --- .../common/dispatch-mul-mat.cpp | 8 +++++++ tests/hrx-backend-test.cpp | 21 +++++++++++++++---- 2 files changed, 25 insertions(+), 4 deletions(-) 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 108553916790..3c9dbe3b0242 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 @@ -655,6 +655,14 @@ static bool match_q6_k_i8_prepacked_dispatch(const DispatchMatchContext & contex (match.output_size > match.input_size / 2 && match.token_count > 5)) { return false; } + // engine #2137: this kernel reads the weight in a layout of its own (i8, k32), a third device copy of a weight that + // 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)) { + return false; + } const size_t block_count = static_cast(match.input_size / ggml_blck_size(GGML_TYPE_Q6_K)); size_t materialized_weight_bytes = static_cast(match.output_size) * block_count; diff --git a/tests/hrx-backend-test.cpp b/tests/hrx-backend-test.cpp index 65064797f6c7..7bd32bb30a72 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -3396,14 +3396,15 @@ 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) { +static void schedule_packed_layout_follows_weight(ggml_type weight_type, bool file_mapped, int64_t token_count = 100) { 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 + // token_count: 100 is off the 256 (Q4_K) and 128 (Q6_K) grids of the aligned prefill kernels; 12 is the batch size + // the Q6_K i8 kernel takes, which reads a layout of its own ggml_backend_hrx_device_context device_context = {}; device_context.name = "HRX0"; @@ -3461,8 +3462,16 @@ static void schedule_packed_layout_follows_weight(ggml_type weight_type, bool fi } REQUIRE(format->second == expected); } - // the case must reach a kernel that takes the format, or it asserts nothing - REQUIRE(found == 1); + if (weight_type == GGML_TYPE_Q6_K && !file_mapped && 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); + REQUIRE(kernel_name_for_id(scheduler.plan().dispatches.front().kernel.kernel_id) == + "loom_libs:ggml_mul_mat_q6_k_i8_prepacked_f16_wmma"); + } else { + // 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); @@ -11933,6 +11942,10 @@ int main(int argc, char ** argv) { [] { 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); } }, + { "packed_layout_follows_weight.q6_k.file_mapped.tokens12", + [] { 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); } }, { "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); } },