hrx: a fused dispatch may not swallow a value the caller asked to read back (engine #1078)

ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.cpp

@@ -20,6 +20,8 @@
 namespace ggml::hrx {
 namespace {
 
+static constexpr size_t kNoNode = static_cast<size_t>(-1);
+
 static constexpr size_t kOpCount = static_cast<size_t>(GGML_OP_COUNT);
 
 static const std::vector<DispatchRegistration> kEmptyRegistrations;
@@ -87,6 +89,62 @@ static bool dispatch_disabled(const DispatchRegistration & registration) {
     return false;
 }
 
+// A tensor that something outside the graph will read back after compute must still exist then. A fused match
+// that *covers* such a node (it is not the match's root) keeps the only copy of that value inside the fused
+// program, so the read-back sees the untouched buffer - measured on the HRX route as all-zero pre-RoPE q/k
+// tensors from the zha dump (engine #1078). Declining the match falls through to a narrower registration that
+// writes the node itself.
+//
+// The marker is the explicit name prefix llama_zha_dump applies (src/llama-zha-dump.cpp, "zha-readback."), not
+// GGML_TENSOR_FLAG_OUTPUT on its own: llama.cpp marks several tensors as outputs for API reasons (result_norm is
+// t_embd), and honouring the bare flag declined the final norm fusion in every ordinary run. Nothing else sets
+// this prefix, so a run without the dump is planned exactly as before. The root is exempt as well - a match
+// always writes its root's output.
+static bool tensor_is_readback_request(const ggml_tensor * tensor) {
+    constexpr const char * k_readback_prefix = "zha-readback.";
+    constexpr size_t       k_prefix_len      = sizeof("zha-readback.") - 1;
+    return tensor != nullptr && (tensor->flags & GGML_TENSOR_FLAG_OUTPUT) != 0 &&
+           std::strncmp(tensor->name, k_readback_prefix, k_prefix_len) == 0;
+}
+
+static size_t first_covered_requested_output(const DispatchMatchContext & context, const DispatchMatch & match) {
+    const std::vector<GraphNode> & nodes = context.graph.nodes();
+    for (const size_t node_index : match.covered_nodes) {
+        if (node_index == context.root_index || node_index >= nodes.size()) {
+            continue;
+        }
+        const Value * value = context.graph.values().find(nodes[node_index].output);
+        if (value != nullptr && tensor_is_readback_request(value->tensor)) {
+            return node_index;
+        }
+    }
+    return kNoNode;
+}
+
+static void dispatch_log_requested_output(const DispatchRegistration & registration, const Graph & graph,
+                                          size_t node_index) {
+    static const bool enabled = std::getenv("GGML_HRX_LOG_DISPATCH") != nullptr;
+    if (!enabled) {
+        return;
+    }
+    static std::mutex  mutex;
+    static long long   logged = 0;
+    std::lock_guard<std::mutex> lock(mutex);
+    if (logged >= 8) {
+        if (logged++ == 8) {
+            std::fprintf(stderr, "ggml_hrx dispatch: further declined registrations suppressed\n");
+        }
+        return;
+    }
+    logged++;
+    const std::vector<GraphNode> & nodes = graph.nodes();
+    const char * node_op   = node_index < nodes.size() ? ggml_op_name(nodes[node_index].op) : "?";
+    const Value * value    = node_index < nodes.size() ? graph.values().find(nodes[node_index].output) : nullptr;
+    const char * tensor_nm = (value != nullptr && value->tensor != nullptr) ? value->tensor->name : "?";
+    std::fprintf(stderr, "ggml_hrx dispatch: %s declined at node %zu (%s) - it covers tensor %s, marked as a graph output\n",
+                 registration.name != nullptr ? registration.name : "", node_index, node_op, tensor_nm);
+}
+
 static void dispatch_log_match(const DispatchRegistration & registration) {
     static const bool enabled = std::getenv("GGML_HRX_LOG_DISPATCH") != nullptr;
     if (!enabled) {
@@ -123,8 +181,26 @@ bool DispatchRegistry::match(const DispatchMatchContext & context,
             continue;
         }
         DispatchMatch candidate;
-        if (registration.matcher != nullptr && registration.matcher(context, candidate) &&
-            fusion_order_accepts(context, candidate, registration.name)) {
+        const bool    matched = registration.matcher != nullptr && registration.matcher(context, candidate);
+        const size_t  covered_output = matched ? first_covered_requested_output(context, candidate) : kNoNode;
+        if (matched && covered_output != kNoNode) {
+            dispatch_log_requested_output(registration, context.graph, covered_output);
+            if (diagnostics != nullptr) {
+                diagnostics->attempts.push_back({
+                    registration.name != nullptr ? registration.name : "",
+                    registration.root_op,
+                    registration.kind,
+                    registration.priority,
+                    registration.source,
+                    false,
+                    candidate.covered_nodes,
+                    candidate.status.errors(),
+                });
+            }
+            match.status.append(candidate.status);
+            continue;
+        }
+        if (matched && fusion_order_accepts(context, candidate, registration.name)) {
             dispatch_log_match(registration);
             if (diagnostics != nullptr) {
                 diagnostics->attempts.push_back({

src/llama-zha-dump.cpp

@@ -116,6 +116,20 @@ void llama_zha_dump_hook(ggml_tensor * t, const char * name, int il) {
     }
     // keep it alive until we read it back after graph_compute
     ggml_set_output(t);
+    // ...and say so in the name, because the output flag alone is not a request to *materialise* the value:
+    // llama.cpp marks several tensors as graph outputs for API reasons (result_norm is t_embd), and ggml-hrx
+    // would then decline a fusion for every one of them. The prefix is the explicit "someone will read this
+    // back" marker that ggml-hrx looks for; nothing else in a run carries it, so a run without the dump is
+    // planned exactly as before.
+    {
+        constexpr const char * k_readback_prefix = "zha-readback.";
+        constexpr size_t       k_prefix_len      = sizeof("zha-readback.") - 1;
+        if (strncmp(t->name, k_readback_prefix, k_prefix_len) != 0) {
+            char renamed[GGML_MAX_NAME];
+            snprintf(renamed, sizeof(renamed), "%s%s", k_readback_prefix, t->name);
+            ggml_set_name(t, renamed);
+        }
+    }
     g_pending.push_back({ t, short_name(name), il });
 }