hrx: the fused conv4 prefill shares the weight copy prefill already reads (engine #2137)

ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-ssm-conv.cpp

@@ -3,6 +3,7 @@
 #include "ubatch-token-limits.h"
 
 #include "../common/dispatch-mul-mat-common.h"
+#include "../common/dispatch-q4k-packed-tail.h"
 
 #include "ggml.h"
 #include "graph/graph-matcher.h"
@@ -608,6 +609,12 @@ static bool match_mul_mat_conv4_dispatch(const DispatchMatchContext & context, D
         !common_mul_mat_uses_k16_major_f16(mat.weight_format, mat.input_size, mat.output_size, mat.token_count)) {
         return false;
     }
+    // engine #2137: this fused kernel reads Q4_K only in the packed rows. For a Q4_K weight that decode already holds
+    // stored, the plain matchers take the projection instead (the stored 8-bit-block prefill kernel, then the separate
+    // conv), so the weight keeps one device copy.
+    if (q4k_aligned_prefill_reads_stored(*mat.weight)) {
+        return false;
+    }
 
     const Value * value = mat.output;
     std::vector<const GraphNode *> layouts;
@@ -680,11 +687,25 @@ static bool match_mul_mat_conv4_dispatch(const DispatchMatchContext & context, D
     dispatch.kernel.integer_parameters.emplace("token_count", mat.token_count);
     set_compile_parameter(dispatch.kernel, "ggml.mul_mat.input_size", mat.input_size);
     set_compile_parameter(dispatch.kernel, "ggml.mul_mat.output_size", mat.output_size);
-    set_compile_parameter(dispatch.kernel, "ggml.mul_mat.weight_format", mat.weight->type == GGML_TYPE_Q4_K ? 4 : 6);
+    // engine #2137: Q6_K in the packed rows (format 46) where prefill reads this weight packed elsewhere, so the fused
+    // kernel shares that copy instead of adding a stored one. Not for a weight a GET_ROWS also reads stored.
+    bool q6_packed = mat.weight->type == GGML_TYPE_Q6_K &&
+                     q6k_packed_rows_weight_ok(*mat.weight, mat.input_size, mat.output_size) &&
+                     q6k_packed_rows_enabled(*mat.weight);
+    for (const GraphNode * consumer : context.graph.index().consumers(mat.weight->id)) {
+        if (consumer != nullptr && consumer->op == GGML_OP_GET_ROWS) {
+            q6_packed = false;
+        }
+    }
+    const int64_t weight_format = mat.weight->type == GGML_TYPE_Q4_K ? 4 : q6_packed ? 46 : 6;
+    set_compile_parameter(dispatch.kernel, "ggml.mul_mat.weight_format", weight_format);
     dispatch.bindings.push_back(activation);
     if (mat.weight->type == GGML_TYPE_Q4_K) {
         dispatch.bindings.push_back({mat.weight->id, 0, mat.weight->byte_count, kQ4KPackedK256Row64Layout,
                                      mat.weight->type, mat.input_size, mat.output_size, mat.weight->byte_count});
+    } else if (q6_packed) {
+        dispatch.bindings.push_back({mat.weight->id, 0, mat.weight->byte_count, kQ6KPackedK256Row64ScaleRowLayout,
+                                     mat.weight->type, mat.input_size, mat.output_size, mat.weight->byte_count});
     } else {
         dispatch.bindings.push_back({mat.weight->id, 0, mat.weight->byte_count});
     }

tests/test-hrx-ops.cpp

@@ -5336,8 +5336,13 @@ static SsmConvPrefillGraph make_ssm_conv_prefill_graph(ggml_type type, int64_t k
     const auto sequence = scheduled_kernel_sequence(c.graph);
     const char * fused = gathered ? "loom_libs:ggml_mul_mat_quantized_f16_wmma_prefill_conv4_interior" :
                                     "loom_libs:ggml_mul_mat_quantized_f16_wmma_prefill_conv4";
-    REQUIRE((std::find(sequence.begin(), sequence.end(), fused) != sequence.end()) == !expose);
-    if (gathered && !expose) {
+    // engine #2137: the fused kernel reads Q4_K only packed, so GGML_HRX_Q4K_PACKED_TAIL=0 (stored for every Q4_K
+    // weight) sends the projection to the stored kernels and the conv runs on its own.
+    const char * q4k_layout     = std::getenv("GGML_HRX_Q4K_PACKED_TAIL");
+    const bool   q4k_stored     = type == GGML_TYPE_Q4_K && q4k_layout != nullptr && std::string(q4k_layout) == "0";
+    const bool   expect_fused   = !expose && !q4k_stored;
+    REQUIRE((std::find(sequence.begin(), sequence.end(), fused) != sequence.end()) == expect_fused);
+    if (gathered && expect_fused) {
         REQUIRE(std::find(sequence.begin(), sequence.end(),
                           "loom_libs:llm_ssm_conv_dconv4_silu_prefill_finish_f32") != sequence.end());
     }