hrx: a tied weight's 2-to-8-token batch stays on the kquant kernel (engine #2137)

ggml/src/ggml-hrx/dispatch_registration/common/dispatch-kquant-decode.cpp

@@ -171,10 +171,18 @@ bool has_q8_alternate(const DispatchMatchContext & context, const Value & input,
 // 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) {
+// Not for a weight a GET_ROWS also reads stored (a tied token embedding, engine #1316): prefill cannot pack that one
+// (dispatch-mul-mat.cpp, weight_has_native_layout_reader), so it reads the same stored copy these kernels read, and
+// declining would only hand the batch to a slower kernel. Same clause as the conv4 matcher in dispatch-ssm-conv.cpp.
+bool few_token_batch_reads_packed(const Graph & graph, const CommonMulMatMatch & match) {
     if (match.token_count <= 1 || match.weight == nullptr) {
         return false;
     }
+    for (const GraphNode * consumer : graph.index().consumers(match.weight->id)) {
+        if (consumer != nullptr && consumer->op == GGML_OP_GET_ROWS) {
+            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);
@@ -203,7 +211,7 @@ bool match_kquant_swiglu_decode(const DispatchMatchContext & context, DispatchMa
         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) ||
-        few_token_batch_reads_packed(root)) {
+        few_token_batch_reads_packed(graph, root)) {
         return false;
     }
 
@@ -228,7 +236,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 || few_token_batch_reads_packed(other)) {
+        other.token_count != root.token_count || few_token_batch_reads_packed(graph, other)) {
         return false;
     }
     const Value * output = common_graph_value(graph, glu->output);
@@ -282,7 +290,7 @@ bool match_kquant_mul_mat_decode_impl(const DispatchMatchContext & context, Disp
         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) ||
-        few_token_batch_reads_packed(root)) {
+        few_token_batch_reads_packed(graph, root)) {
         return false;
     }
 

tests/hrx-backend-test.cpp

@@ -3571,6 +3571,75 @@ static void schedule_packed_layout_follows_weight(ggml_type weight_type, bool fi
     ggml_free(ctx);
 }
 
+// engine #2137: a weight a GET_ROWS also reads (a tied token embedding) is read stored by every kernel, because the
+// gather needs the stored layout and a program holds one layout per value (engine #1316). A 2-to-8-token batch on it
+// stays on the kquant few-token kernel, which reads that same stored copy; handing it to the packed prefill kernel
+// would only pick the stored tiled path, which is slower.
+static void schedule_tied_weight_few_tokens_keeps_kquant(ggml_type weight_type, int64_t token_count) {
+    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;
+
+    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);
+    ggml_tensor * ids    = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count);
+    REQUIRE(weight != nullptr);
+    REQUIRE(input != nullptr);
+    REQUIRE(ids != nullptr);
+
+    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 = ggml::hrx::model_buffer_from_mmap(&device_context, bytes, byte_count);
+    REQUIRE(buffer != nullptr);
+    REQUIRE(ggml::hrx::buffer_is_file_mapped(buffer));
+    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 * gathered = ggml_get_rows(ctx, weight, ids);
+    ggml_tensor * output   = ggml_mul_mat(ctx, weight, input);
+    REQUIRE(gathered != nullptr);
+    REQUIRE(output != nullptr);
+    ggml_cgraph * graph = ggml_new_graph(ctx);
+    REQUIRE(graph != nullptr);
+    ggml_build_forward_expand(graph, gathered);
+    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());
+
+    size_t kquant = 0;
+    for (const ggml::hrx::Dispatch & dispatch : scheduler.plan().dispatches) {
+        const std::string name = kernel_name_for_id(dispatch.kernel.kernel_id);
+        if (name == "loom_libs:ggml_kquant_mul_mat_decode_tokens_f32") {
+            ++kquant;
+        } else {
+            std::printf("       %s, %lld tokens: dispatch %s\n", ggml_type_name(weight_type),
+                        static_cast<long long>(token_count), name.c_str());
+        }
+    }
+    REQUIRE(kquant == 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
@@ -12054,6 +12123,12 @@ int main(int argc, char ** argv) {
           [] { 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); } },
+        // ...but not for a weight a GET_ROWS also reads (a tied embedding): prefill cannot pack it either. (Q6_K: the
+        // i8 kernel takes a tied weight at 6 to 16 tokens before this matcher, dispatch-mul-mat.cpp)
+        { "tied_weight_few_tokens_keeps_kquant.q4_k.tokens6",
+          [] { schedule_tied_weight_few_tokens_keeps_kquant(GGML_TYPE_Q4_K, 6); } },
+        { "tied_weight_few_tokens_keeps_kquant.q4_k.tokens8",
+          [] { schedule_tied_weight_few_tokens_keeps_kquant(GGML_TYPE_Q4_K, 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(); } },