diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-kquant-decode.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-kquant-decode.cpp index 700ceed9d0f6..10da563d8e2b 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-kquant-decode.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-kquant-decode.cpp @@ -25,6 +25,7 @@ #include "dispatch-kquant-decode.h" #include "dispatch-mul-mat-common.h" +#include "dispatch-q4k-packed-tail.h" #include "dispatch/ternary-q4-0.h" #include "graph/graph-matcher.h" @@ -164,6 +165,27 @@ bool has_q8_alternate(const DispatchMatchContext & context, const Value & input, return find_alternate_value(context.graph, context.plan, input.id, GGML_TYPE_Q8_1, bytes) != nullptr; } +// engine #2137: a 2-to-8-token batch (a short prompt, or a speculative verify batch) where prefill reads this Q4_K / +// Q6_K weight in the packed rows. These kernels read the stored layout, so taking the batch would add a second device +// copy of the weight; the packed prefill kernels take it instead. A single token, and a weight prefill reads stored (for +// example one decode already holds stored, dispatch-q4k-packed-tail.h), stay here. +// The upper end is kMaxTokens in this file: the callers reject larger batches before asking. The predicate would still +// be correct above it, because prefill reads these weights packed there too. +bool few_token_batch_reads_packed(const CommonMulMatMatch & match) { + if (match.token_count <= 1 || match.weight == nullptr) { + return false; + } + if (match.weight->type == GGML_TYPE_Q4_K) { + return q4k_packed_tail_weight_ok(*match.weight, match.input_size, match.output_size) && + q4k_packed_tail_enabled(*match.weight); + } + if (match.weight->type == GGML_TYPE_Q6_K) { + return q6k_packed_rows_weight_ok(*match.weight, match.input_size, match.output_size) && + q6k_packed_rows_enabled(*match.weight); + } + return false; +} + bool uncovered(const DispatchMatchContext & context, const GraphNode * node) { size_t index = 0; return context.graph.index().node_index(node, index) && index < context.covered_nodes.size() && @@ -180,7 +202,8 @@ bool match_kquant_swiglu_decode(const DispatchMatchContext & context, DispatchMa if (!root.matched() || root.token_count < 1 || root.token_count > kMaxTokens || !kquant_supported(root) || root.input_size % 256 != 0 || root.input_size > kMaxInputSize || root.output_size > kMaxOutputSize || !value_ready(context, root.input->id) || - has_q8_alternate(context, *root.input, root.input_size, root.token_count)) { + has_q8_alternate(context, *root.input, root.input_size, root.token_count) || + few_token_batch_reads_packed(root)) { return false; } @@ -205,7 +228,7 @@ bool match_kquant_swiglu_decode(const DispatchMatchContext & context, DispatchMa const CommonMulMatMatch other = match_few_token_mul_mat(graph, peer, kKQuantSwiGLUDecodeKernel); if (!other.matched() || !kquant_supported(other) || other.input->id != root.input->id || other.input_size != root.input_size || other.output_size != root.output_size || - other.token_count != root.token_count) { + other.token_count != root.token_count || few_token_batch_reads_packed(other)) { return false; } const Value * output = common_graph_value(graph, glu->output); @@ -258,7 +281,8 @@ bool match_kquant_mul_mat_decode_impl(const DispatchMatchContext & context, Disp if (!root.matched() || root.token_count < 1 || root.token_count > kMaxTokens || !kquant_supported(root) || root.input_size % 256 != 0 || root.input_size > kMaxInputSize || root.output_size > kMaxOutputSize || !value_ready(context, root.input->id) || - has_q8_alternate(context, *root.input, root.input_size, root.token_count)) { + has_q8_alternate(context, *root.input, root.input_size, root.token_count) || + few_token_batch_reads_packed(root)) { return false; } diff --git a/tests/hrx-backend-test.cpp b/tests/hrx-backend-test.cpp index cc86a7135fe0..763587597c4a 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -12048,6 +12048,12 @@ int main(int argc, char ** argv) { [] { 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); } }, + // engine #2137: 6 to 8 tokens is also the kquant few-token kernels' range, which read the stored layout; for a + // file-mapped weight prefill reads packed, the packed prefill kernel takes the batch + { "packed_layout_follows_weight.q4_k.file_mapped.tokens6", + [] { schedule_packed_layout_follows_weight(GGML_TYPE_Q4_K, true, 6); } }, + { "packed_layout_follows_weight.q6_k.file_mapped.tokens8", + [] { schedule_packed_layout_follows_weight(GGML_TYPE_Q6_K, true, 8); } }, // engine #2137: a file-mapped weight that already has a stored device copy (a decode kernel read it stored) is // read stored by prefill too, so it is not copied a second time; and not in the i8 layout at 6 to 16 tokens { "stored_copy_entries_forgotten_on_remap", [] { run_stored_copy_entries_forgotten_on_remap(); } },