ggml-hrx: restore correctness after upstream interface rebase

ggml/src/ggml-hrx/dispatch/command-plan.h

@@ -33,6 +33,8 @@ struct CommandPlanCompletionCounterRequest {
 };
 
 struct CommandPlan {
+    // Initializers prepare device metadata; their writes alone do not mutate caller-owned host inputs.
+    // Caller-visible outputs and in-place updates belong in dispatches.
     std::vector<Dispatch>                            initialization_dispatches;
     std::vector<Dispatch>                            dispatches;
     std::vector<CommandPlanTransient>                transients;

ggml/src/ggml-hrx/dispatch/dispatch-scheduler.cpp

@@ -174,6 +174,9 @@ bool DispatchScheduler::schedule_graph(Graph &                       graph,
         return false;
     }
     std::vector<bool>         covered_nodes(nodes.size(), false);
+    for (size_t i = 0; i < nodes.size(); ++i) {
+        covered_nodes[i] = !nodes[i].compute;
+    }
     Status                    pending_diagnostics;
     const GraphTraversalOrder traversal = GraphTraversalOrder::build(graph);
     for (const GraphNode * node : traversal.nodes()) {
@@ -299,6 +302,9 @@ bool DispatchScheduler::supports_node(const Graph & graph, const GraphNode * nod
     if (!graph.index().node_index(node, node_index)) {
         return false;
     }
+    if (!node->compute) {
+        return true;
+    }
     const std::vector<bool> covered_nodes(graph.nodes().size(), false);
     DispatchMatch           match;
     const ValueId           next_plan_value(static_cast<int32_t>(graph.values().size()));

ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rmsnorm.cpp

@@ -299,6 +299,10 @@ static RmsNormMulRopeMatch match_rmsnorm_mul_rope_f32(const Graph & graph, const
         }
     }
 
+    if (rope_params->n_offs != 0) {
+        return {};
+    }
+
     const bool sectioned_mode = rope_params->mode == GGML_ROPE_TYPE_MROPE || rope_params->mode == GGML_ROPE_TYPE_IMROPE;
     const bool lfm_neox_mode  = rope_params->mode == GGML_ROPE_TYPE_NEOX && input->ne[0] == 64 &&
                                rope_params->n_dims == 64 && !rope_has_set_rows_consumer(graph, *rope);

ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rope-set-rows.cpp

@@ -126,8 +126,12 @@ static bool is_supported_rope_input_layout(const Value & value, size_t & span_el
 }
 
 static bool is_supported_rope_params(const RopeParams & params, int64_t head_size) {
+    const bool nonnegative_offset = params.n_offs >= 0;
+    const bool aligned_offset = params.n_offs % 4 == 0;
+    const bool fits_head = params.n_offs + params.n_dims <= head_size;
+    const bool supported_offset = nonnegative_offset && aligned_offset && fits_head;
     const bool supported_mode = params.mode == GGML_ROPE_TYPE_NORMAL || params.mode == GGML_ROPE_TYPE_NEOX;
-    return params.n_dims >= 4 && params.n_dims <= head_size && params.n_dims % 4 == 0 && supported_mode &&
+    return supported_offset && params.n_dims >= 4 && params.n_dims <= head_size && params.n_dims % 4 == 0 && supported_mode &&
            std::isfinite(params.freq_base) && params.freq_base > 0.0f && std::isfinite(params.freq_scale) &&
            params.freq_scale > 0.0f && std::isfinite(params.ext_factor) && std::isfinite(params.attn_factor);
 }
@@ -244,6 +248,7 @@ struct RopeMatch {
     int64_t              head_count         = 0;
     int64_t              head_size          = 0;
     int64_t              n_dims             = 0;
+    int64_t              n_offs             = 0;
     int64_t              input_stride1      = 0;
     int64_t              input_stride2      = 0;
     size_t               input_span         = 0;
@@ -382,6 +387,7 @@ static RopeMatch match_rope_f32(const Graph & graph, const GraphNode * node, Sta
     match.head_count         = head_count;
     match.head_size          = head_size;
     match.n_dims             = n_dims;
+    match.n_offs             = params->n_offs;
     match.input_stride1      = static_cast<int64_t>(input->nb[1] / sizeof(float));
     match.input_stride2      = static_cast<int64_t>(input->nb[2] / sizeof(float));
     match.input_span         = input_span;
@@ -456,6 +462,9 @@ static void add_rope_compile_parameters(Dispatch & dispatch,
                                         const RopeMatch & match,
                                         const char *      prefix,
                                         bool              include_input_span = true) {
+    if (std::string(prefix) == "ggml.rope_f32") {
+        dispatch.kernel.compile_parameters.emplace(std::string(prefix) + ".n_offs", to_config_value(match.n_offs));
+    }
     dispatch.kernel.compile_parameters.emplace(std::string(prefix) + ".head_size", to_config_value(match.head_size));
     dispatch.kernel.compile_parameters.emplace(std::string(prefix) + ".n_dims", to_config_value(match.n_dims));
     dispatch.kernel.compile_parameters.emplace(std::string(prefix) + ".head_count", to_config_value(match.head_count));
@@ -633,6 +642,9 @@ static Dispatch make_rope_concat_dispatch(const DispatchMatchContext & context,
 static RopeConcatMatch match_rope_concat_f32(const Graph & graph, const GraphNode * node) {
     RopeConcatMatch match;
     match.rope = match_rope_f32(graph, node);
+    if (match.rope.n_offs != 0) {
+        return {};
+    }
     if (!match.rope.matched() || match.rope.mode != GGML_ROPE_TYPE_NEOX || !graph.has_index() ||
         !graph.index().has_single_consumer(node->output)) {
         return {};
@@ -722,6 +734,9 @@ static bool find_rope_set_rows_consumer(const Graph & graph, const GraphNode & r
 static RopeSetRowsMatch match_rope_set_rows_f32(const Graph & graph, const GraphNode * node) {
     RopeSetRowsMatch match;
     match.rope = match_rope_f32(graph, node);
+    if (match.rope.n_offs != 0) {
+        return {};
+    }
     if (!match.rope.matched() || !find_rope_set_rows_consumer(graph, *node, match)) {
         return {};
     }

ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-attention-qkv.cpp

@@ -112,6 +112,9 @@ static std::string to_config_value(int64_t value) {
 }
 
 static bool is_supported_attention_rope_params(const RopeParams & params, int64_t head_size) {
+    if (params.n_offs != 0) {
+        return false;
+    }
     return params.n_dims == head_size && params.mode == GGML_ROPE_TYPE_NORMAL && std::isfinite(params.freq_base) &&
            params.freq_base > 0.0f && std::isfinite(params.freq_scale) && params.freq_scale > 0.0f &&
            params.ext_factor == 0.0f && params.attn_factor == 1.0f;

ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-gated-delta-net.cpp

@@ -78,6 +78,90 @@ static bool append_covered_node(const DispatchMatchContext & context, const Grap
     return append_covered_node_index_once(context.graph, context.covered_nodes, node, match.covered_nodes);
 }
 
+struct GatedDeltaNetNormMatch {
+    const GraphNode * scale   = nullptr;
+    const Value *     input   = nullptr;
+    const Value *     output  = nullptr;
+    float             epsilon = 0.0f;
+};
+
+static const GraphNode * gated_delta_net_norm_root(const Graph & graph, ValueId output) {
+    const GraphNode * node = graph.index().producer(output);
+    if (node == nullptr) {
+        return nullptr;
+    }
+    const bool is_scale = node->op == GGML_OP_SCALE;
+    const bool unary_input = node->inputs.size() == 1;
+    const bool scaled_norm = is_scale && unary_input;
+    return scaled_norm ? producer_with_op(graph, node->inputs[0], GGML_OP_RMS_NORM) : node;
+}
+
+static GatedDeltaNetNormMatch match_gated_delta_net_norm(const Graph & graph, const GraphNode * node) {
+    if (node == nullptr) {
+        return {};
+    }
+    const bool is_l2 = node->op == GGML_OP_L2_NORM;
+    const bool is_rms = node->op == GGML_OP_RMS_NORM;
+    const bool unary_input = node->inputs.size() == 1;
+    const bool supported_norm = (is_l2 || is_rms) && unary_input;
+    if (!supported_norm) {
+        return {};
+    }
+    const auto * params = op_params_as<RmsNormParams>(node->params);
+    if (params == nullptr) {
+        return {};
+    }
+    const bool finite_epsilon = std::isfinite(params->eps);
+    const bool positive_epsilon = params->eps > 0.0f;
+    const bool supported_epsilon = finite_epsilon && positive_epsilon;
+    if (!supported_epsilon) {
+        return {};
+    }
+    GatedDeltaNetNormMatch match;
+    match.input   = graph_value(graph, node->inputs[0]);
+    match.output  = graph_value(graph, node->output);
+    match.epsilon = params->eps;
+    const bool f32_input = is_f32(match.input);
+    const bool f32_output = is_f32(match.output);
+    const bool f32_norm = f32_input && f32_output;
+    if (!f32_norm) {
+        return {};
+    }
+    if (is_rms) {
+        if (match.output->kind != ValueKind::Transient) {
+            return {};
+        }
+        match.scale = common_find_only_consumer_with_op(graph, node->output, GGML_OP_SCALE);
+        if (match.scale == nullptr) {
+            return {};
+        }
+        const auto * scale_params = op_params_as<ScaleParams>(match.scale->params);
+        if (scale_params == nullptr) {
+            return {};
+        }
+        const float width = static_cast<float>(match.input->ne[0]);
+        const bool unit_l2_scale = scale_params->scale == 1.0f / std::sqrt(width);
+        const bool zero_bias = scale_params->bias == 0.0f;
+        const bool l2_normalization = unit_l2_scale && zero_bias;
+        if (!l2_normalization) {
+            return {};
+        }
+        // RMS(eps / width) / sqrt(width) uses sum(x*x) + eps, not the L2 clamp.
+        match.epsilon *= width;
+        match.output = graph_value(graph, match.scale->output);
+        if (!is_f32(match.output)) {
+            return {};
+        }
+        const bool internal_output = match.output->kind == ValueKind::Transient;
+        const bool single_consumer = graph.index().has_single_consumer(match.output->id);
+        const bool can_elide_output = internal_output && single_consumer;
+        if (!can_elide_output) {
+            return {};
+        }
+    }
+    return match;
+}
+
 struct GatedDeltaNetMatch {
     std::vector<const GraphNode *> covered;
     const Value *                  alpha_raw      = nullptr;
@@ -102,6 +186,7 @@ struct GatedDeltaNetMatch {
     int64_t                        sequence_count = 0;
     int64_t                        snapshot_count = 0;
     float                          l2_epsilon     = 0.0f;
+    bool                           additive_epsilon = false;
 
     bool matched() const {
         return !covered.empty() && gate != nullptr && beta != nullptr && raw_q != nullptr && raw_k != nullptr &&
@@ -136,12 +221,13 @@ struct GatedDeltaNetProjectionPairMatch {
 
 static GatedDeltaNetMatch match_gated_delta_net(const Graph & graph, const GraphNode * node) {
     GatedDeltaNetMatch match;
-    if (node == nullptr || node->op != GGML_OP_L2_NORM || node->inputs.size() != 1 || !graph.has_index()) {
+    if (!graph.has_index()) {
         return match;
     }
 
-    const Value * raw_q  = graph_value(graph, node->inputs[0]);
-    const Value * q_norm = graph_value(graph, node->output);
+    const auto q_normalization = match_gated_delta_net_norm(graph, node);
+    const Value * raw_q  = q_normalization.input;
+    const Value * q_norm = q_normalization.output;
     if (!is_f32(raw_q) || !is_f32(q_norm)) {
         return {};
     }
@@ -159,17 +245,17 @@ static GatedDeltaNetMatch match_gated_delta_net(const Graph & graph, const Graph
     const Value *     beta        = graph_value(graph, gdn->inputs[4]);
     const Value *     state       = graph_value(graph, gdn->inputs[5]);
     const Value *     gdn_output  = graph_value(graph, gdn->output);
-    const GraphNode * k_norm_node = k_norm != nullptr ? producer_with_op(graph, k_norm->id, GGML_OP_L2_NORM) : nullptr;
-    const Value *     raw_k       = k_norm_node != nullptr && k_norm_node->inputs.size() == 1 ?
-                                        graph_value(graph, k_norm_node->inputs[0]) :
-                                        nullptr;
+    const GraphNode * k_norm_node = gated_delta_net_norm_root(graph, gdn->inputs[1]);
+    const auto k_normalization = match_gated_delta_net_norm(graph, k_norm_node);
+    const Value * raw_k = k_normalization.input;
     if (!is_f32(k_norm) || !is_f32(raw_k) || !is_f32(v) || !is_f32(gate) || !is_f32(beta) || !is_f32(state) ||
         !is_f32(gdn_output)) {
         return {};
     }
-    const RmsNormParams * l2_params = op_params_as<RmsNormParams>(node->params);
-    if (l2_params == nullptr || !std::isfinite(l2_params->eps) || l2_params->eps <= 0.0f ||
-        !op_params_equivalent(GGML_OP_L2_NORM, node->params, k_norm_node->params)) {
+    const bool matching_norm_ops = node->op == k_norm_node->op;
+    const bool matching_epsilon = q_normalization.epsilon == k_normalization.epsilon;
+    const bool matching_norms = matching_norm_ops && matching_epsilon;
+    if (!matching_norms) {
         return {};
     }
 
@@ -349,6 +435,10 @@ static GatedDeltaNetMatch match_gated_delta_net(const Graph & graph, const Graph
     match.covered        = { alpha_reshape, add_bias,       softplus,   gate_mul,    gate_reshape,
                              beta_reshape,  sigmoid,        node,       k_norm_node, state_reshape,
                              gdn,           new_state_view, cache_view, cache_copy,  attention_view };
+    if (q_normalization.scale != nullptr) {
+        match.covered.push_back(q_normalization.scale);
+        match.covered.push_back(k_normalization.scale);
+    }
     match.alpha_raw      = alpha_raw;
     match.beta_raw       = beta_raw;
     match.bias           = bias;
@@ -370,7 +460,8 @@ static GatedDeltaNetMatch match_gated_delta_net(const Graph & graph, const Graph
     match.token_count    = token_count;
     match.sequence_count = sequence_count;
     match.snapshot_count = snapshot_count;
-    match.l2_epsilon     = l2_params->eps;
+    match.l2_epsilon     = q_normalization.epsilon;
+    match.additive_epsilon = q_normalization.scale != nullptr;
     return match;
 }
 
@@ -400,7 +491,7 @@ static GatedDeltaNetProjectionPairMatch match_gated_delta_net_projection_pair(co
         return {};
     }
 
-    const GraphNode *        q_norm    = producer_with_op(context.graph, gdn->inputs[0], GGML_OP_L2_NORM);
+    const GraphNode *        q_norm    = gated_delta_net_norm_root(context.graph, gdn->inputs[0]);
     const GatedDeltaNetMatch gdn_match = match_gated_delta_net(context.graph, q_norm);
     if (!gdn_match.matched() || !gdn_match.has_projection_epilogue() || gdn_match.alpha_raw->id != first_node->output) {
         return {};
@@ -664,6 +755,7 @@ static void configure_gated_delta_net_kernel(Dispatch &                 dispatch
     set_compile_parameter(dispatch.kernel, "llm.gated_delta_net.query_head_count", match.q_head_count);
     set_compile_parameter(dispatch.kernel, "llm.gated_delta_net.query_sequence_ratio", 1);
     set_float_compile_parameter(dispatch.kernel, "llm.gated_delta_net.l2_epsilon", match.l2_epsilon);
+    set_compile_parameter(dispatch.kernel, "llm.gated_delta_net.additive_epsilon", match.additive_epsilon ? 1 : 0);
     set_compile_parameter(dispatch.kernel, "llm.gated_delta_net.workgroup_size", 256);
 }
 
@@ -1058,7 +1150,7 @@ static bool match_selected_gated_delta_net_dispatch(const DispatchMatchContext &
         gdn->inputs[5] != reshape->output) {
         return false;
     }
-    const GraphNode * q_norm = producer_with_op(graph, gdn->inputs[0], GGML_OP_L2_NORM);
+    const GraphNode * q_norm = gated_delta_net_norm_root(graph, gdn->inputs[0]);
     const GatedDeltaNetMatch match = match_gated_delta_net(graph, q_norm);
     if (!match.matched() || !match.has_projection_epilogue() || match.sequence_count != 1 ||
         match.token_count < 1 || match.token_count > 5 || match.snapshot_count < match.token_count ||
@@ -1196,6 +1288,14 @@ static bool match_selected_gated_delta_net_dispatch(const DispatchMatchContext &
 }  // namespace
 
 void register_llm_gated_delta_net_dispatch(DispatchRegistryBuilder & registry) {
+    registry.add({
+        "llm.gated_delta_net.rms_scale.f32_wmma_head128",
+        GGML_OP_RMS_NORM,
+        DispatchMatchKind::Fused,
+        200,
+        DispatchSource::Llm,
+        match_gated_delta_net_dispatch,
+    });
     registry.add({
         "llm.gated_delta_net.selected_snapshot.f32_wmma_head128",
         GGML_OP_GET_ROWS,

ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen-attention-postprocess.cpp

@@ -59,6 +59,9 @@ static bool is_qwen_rms_norm_epsilon(float eps) {
 }
 
 static bool is_qwen_implicit_rope_contract(const RopeParams & params) {
+    if (params.n_offs != 0) {
+        return false;
+    }
     return params.n_dims == kQwenAttentionHeadSize && params.mode == GGML_ROPE_TYPE_NEOX &&
            std::isfinite(params.freq_base) && params.freq_base > 0.0f && std::isfinite(params.freq_scale) &&
            params.freq_scale > 0.0f && std::isfinite(params.ext_factor) && std::isfinite(params.attn_factor);

ggml/src/ggml-hrx/graph/graph-diagnostics.cpp

@@ -269,6 +269,7 @@ json node_json(const GraphNode & node) {
         { "op",      static_cast<int>(node.op)   },
         { "op_name", ggml_op_name(node.op)       },
         { "output",  node.output.value           },
+        { "compute", node.compute                },
         { "inputs",  std::move(inputs)           },
         { "params",  op_params_json(node.params) },
     };
@@ -339,7 +340,7 @@ std::string serialize_graph_snapshot_json(const Graph & graph, const std::string
         nodes.push_back(node_json(node));
     }
     json root = {
-        { "schema",   "ggml-hrx-graph-snapshot-v1" },
+        { "schema",   "ggml-hrx-graph-snapshot-v2" },
         { "uid",      uid                          },
         { "target",   target                       },
         { "storages", std::move(storages)          },
@@ -351,13 +352,13 @@ std::string serialize_graph_snapshot_json(const Graph & graph, const std::string
 
 std::string format_graph_snapshot_text(const Graph & graph, const std::string & target, uint64_t uid) {
     std::ostringstream out;
-    out << "schema=ggml-hrx-graph-snapshot-v1\n";
+    out << "schema=ggml-hrx-graph-snapshot-v2\n";
     out << "uid=" << uid << "\ntarget=" << target << "\nvalues=" << graph.values().size()
         << "\nnodes=" << graph.nodes().size() << '\n';
     for (size_t i = 0; i < graph.nodes().size(); ++i) {
         const GraphNode & node = graph.nodes()[i];
         out << "node " << i << " " << ggml_op_name(node.op) << " output=" << value_summary(graph, node.output)
-            << " inputs=[";
+            << " compute=" << node.compute << " inputs=[";
         for (size_t j = 0; j < node.inputs.size(); ++j) {
             if (j > 0) {
                 out << ", ";
@@ -384,7 +385,7 @@ GraphSnapshotLoadResult load_graph_snapshot_json(const std::string & contents) {
     GraphSnapshotLoadResult result;
     try {
         const json root = json::parse(contents);
-        if (root.value("schema", "") != "ggml-hrx-graph-snapshot-v1") {
+        if (root.value("schema", "") != "ggml-hrx-graph-snapshot-v2") {
             result.status.log("unsupported graph snapshot schema");
             return result;
         }
@@ -448,6 +449,7 @@ GraphSnapshotLoadResult load_graph_snapshot_json(const std::string & contents) {
             }
             GraphNode & node = result.graph.add_node(static_cast<ggml_op>(item.at("op").get<int>()),
                                                      ValueId(item.at("output").get<int32_t>()), std::move(inputs));
+            node.compute     = item.at("compute").get<bool>();
             node.params      = parse_op_params(item.at("params"));
         }
         result.status.append(result.graph.build_index());