diff --git a/ggml/src/ggml-hrx/dispatch/command-plan.h b/ggml/src/ggml-hrx/dispatch/command-plan.h index aa6693342086..9807df49c534 100644 --- a/ggml/src/ggml-hrx/dispatch/command-plan.h +++ b/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 initialization_dispatches; std::vector dispatches; std::vector transients; diff --git a/ggml/src/ggml-hrx/dispatch/dispatch-scheduler.cpp b/ggml/src/ggml-hrx/dispatch/dispatch-scheduler.cpp index 21990fe038fa..12d4f5fa338a 100644 --- a/ggml/src/ggml-hrx/dispatch/dispatch-scheduler.cpp +++ b/ggml/src/ggml-hrx/dispatch/dispatch-scheduler.cpp @@ -174,6 +174,9 @@ bool DispatchScheduler::schedule_graph(Graph & graph, return false; } std::vector 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 covered_nodes(graph.nodes().size(), false); DispatchMatch match; const ValueId next_plan_value(static_cast(graph.values().size())); diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rmsnorm.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rmsnorm.cpp index 741bd7c110f1..15e8965ea721 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rmsnorm.cpp +++ b/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); diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rope-set-rows.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rope-set-rows.cpp index 3b668b05bb28..a1976d98e4e5 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rope-set-rows.cpp +++ b/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(input->nb[1] / sizeof(float)); match.input_stride2 = static_cast(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 {}; } diff --git a/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-attention-qkv.cpp b/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-attention-qkv.cpp index 75256bc4ec7d..98d90b2957d3 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-attention-qkv.cpp +++ b/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; diff --git a/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-gated-delta-net.cpp b/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-gated-delta-net.cpp index 56b3888dc82f..06ed2f48b789 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-gated-delta-net.cpp +++ b/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(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(match.scale->params); + if (scale_params == nullptr) { + return {}; + } + const float width = static_cast(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 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(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, diff --git a/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen-attention-postprocess.cpp b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen-attention-postprocess.cpp index 9b0dfd1cff88..798cfcf91596 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen-attention-postprocess.cpp +++ b/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); diff --git a/ggml/src/ggml-hrx/graph/graph-diagnostics.cpp b/ggml/src/ggml-hrx/graph/graph-diagnostics.cpp index 36df4bf405b8..2ead551ffe26 100644 --- a/ggml/src/ggml-hrx/graph/graph-diagnostics.cpp +++ b/ggml/src/ggml-hrx/graph/graph-diagnostics.cpp @@ -269,6 +269,7 @@ json node_json(const GraphNode & node) { { "op", static_cast(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(item.at("op").get()), ValueId(item.at("output").get()), std::move(inputs)); + node.compute = item.at("compute").get(); node.params = parse_op_params(item.at("params")); } result.status.append(result.graph.build_index()); diff --git a/ggml/src/ggml-hrx/graph/graph.cpp b/ggml/src/ggml-hrx/graph/graph.cpp index f88f34faba02..b9b2c1124765 100644 --- a/ggml/src/ggml-hrx/graph/graph.cpp +++ b/ggml/src/ggml-hrx/graph/graph.cpp @@ -56,6 +56,9 @@ GraphIndex GraphIndex::build(const Graph & graph) { for (size_t i = 0; i < nodes.size(); ++i) { const GraphNode & node = nodes[i]; index.node_indices_.emplace(&node, i); + if (!node.compute) { + continue; + } index.producers_.emplace(node.output.value, &node); for (ValueId input : node.inputs) { index.consumers_[input.value].push_back(&node); @@ -157,6 +160,10 @@ GraphImportResult import_ggml_graph(const ggml_cgraph & graph) { result.status.log("ggml graph contains a null node"); return result; } + const bool compute = (node->flags & GGML_TENSOR_FLAG_COMPUTE) != 0; + if (!compute) { + continue; + } graph_nodes.insert(node); for (const ggml_tensor * source : node->src) { if (source != nullptr) { @@ -168,6 +175,13 @@ GraphImportResult import_ggml_graph(const ggml_cgraph & graph) { ValueMap & values = result.graph.values(); for (int i = 0; i < graph.n_nodes; ++i) { const ggml_tensor * node = graph.nodes[i]; + const bool compute = (node->flags & GGML_TENSOR_FLAG_COMPUTE) != 0; + if (!compute) { + // Keep node slots aligned with ggml for cached external bindings, but do not import inactive branch storage. + GraphNode & inactive = result.graph.add_node(node->op, ValueId{}, {}); + inactive.compute = false; + continue; + } std::vector inputs; for (const ggml_tensor * source : node->src) { if (source == nullptr) { diff --git a/ggml/src/ggml-hrx/graph/graph.h b/ggml/src/ggml-hrx/graph/graph.h index 60b26a59daa7..b3783a2178c3 100644 --- a/ggml/src/ggml-hrx/graph/graph.h +++ b/ggml/src/ggml-hrx/graph/graph.h @@ -17,6 +17,7 @@ struct ggml_tensor; namespace ggml::hrx { struct GraphNode { + bool compute = true; ggml_op op; ValueId output; std::vector inputs; diff --git a/ggml/src/ggml-hrx/graph/op-params.cpp b/ggml/src/ggml-hrx/graph/op-params.cpp index 015b6b7b77e3..bf3374f973bf 100644 --- a/ggml/src/ggml-hrx/graph/op-params.cpp +++ b/ggml/src/ggml-hrx/graph/op-params.cpp @@ -84,7 +84,8 @@ static bool rope_params_equivalent(const OpParams & lhs, const OpParams & rhs) { nearly_equal(lhs_params->ext_factor, rhs_params->ext_factor) && nearly_equal(lhs_params->attn_factor, rhs_params->attn_factor) && nearly_equal(lhs_params->beta_fast, rhs_params->beta_fast) && - nearly_equal(lhs_params->beta_slow, rhs_params->beta_slow) && lhs_params->sections == rhs_params->sections; + nearly_equal(lhs_params->beta_slow, rhs_params->beta_slow) && lhs_params->sections == rhs_params->sections && + lhs_params->n_offs == rhs_params->n_offs; } } // namespace @@ -302,6 +303,7 @@ OpParams import_op_params(const ggml_tensor & tensor) { ggml_get_op_params_i32(&tensor, 13), ggml_get_op_params_i32(&tensor, 14), }, + ggml_get_op_params_i32(&tensor, 15), }; default: return std::monostate{}; diff --git a/ggml/src/ggml-hrx/graph/op-params.h b/ggml/src/ggml-hrx/graph/op-params.h index fc6a7354f60f..73f5a9122322 100644 --- a/ggml/src/ggml-hrx/graph/op-params.h +++ b/ggml/src/ggml-hrx/graph/op-params.h @@ -109,6 +109,7 @@ struct RopeParams { float beta_fast = 0.0f; float beta_slow = 0.0f; std::array sections = {}; + int n_offs = 0; }; // clang-format off diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/gated_delta_net/scan_state_head128.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/gated_delta_net/scan_state_head128.loom index 965fa1b1a86f..397b3350d807 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/gated_delta_net/scan_state_head128.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/gated_delta_net/scan_state_head128.loom @@ -68,6 +68,8 @@ config.decl @llm.gated_delta_net.snapshot_stride : %value: index where [range(%v config.decl @llm.gated_delta_net.l2_epsilon : f32 +config.decl @llm.gated_delta_net.additive_epsilon : %value: index where [range(%value, 0, 1)] + config.decl @llm.gated_delta_net.workgroup_size : %value: index where [range(%value, 32, 1024), multiple_of(%value, 32)] template.def<@llm.gated_delta_net.f32_wmma_head128.body> device @llm_gated_delta_net_f32_wmma_head128_body(%state_inplace: i1, %publish_snapshots: i1, %projection_epilogue: i1, %snapshot_stride: index, %scale: f32, %q: buffer, %k: buffer, %v: buffer, %g: buffer, %beta: buffer, %bias: buffer, %a_scale: buffer, %state_in: buffer, %dst: buffer, %snapshot_cache: buffer, %rmsnorm_gate: i1, %rms_epsilon: f32, %rms_gate_op: index, %rms_weight: buffer, %raw_gate: buffer, %norm_output: buffer, %half_output: buffer, %select_state: i1, %state_row_count: index, %state_ids: buffer, %publish_q8: i1) { @@ -448,6 +450,8 @@ template.def<@llm.gated_delta_net.f32_wmma_head128.body> device @llm_gated_delta %f0 = scalar.constant 0.0 : f32 %norm_eps = config.get @llm.gated_delta_net.l2_epsilon : f32 %norm_eps_sq = scalar.mulf %norm_eps, %norm_eps : f32 + %additive_epsilon = config.get @llm.gated_delta_net.additive_epsilon : index + %use_additive_epsilon = index.cmp eq, %additive_epsilon, %c1 : index %fz8 = vector.constant 0.0 : vector<8xf32> %zv8 = vector.constant 0.0 : vector<8xf32> %fzh8 = vector.constant 0.0 : vector<8xf16> @@ -674,7 +678,9 @@ template.def<@llm.gated_delta_net.f32_wmma_head128.body> device @llm_gated_delta %n0_k_hi = scalar.addf %n0_k_hi46, %n0_k_hi57 : f32 %n0_k_sum = scf.select %lane_lo, %n0_k_lo, %n0_k_hi : f32 %n0_k_clamped = scalar.maxnumf %n0_k_sum, %norm_eps_sq : f32 - %n0_k_scale = scalar.rsqrtf %n0_k_clamped : f32 + %n0_k_biased = scalar.addf %n0_k_sum, %norm_eps : f32 + %n0_k_norm = scf.select %use_additive_epsilon, %n0_k_biased, %n0_k_clamped : f32 + %n0_k_scale = scalar.rsqrtf %n0_k_norm : f32 %n0_k_scale_v = vector.splat %n0_k_scale : vector<8xf32> %n0_kn = vector.mulf %n0_k_scale_v, %kvs0 : vector<8xf32> %kh0 = vector.fptrunc %n0_kn : vector<8xf32> to vector<8xf16> @@ -757,7 +763,9 @@ template.def<@llm.gated_delta_net.f32_wmma_head128.body> device @llm_gated_delta %n0_q_hi = scalar.addf %n0_q_hi46, %n0_q_hi57 : f32 %n0_q_sum = scf.select %lane_lo, %n0_q_lo, %n0_q_hi : f32 %n0_q_clamped = scalar.maxnumf %n0_q_sum, %norm_eps_sq : f32 - %n0_q_scale = scalar.rsqrtf %n0_q_clamped : f32 + %n0_q_biased = scalar.addf %n0_q_sum, %norm_eps : f32 + %n0_q_norm = scf.select %use_additive_epsilon, %n0_q_biased, %n0_q_clamped : f32 + %n0_q_scale = scalar.rsqrtf %n0_q_norm : f32 %n0_q_scale_v = vector.splat %n0_q_scale : vector<8xf32> %n0_qn = vector.mulf %n0_q_scale_v, %qvs0 : vector<8xf32> %qh0 = vector.fptrunc %n0_qn : vector<8xf32> to vector<8xf16> diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/rope/rope_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/rope/rope_f32.loom index cdb99471c0c4..78192197f529 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/rope/rope_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/rope/rope_f32.loom @@ -55,6 +55,8 @@ template.decl @ggml.rows.cache_publish.f32_vector4(%arg0: index, %arg1: i1, %arg config.decl @ggml.rope_f32.head_size : %value: index where [range(%value, 4, 1024), multiple_of(%value, 4)] +config.decl @ggml.rope_f32.n_offs : %value: index where [range(%value, 0, 1024), multiple_of(%value, 4)] + config.decl @ggml.rope_f32.n_dims : %value: index where [range(%value, 4, 1024), multiple_of(%value, 4)] config.decl @ggml.rope_f32.head_count : %value: index where [range(%value, 1, 64)] @@ -155,6 +157,7 @@ template.def<@ggml.rope_f32.scaled_body> device @ggml_rope_f32_scaled_body(%toke %token_capacity = config.get @ggml.rope_f32.token_capacity : index %head_count0 = config.get @ggml.rope_f32.head_count : index %head_size0 = config.get @ggml.rope_f32.head_size : index + %n_offs = config.get @ggml.rope_f32.n_offs : index %n_dims0 = config.get @ggml.rope_f32.n_dims : index %input_stride1_0 = config.get @ggml.rope_f32.input_stride1 : index %input_stride2_0 = config.get @ggml.rope_f32.input_stride2 : index @@ -193,7 +196,8 @@ template.def<@ggml.rope_f32.scaled_body> device @ggml_rope_f32_scaled_body(%toke scf.if %publish_rotary { %channel = index.assume %channel0 [lt(%channel0, %pair_packet_count)] : index %pair_channel = index.mul %channel, %c2 : index - %normal_channel = index.mul %channel, %c4 : index + %normal_channel0 = index.mul %channel, %c4 : index + %normal_channel = index.add %normal_channel0, %n_offs : index %positions_view = buffer.view %positions_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]xi32> %theta_view = buffer.view %theta_noalias[%c0_offset] : buffer -> view<[%half_n_dims]xf32> %freq_factors_view = buffer.view %freq_factors_noalias[%c0_offset] : buffer -> view<[%half_n_dims]xf32> @@ -206,11 +210,12 @@ template.def<@ggml.rope_f32.scaled_body> device @ggml_rope_f32_scaled_body(%toke %theta_packet = vector.load %theta_view[%pair_channel] : view<[%half_n_dims]xf32> -> vector<2xf32> %freq_factors_packet = vector.load %freq_factors_view[%pair_channel] : view<[%half_n_dims]xf32> -> vector<2xf32> scf.if %is_neox { - %paired_channel = index.add %pair_channel, %half_n_dims : index - %scaled_low, %scaled_high = template.apply<@ggml.rope_f32.rotate_neox_packet>(%position, %input_base, %pair_channel, %half_n_dims, %bounded_input_span, %mscale, %input_noalias, %theta_packet, %freq_factors_packet) : (f32, index, index, index, index, f32, buffer, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) + %rotary_channel = index.add %pair_channel, %n_offs : index + %paired_channel = index.add %rotary_channel, %half_n_dims : index + %scaled_low, %scaled_high = template.apply<@ggml.rope_f32.rotate_neox_packet>(%position, %input_base, %rotary_channel, %half_n_dims, %bounded_input_span, %mscale, %input_noalias, %theta_packet, %freq_factors_packet) : (f32, index, index, index, index, f32, buffer, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) %output_low = vector.mulf %scaled_low, %output_scale2 : vector<2xf32> %output_high = vector.mulf %scaled_high, %output_scale2 : vector<2xf32> - vector.store %output_low, %output_view[%token, %head, %pair_channel] : vector<2xf32>, view<[%bounded_token_count]x[%head_count]x[%head_size]xf32> + vector.store %output_low, %output_view[%token, %head, %rotary_channel] : vector<2xf32>, view<[%bounded_token_count]x[%head_count]x[%head_size]xf32> vector.store %output_high, %output_view[%token, %head, %paired_channel] : vector<2xf32>, view<[%bounded_token_count]x[%head_count]x[%head_size]xf32> } scf.if %is_normal { @@ -221,7 +226,11 @@ template.def<@ggml.rope_f32.scaled_body> device @ggml_rope_f32_scaled_body(%toke } scf.if %publish_copy { %channel = index.assume %channel0 [lt(%channel0, %pair_packet_count)] : index - %normal_channel = index.mul %channel, %c4 : index + // Copy packets cover the unchanged prefix first, then the suffix after RoPE. + %copy_channel0 = index.mul %channel, %c4 : index + %prefix_channel = index.sub %copy_channel0, %n_dims : index + %is_prefix = index.cmp ult, %prefix_channel, %n_offs : index + %normal_channel = scf.select %is_prefix, %prefix_channel, %copy_channel0 : index %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%head_count]x[%head_size]xf32> %input_token_offset = index.mul %token, %input_stride2 : index %input_head_offset = index.mul %head, %input_stride1 : index diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rope/rope_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rope/rope_f32.loom index 5de26f97c55a..132f3cfa916a 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rope/rope_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rope/rope_f32.loom @@ -7,6 +7,7 @@ template.decl @ggml.rope_f32.body(%token_count: index, %input_span: index, %posi amdgpu.target @ggml_rope_f32_gfx11_wave32 {subgroup_size = 32} config.decl @ggml.rope_f32.head_size : %value: index where [range(%value, 4, 1024), multiple_of(%value, 4)] +config.decl @ggml.rope_f32.n_offs : %value: index where [range(%value, 0, 1024), multiple_of(%value, 4)] config.decl @ggml.rope_f32.n_dims : %value: index where [range(%value, 4, 1024), multiple_of(%value, 4)] config.decl @ggml.rope_f32.head_count : %value: index where [range(%value, 1, 64)] config.decl @ggml.rope_f32.token_capacity : %value: index where [range(%value, 1, 2048)] diff --git a/ggml/src/ggml-hrx/runtime/command-program-executor.cpp b/ggml/src/ggml-hrx/runtime/command-program-executor.cpp index 8c30489611a2..7723246b29f9 100644 --- a/ggml/src/ggml-hrx/runtime/command-program-executor.cpp +++ b/ggml/src/ggml-hrx/runtime/command-program-executor.cpp @@ -701,6 +701,7 @@ static Dispatch build_dispatch(const ResolvedCommand & command) { struct GraphValueAccess { bool read = false; bool write = false; + bool main_write = false; bool native_read = false; bool transformed_read = false; bool layout_conflict = false; @@ -714,13 +715,14 @@ struct GraphValueAccess { static std::unordered_map collect_graph_value_access(const CommandProgram & commands) { std::unordered_map access_by_value; - auto append_command_list_access = [&](const std::vector & command_list) { + auto append_command_list_access = [&](const std::vector & command_list, bool main_commands) { for (const Command & command : command_list) { for (const CommandBinding & binding : command.bindings) { if (binding.origin != CommandBindingOrigin::GraphValue) { continue; } GraphValueAccess & access = access_by_value[binding.value.value]; + access.main_write |= main_commands && resource_access_writes(binding.access); switch (binding.access) { case ResourceAccess::Read: access.read = true; @@ -768,8 +770,8 @@ static std::unordered_map collect_graph_value_access( } } }; - append_command_list_access(commands.initialization_commands); - append_command_list_access(commands.commands); + append_command_list_access(commands.initialization_commands, false); + append_command_list_access(commands.commands, true); return access_by_value; } @@ -851,7 +853,8 @@ static CommandProgramBindings materialize_host_bindings(const CommandProgramExec staging.value = binding.value.value; staging.host_data = static_cast(binding.host_data) + binding.offset; staging.upload = access.read; - staging.download = access.write; + // Initializers may synthesize device metadata from caller inputs, but only main graph writes are observable. + staging.download = access.main_write; if (staging.upload && context.host_buffers != nullptr) { staging.source_host_buffer = context.host_buffers->find(staging.host_data, staging.length); } diff --git a/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp b/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp index fb9af040c528..666602f8fcd3 100644 --- a/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp +++ b/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp @@ -385,6 +385,19 @@ GraphProgramMatch GraphProgram::match_trusted_graph(const ggml_cgraph & current_ return result; } } + // The scheduler can retain node slots while changing which select branch computes. + for (size_t i = 0; i < graph_->nodes().size(); ++i) { + const ggml_tensor * node = current_graph.nodes[i]; + if (node == nullptr) { + result.status.log("current graph node %zu is null", i); + return result; + } + const bool compute = (node->flags & GGML_TENSOR_FLAG_COMPUTE) != 0; + if (compute != graph_->nodes()[i].compute) { + result.status.log("node %zu compute mask does not match cached HRX graph", i); + return result; + } + } if (!bind_external) { return result; } @@ -437,6 +450,9 @@ Status GraphProgram::capture_external_slots(const ggml_cgraph & graph, const Gra if (node == nullptr) { continue; } + if (!graph_->nodes()[static_cast(i)].compute) { + continue; + } if (node == binding.tensor) { slot.kind = GraphProgramExternalSlotKind::Node; slot.node_index = static_cast(i); @@ -574,6 +590,14 @@ GraphProgramMatch GraphProgram::match_current_graph(const ggml_cgraph & current_ ggml_op_name(cached_node.op), ggml_op_name(current_node->op)); return result; } + const bool compute = (current_node->flags & GGML_TENSOR_FLAG_COMPUTE) != 0; + if (compute != cached_node.compute) { + result.status.log("node %zu compute mask does not match cached HRX graph", node_index); + return result; + } + if (!compute) { + continue; + } if (!graph_node_params_match(cached_node, current_node)) { result.status.log("node %zu cached op params do not match current graph", node_index); return result; diff --git a/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_draft_programs.tg_c1.json b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_draft_programs.tg_c1.json index d6f812883b17..f12421e4bb6f 100644 --- a/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_draft_programs.tg_c1.json +++ b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_draft_programs.tg_c1.json @@ -393,6 +393,7 @@ "llm.gated_delta_net.head_count": "48", "llm.gated_delta_net.head_width": "128", "llm.gated_delta_net.l2_epsilon": "9.99999997e-07", + "llm.gated_delta_net.additive_epsilon": "0", "llm.gated_delta_net.qk_stride1": "128", "llm.gated_delta_net.qk_stride2": "10240", "llm.gated_delta_net.qk_stride3": "20480", @@ -1308,6 +1309,7 @@ "llm.gated_delta_net.head_count": "48", "llm.gated_delta_net.head_width": "128", "llm.gated_delta_net.l2_epsilon": "9.99999997e-07", + "llm.gated_delta_net.additive_epsilon": "0", "llm.gated_delta_net.qk_stride1": "128", "llm.gated_delta_net.qk_stride2": "10240", "llm.gated_delta_net.qk_stride3": "40960", @@ -2123,6 +2125,7 @@ "llm.gated_delta_net.head_count": "48", "llm.gated_delta_net.head_width": "128", "llm.gated_delta_net.l2_epsilon": "9.99999997e-07", + "llm.gated_delta_net.additive_epsilon": "0", "llm.gated_delta_net.qk_stride1": "128", "llm.gated_delta_net.qk_stride2": "10240", "llm.gated_delta_net.qk_stride3": "51200", @@ -2990,6 +2993,7 @@ "llm.gated_delta_net.head_count": "48", "llm.gated_delta_net.head_width": "128", "llm.gated_delta_net.l2_epsilon": "9.99999997e-07", + "llm.gated_delta_net.additive_epsilon": "0", "llm.gated_delta_net.qk_stride1": "128", "llm.gated_delta_net.qk_stride2": "10240", "llm.gated_delta_net.qk_stride3": "30720", @@ -3857,6 +3861,7 @@ "llm.gated_delta_net.head_count": "48", "llm.gated_delta_net.head_width": "128", "llm.gated_delta_net.l2_epsilon": "9.99999997e-07", + "llm.gated_delta_net.additive_epsilon": "0", "llm.gated_delta_net.qk_stride1": "128", "llm.gated_delta_net.qk_stride2": "10240", "llm.gated_delta_net.qk_stride3": "163840", @@ -3919,6 +3924,7 @@ "llm.gated_delta_net.head_count": "48", "llm.gated_delta_net.head_width": "128", "llm.gated_delta_net.l2_epsilon": "9.99999997e-07", + "llm.gated_delta_net.additive_epsilon": "0", "llm.gated_delta_net.qk_stride1": "128", "llm.gated_delta_net.qk_stride2": "10240", "llm.gated_delta_net.qk_stride3": "163840", @@ -5601,6 +5607,7 @@ "llm.gated_delta_net.head_count": "48", "llm.gated_delta_net.head_width": "128", "llm.gated_delta_net.l2_epsilon": "9.99999997e-07", + "llm.gated_delta_net.additive_epsilon": "0", "llm.gated_delta_net.qk_stride1": "128", "llm.gated_delta_net.qk_stride2": "10240", "llm.gated_delta_net.qk_stride3": "10240", diff --git a/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_companion_all.tg_c5.json b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_companion_all.tg_c5.json index c25c23af3c4d..b700589d3cbf 100644 --- a/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_companion_all.tg_c5.json +++ b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_companion_all.tg_c5.json @@ -393,6 +393,7 @@ "llm.gated_delta_net.head_count": "48", "llm.gated_delta_net.head_width": "128", "llm.gated_delta_net.l2_epsilon": "9.99999997e-07", + "llm.gated_delta_net.additive_epsilon": "0", "llm.gated_delta_net.qk_stride1": "128", "llm.gated_delta_net.qk_stride2": "10240", "llm.gated_delta_net.qk_stride3": "20480", @@ -1308,6 +1309,7 @@ "llm.gated_delta_net.head_count": "48", "llm.gated_delta_net.head_width": "128", "llm.gated_delta_net.l2_epsilon": "9.99999997e-07", + "llm.gated_delta_net.additive_epsilon": "0", "llm.gated_delta_net.qk_stride1": "128", "llm.gated_delta_net.qk_stride2": "10240", "llm.gated_delta_net.qk_stride3": "40960", @@ -2123,6 +2125,7 @@ "llm.gated_delta_net.head_count": "48", "llm.gated_delta_net.head_width": "128", "llm.gated_delta_net.l2_epsilon": "9.99999997e-07", + "llm.gated_delta_net.additive_epsilon": "0", "llm.gated_delta_net.qk_stride1": "128", "llm.gated_delta_net.qk_stride2": "10240", "llm.gated_delta_net.qk_stride3": "51200", @@ -2990,6 +2993,7 @@ "llm.gated_delta_net.head_count": "48", "llm.gated_delta_net.head_width": "128", "llm.gated_delta_net.l2_epsilon": "9.99999997e-07", + "llm.gated_delta_net.additive_epsilon": "0", "llm.gated_delta_net.qk_stride1": "128", "llm.gated_delta_net.qk_stride2": "10240", "llm.gated_delta_net.qk_stride3": "30720", @@ -3857,6 +3861,7 @@ "llm.gated_delta_net.head_count": "48", "llm.gated_delta_net.head_width": "128", "llm.gated_delta_net.l2_epsilon": "9.99999997e-07", + "llm.gated_delta_net.additive_epsilon": "0", "llm.gated_delta_net.qk_stride1": "128", "llm.gated_delta_net.qk_stride2": "10240", "llm.gated_delta_net.qk_stride3": "163840", @@ -3919,6 +3924,7 @@ "llm.gated_delta_net.head_count": "48", "llm.gated_delta_net.head_width": "128", "llm.gated_delta_net.l2_epsilon": "9.99999997e-07", + "llm.gated_delta_net.additive_epsilon": "0", "llm.gated_delta_net.qk_stride1": "128", "llm.gated_delta_net.qk_stride2": "10240", "llm.gated_delta_net.qk_stride3": "163840", @@ -5621,6 +5627,7 @@ "llm.gated_delta_net.head_count": "48", "llm.gated_delta_net.head_width": "128", "llm.gated_delta_net.l2_epsilon": "9.99999997e-07", + "llm.gated_delta_net.additive_epsilon": "0", "llm.gated_delta_net.qk_stride1": "128", "llm.gated_delta_net.qk_stride2": "10240", "llm.gated_delta_net.qk_stride3": "10240", diff --git a/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_program4.tg_c4.json b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_program4.tg_c4.json index 5fb4222c0d2b..a5b18a608763 100644 --- a/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_program4.tg_c4.json +++ b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_program4.tg_c4.json @@ -339,6 +339,7 @@ "llm.gated_delta_net.head_count": "48", "llm.gated_delta_net.head_width": "128", "llm.gated_delta_net.l2_epsilon": "9.99999997e-07", + "llm.gated_delta_net.additive_epsilon": "0", "llm.gated_delta_net.qk_stride1": "128", "llm.gated_delta_net.qk_stride2": "10240", "llm.gated_delta_net.qk_stride3": "40960", diff --git a/tests/hrx-backend-test.cpp b/tests/hrx-backend-test.cpp index 7720806cf54c..659cff9bb0af 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -1344,6 +1344,7 @@ static void run_graph_import_mixed_backend_boundary_checks() { ggml_cgraph * cpu_to_hrx = ggml_new_graph(ctx); REQUIRE(cpu_to_hrx != nullptr); + hrx_norm->flags |= GGML_TENSOR_FLAG_COMPUTE; ggml_graph_add_node(cpu_to_hrx, hrx_norm); ggml::hrx::GraphImportResult imported_cpu_to_hrx = ggml::hrx::import_ggml_graph(*cpu_to_hrx); @@ -1365,6 +1366,7 @@ static void run_graph_import_mixed_backend_boundary_checks() { ggml_cgraph * hrx_to_cpu = ggml_new_graph(ctx); REQUIRE(hrx_to_cpu != nullptr); + hrx_norm_for_cpu->flags |= GGML_TENSOR_FLAG_COMPUTE; ggml_graph_add_node(hrx_to_cpu, hrx_norm_for_cpu); ggml::hrx::GraphImportResult imported_hrx_to_cpu = ggml::hrx::import_ggml_graph(*hrx_to_cpu); @@ -1388,7 +1390,9 @@ static void run_graph_import_mixed_backend_boundary_checks() { ggml_cgraph * cpu_view_to_hrx = ggml_new_graph(ctx); REQUIRE(cpu_view_to_hrx != nullptr); + cpu_sin_view->flags |= GGML_TENSOR_FLAG_COMPUTE; ggml_graph_add_node(cpu_view_to_hrx, cpu_sin_view); + hrx_norm_from_view->flags |= GGML_TENSOR_FLAG_COMPUTE; ggml_graph_add_node(cpu_view_to_hrx, hrx_norm_from_view); ggml::hrx::GraphImportResult imported_cpu_view_to_hrx = ggml::hrx::import_ggml_graph(*cpu_view_to_hrx); @@ -2084,7 +2088,7 @@ static void run_graph_snapshot_diagnostics_checks() { REQUIRE(imported.valid()); const std::string snapshot_json = ggml::hrx::serialize_graph_snapshot_json(imported.graph, "gfx1151", 42); - REQUIRE(string_contains(snapshot_json, "ggml-hrx-graph-snapshot-v1")); + REQUIRE(string_contains(snapshot_json, "ggml-hrx-graph-snapshot-v2")); REQUIRE(string_contains(snapshot_json, "ADD")); ggml::hrx::GraphSnapshotLoadResult loaded = ggml::hrx::load_graph_snapshot_json(snapshot_json); @@ -7483,7 +7487,53 @@ static void run_lfm_dconv3_dispatch_checks() { } } -static void run_gdn_native_projection_pair_dispatch_checks() { +static ggml_tensor * build_gdn_norm(ggml_context * ctx, ggml_tensor * input, bool rms_scale_norm) { + if (rms_scale_norm) { + const float width = static_cast(input->ne[0]); + return ggml_scale(ctx, ggml_rms_norm(ctx, input, 1.e-6f / width), 1.0f / std::sqrt(width)); + } + return ggml_l2_norm(ctx, input, 1.e-6f); +} + +static void require_gdn_rms_scale_match_constraints(ggml_cgraph * graph, ggml_tensor * q, ggml_tensor * k) { + const auto matches = [&] { + auto imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + const size_t root = producer_index_for_tensor(imported.graph, q->src[0]); + std::vector covered(imported.graph.nodes().size(), false); + std::fill(covered.begin(), covered.begin() + root, true); + ggml::hrx::CommandPlan plan; + ggml::hrx::DispatchMatch match; + const bool matched = match_dispatch_at_index(imported.graph, plan, covered, root, match); + if (matched) { + require_compile_parameter(match.dispatches.front(), "llm.gated_delta_net.additive_epsilon", "1"); + } + return matched; + }; + REQUIRE(matches()); + const float scale = ggml_get_op_params_f32(q, 0); + ggml_set_op_params_f32(q, 0, scale * 0.5f); + REQUIRE(!matches()); + ggml_set_op_params_f32(q, 0, scale); + ggml_set_op_params_f32(q, 1, 1.0f); + REQUIRE(!matches()); + ggml_set_op_params_f32(q, 1, 0.0f); + const float epsilon = ggml_get_op_params_f32(k->src[0], 0); + ggml_set_op_params_f32(k->src[0], 0, epsilon * 2.0f); + REQUIRE(!matches()); + ggml_set_op_params_f32(k->src[0], 0, epsilon); + const auto rms_flags = q->src[0]->flags; + ggml_set_output(q->src[0]); + REQUIRE(!matches()); + q->src[0]->flags = rms_flags; + const auto scale_flags = k->flags; + ggml_set_output(k); + REQUIRE(!matches()); + k->flags = scale_flags; + REQUIRE(matches()); +} + +static void run_gdn_native_projection_pair_dispatch_checks(bool rms_scale_norm = false) { ggml_init_params params = {}; params.mem_size = 8 * 1024 * 1024; params.no_alloc = true; @@ -7544,8 +7594,8 @@ static void run_gdn_native_projection_pair_dispatch_checks() { return ggml_view_4d(ctx, raw, width, heads, test.tokens, 1, width * sizeof(float), hidden * sizeof(float), hidden * test.tokens * sizeof(float), offset); }; - ggml_tensor * q = ggml_l2_norm(ctx, view(qheads, 0), 1.e-6f); - ggml_tensor * k = ggml_l2_norm(ctx, view(qheads, width * qheads * sizeof(float)), 1.e-6f); + ggml_tensor * q = build_gdn_norm(ctx, view(qheads, 0), rms_scale_norm); + ggml_tensor * k = build_gdn_norm(ctx, view(qheads, width * qheads * sizeof(float)), rms_scale_norm); ggml_tensor * v = view(test.heads, 2 * width * qheads * sizeof(float)); ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, state_elements, 20); ggml_tensor * ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); @@ -7606,7 +7656,7 @@ static void run_gdn_native_projection_pair_dispatch_checks() { ggml_free(ctx); } -static void run_gdn_selected_snapshot_dispatch_checks() { +static void run_gdn_selected_snapshot_dispatch_checks(bool rms_scale_norm = false) { ggml_init_params params = {}; params.mem_size = 8 * 1024 * 1024; params.no_alloc = true; @@ -7658,8 +7708,8 @@ static void run_gdn_selected_snapshot_dispatch_checks() { return ggml_view_4d(ctx, raw, width, nheads, test.tokens, test.sequences, width*sizeof(float), hidden*sizeof(float), hidden*test.tokens*sizeof(float), offset); }; - ggml_tensor * q=ggml_l2_norm(ctx, view(qheads,0), 1.e-6f); - ggml_tensor * k=ggml_l2_norm(ctx, view(qheads,width*qheads*sizeof(float)), 1.e-6f); + ggml_tensor * q=build_gdn_norm(ctx, view(qheads,0), rms_scale_norm); + ggml_tensor * k=build_gdn_norm(ctx, view(qheads,width*qheads*sizeof(float)), rms_scale_norm); ggml_tensor * v=view(heads,2*width*qheads*sizeof(float)); ggml_tensor * gdn=ggml_gated_delta_net(ctx,q,k,v,gate,beta,state,5); const size_t attention_bytes=width*heads*test.tokens*test.sequences*sizeof(float); @@ -7685,6 +7735,11 @@ static void run_gdn_selected_snapshot_dispatch_checks() { if (test.side_use) { ggml_build_forward_expand(graph,ggml_scale(ctx,gathered,0.5f)); } + const bool first_case = test.tokens == 2; + const bool check_norm_constraints = rms_scale_norm && first_case; + if (check_norm_constraints) { + require_gdn_rms_scale_match_constraints(graph, q, k); + } auto imported=ggml::hrx::import_ggml_graph(*graph); REQUIRE(imported.valid()); ggml::hrx::DispatchScheduler scheduler; @@ -12081,6 +12136,110 @@ static void run_command_shape_hash_checks() { REQUIRE(program.command_shape_hash() == ggml::hrx::command_program_shape_hash(program.command_shape())); } +static void run_graph_compute_mask_checks() { + ggml_context * ctx = ggml_init({ 1024 * 1024, nullptr, true }); + REQUIRE(ctx != nullptr); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 128, 1); + ggml_tensor * rows = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 128, 1); + ggml_tensor * ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, 1); + ggml_tensor * branches[] = { ggml_scale(ctx, input, 2.0f), ggml_set_rows(ctx, input, rows, ids), + ggml_sin(ctx, input) }; + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(ggml_build_forward_select(graph, branches, 3, 0) == branches[0]); + graph->uid = 81001; + + const auto imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 3); + REQUIRE(imported.graph.values().find_tensor(rows) == nullptr); + REQUIRE(imported.graph.values().find_tensor(ids) == nullptr); + const auto * value = imported.graph.values().find_tensor(input); + REQUIRE(value != nullptr); + REQUIRE(imported.graph.index().consumers(value->id).size() == 1); + const auto snapshot = ggml::hrx::serialize_graph_snapshot_json(imported.graph, "gfx1151", graph->uid); + auto loaded = ggml::hrx::load_graph_snapshot_json(snapshot); + REQUIRE(loaded.valid()); + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(loaded.graph, { "gfx1151" })); + REQUIRE(scheduler.plan().dispatches.size() == 1); + + + ggml::hrx::GraphProgramCache cache; + const auto & corpus = ggml::hrx::get_qwen_kernel_corpus(); + auto lookup = cache.get_or_build(*graph, corpus, "gfx1151"); + REQUIRE(lookup.valid()); + REQUIRE(lookup.program->commands().commands.size() == 1); + REQUIRE(lookup.program->match_trusted_graph(*graph).valid()); + REQUIRE(lookup.program->match_current_graph(*graph).valid()); + + // A stable node array and UID must not reuse the formerly selected branch. + branches[0]->flags &= ~GGML_TENSOR_FLAG_COMPUTE; + branches[1]->flags |= GGML_TENSOR_FLAG_COMPUTE; + REQUIRE(!lookup.program->match_trusted_graph(*graph).valid()); + REQUIRE(!lookup.program->match_current_graph(*graph).valid()); + lookup = cache.get_or_build(*graph, corpus, "gfx1151"); + REQUIRE(lookup.valid()); + REQUIRE(cache.stats().builds == 2); + REQUIRE(lookup.program->commands().commands.size() == 1); + REQUIRE(kernel_name_for_id(lookup.program->commands().commands[0].kernel.kernel_id) == + "loom_libs:ggml_set_rows"); + + graph->uid = 81002; + REQUIRE(cache.get_or_build(*graph, corpus, "gfx1151").valid()); + REQUIRE(cache.stats().builds == 2); + branches[1]->flags &= ~GGML_TENSOR_FLAG_COMPUTE; + branches[2]->flags |= GGML_TENSOR_FLAG_COMPUTE; + REQUIRE(!cache.get_or_build(*graph, corpus, "gfx1151").valid()); + branches[2]->flags &= ~GGML_TENSOR_FLAG_COMPUTE; + lookup = cache.get_or_build(*graph, corpus, "gfx1151"); + REQUIRE(lookup.valid()); + REQUIRE(lookup.program->commands().commands.empty()); + ggml_free(ctx); +} + +static void run_inactive_set_rows_replay() { + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + ggml_context * ctx = ggml_init({ 1024 * 1024, nullptr, true }); + REQUIRE(ctx != nullptr); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 128, 1); + ggml_tensor * rows = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 128, 1); + ggml_tensor * ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, 1); + ggml_tensor * branches[] = { ggml_scale(ctx, input, 2.0f), ggml_set_rows(ctx, input, rows, ids) }; + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_select(graph, branches, 2, 0); + graph->uid = 82001; + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + REQUIRE(buffer != nullptr); + const std::vector original(128, 1.0f); + const std::vector replacement(128, 9.0f); + const int64_t row = 0; + ggml_backend_tensor_set(rows, replacement.data(), 0, replacement.size() * sizeof(float)); + ggml_backend_tensor_set(ids, &row, 0, sizeof(row)); + + for (int selected : { 0, 1, -1, 0 }) { + branches[0]->flags &= ~GGML_TENSOR_FLAG_COMPUTE; + branches[1]->flags &= ~GGML_TENSOR_FLAG_COMPUTE; + if (selected >= 0) { + branches[selected]->flags |= GGML_TENSOR_FLAG_COMPUTE; + } + ggml_backend_tensor_set(input, original.data(), 0, original.size() * sizeof(float)); + for (int replay = 0; replay < 2; ++replay) { + REQUIRE(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); + std::vector actual(128); + ggml_backend_tensor_get(input, actual.data(), 0, actual.size() * sizeof(float)); + REQUIRE(actual == (selected == 1 ? replacement : original)); + if (selected == 0) { + ggml_backend_tensor_get(branches[0], actual.data(), 0, actual.size() * sizeof(float)); + REQUIRE(actual == std::vector(128, 2.0f)); + } + } + } + ggml_backend_free(backend); + ggml_backend_buffer_free(buffer); + ggml_free(ctx); +} + static void run_graph_program_cache_uid_mismatch_checks() { ggml::hrx::GraphProgramCache cache; const ggml::hrx::KernelCorpus & corpus = ggml::hrx::get_qwen_kernel_corpus(); @@ -12917,7 +13076,7 @@ static void run_add_f32() { ggml_backend_free(backend); } -static std::vector run_rope_scale_numerical_case(ggml_backend_t backend, bool per_token, int mode) { +static std::vector run_rope_scale_numerical_case(ggml_backend_t backend, bool per_token, int mode, int offset) { ggml_init_params params = {}; params.mem_size = 512 * 1024; params.no_alloc = true; @@ -12935,6 +13094,7 @@ static std::vector run_rope_scale_numerical_case(ggml_backend_t backend, ggml_tensor * freqs = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_dims / 2); ggml_tensor * rope = ggml_rope_ext(ctx, input, pos, freqs, n_dims, mode, 0, 10000.0f, 1.0f, 0.0f, 1.19024f, 32.0f, 1.0f); + ggml_rope_set_offset(rope, offset); ggml_tensor * scales = per_token ? ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, 1, token_count) : nullptr; ggml_tensor * output = per_token ? ggml_mul(ctx, rope, scales) : ggml_scale(ctx, rope, 0.0883883461f); REQUIRE(input != nullptr); @@ -12984,12 +13144,14 @@ static void run_rope_scale_f32_numerics() { for (const bool per_token : { false, true }) { for (const int mode : { static_cast(GGML_ROPE_TYPE_NORMAL), static_cast(GGML_ROPE_TYPE_NEOX) }) { - const std::vector expected = run_rope_scale_numerical_case(cpu, per_token, mode); - const std::vector actual = run_rope_scale_numerical_case(hrx, per_token, mode); - REQUIRE(actual.size() == expected.size()); - for (size_t i = 0; i < actual.size(); ++i) { - const float tolerance = 2.0e-4f + 2.0e-4f * std::fabs(expected[i]); - REQUIRE(std::fabs(actual[i] - expected[i]) <= tolerance); + for (const int offset : { 0, 16, 32 }) { + const std::vector expected = run_rope_scale_numerical_case(cpu, per_token, mode, offset); + const std::vector actual = run_rope_scale_numerical_case(hrx, per_token, mode, offset); + REQUIRE(actual.size() == expected.size()); + for (size_t i = 0; i < actual.size(); ++i) { + const float tolerance = 2.0e-4f + 2.0e-4f * std::fabs(expected[i]); + REQUIRE(std::fabs(actual[i] - expected[i]) <= tolerance); + } } } } @@ -13653,7 +13815,7 @@ static void run_lowtoken_bias_dispatch_checks() { ggml_free(ctx); } -static void run_gdn_selected_rms_q8_dispatch_checks() { +static void run_gdn_selected_rms_q8_dispatch_checks(bool rms_scale_norm = false) { ggml_init_params params = {}; params.mem_size = 32 * 1024 * 1024; params.no_alloc = true; @@ -13736,8 +13898,8 @@ static void run_gdn_selected_rms_q8_dispatch_checks() { return ggml_view_4d(ctx, raw, width, nheads, test.tokens, test.sequences, width * sizeof(float), hidden * sizeof(float), hidden * test.tokens * sizeof(float), offset); }; - ggml_tensor * q = ggml_l2_norm(ctx, view(qheads, 0), 1.e-6f); - ggml_tensor * k = ggml_l2_norm(ctx, view(qheads, width * qheads * sizeof(float)), 1.e-6f); + ggml_tensor * q = build_gdn_norm(ctx, view(qheads, 0), rms_scale_norm); + ggml_tensor * k = build_gdn_norm(ctx, view(qheads, width * qheads * sizeof(float)), rms_scale_norm); ggml_tensor * v = view(heads, 2 * width * qheads * sizeof(float)); ggml_tensor * gdn = ggml_gated_delta_net(ctx, q, k, v, gate, beta, state, 5); if (test.flags & GdnObserved) { @@ -13833,11 +13995,60 @@ static void run_gdn_selected_rms_q8_dispatch_checks() { ggml_free(ctx); } +static void run_rope_offset_import_checks() { + ggml_init_params params = {}; + params.mem_size = 256 * 1024; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 96, 40, 1); + ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); + ggml_tensor * rope0 = ggml_rope(ctx, input, pos, 32, GGML_ROPE_TYPE_NEOX); + ggml_tensor * rope64 = ggml_rope(ctx, input, pos, 32, GGML_ROPE_TYPE_NEOX); + ggml_rope_set_offset(rope64, 64); + + const ggml::hrx::OpParams imported0 = ggml::hrx::import_op_params(*rope0); + const ggml::hrx::OpParams imported64 = ggml::hrx::import_op_params(*rope64); + const auto * params0 = ggml::hrx::op_params_as(imported0); + const auto * params64 = ggml::hrx::op_params_as(imported64); + REQUIRE(params0 != nullptr); + REQUIRE(params64 != nullptr); + REQUIRE(params0->n_offs == 0); + REQUIRE(params64->n_offs == 64); + REQUIRE(ggml::hrx::op_params_equivalent(GGML_OP_ROPE, imported0, *rope0)); + REQUIRE(!ggml::hrx::op_params_equivalent(GGML_OP_ROPE, imported0, imported64)); + REQUIRE(!ggml::hrx::op_params_equivalent(GGML_OP_ROPE, imported0, *rope64)); + + ggml_tensor * rows = ggml_reshape_2d(ctx, rope64, 96 * 40, 1); + ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 96 * 40, 8); + ggml_tensor * ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, 1); + ggml_tensor * output = ggml_set_rows(ctx, cache, rows, ids); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, output); + ggml::hrx::GraphImportResult imported_graph = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported_graph.valid()); + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported_graph.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + const auto & dispatches = scheduler.plan().dispatches; + const auto rope_dispatch = std::find_if(dispatches.begin(), dispatches.end(), [](const auto & dispatch) { + return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rope_f32"; + }); + REQUIRE(rope_dispatch != dispatches.end()); + require_compile_parameter(*rope_dispatch, "ggml.rope_f32.n_offs", "64"); + for (const auto & dispatch : dispatches) { + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) != "loom_libs:ggml_rope_set_rows_f32"); + } + + ggml_free(ctx); +} + static void register_hrx_backend_host_cases(test_runner::Suite & suite) { suite.host_case("status", [] { run_status_checks(); }); suite.host_case("command_plan_metadata", [] { run_command_plan_metadata_checks(); }); suite.host_case("dispatch_registry", [] { run_dispatch_registry_checks(); }); suite.host_case("graph_import", [] { run_graph_import_checks(); }); + suite.host_case("rope_offset_import", [] { run_rope_offset_import_checks(); }); suite.host_case("graph_import_mixed_backend_boundary", [] { run_graph_import_mixed_backend_boundary_checks(); }); suite.host_case("graph_view_external_use", [] { run_graph_view_external_use_checks(); }); suite.host_case("scale_f32_dispatch", [] { run_scale_f32_dispatch_checks(); }); @@ -13886,6 +14097,9 @@ static void register_hrx_backend_host_cases(test_runner::Suite & suite) { suite.host_case("gdn_native_projection_pair_dispatch", [] { run_gdn_native_projection_pair_dispatch_checks(); }); suite.host_case("gdn_selected_snapshot_dispatch", [] { run_gdn_selected_snapshot_dispatch_checks(); }); suite.host_case("gdn_selected_rms_q8_dispatch", [] { run_gdn_selected_rms_q8_dispatch_checks(); }); + suite.host_case("gdn_native_projection_pair_dispatch.rms_scale", [] { run_gdn_native_projection_pair_dispatch_checks(true); }); + suite.host_case("gdn_selected_snapshot_dispatch.rms_scale", [] { run_gdn_selected_snapshot_dispatch_checks(true); }); + suite.host_case("gdn_selected_rms_q8_dispatch.rms_scale", [] { run_gdn_selected_rms_q8_dispatch_checks(true); }); suite.host_case("llama_attention_matmul_dispatch", [] { run_llama_attention_matmul_dispatch_checks(); }); suite.host_case("quantized_value_projection_dispatch", [] { run_quantized_value_projection_dispatch_checks(); }); suite.host_case("prefill_value_projection_cache_dispatch", @@ -13922,6 +14136,7 @@ static void register_hrx_backend_host_cases(test_runner::Suite & suite) { suite.host_case("disjoint_transient_plan_packing", [] { run_disjoint_transient_plan_packing_checks(); }); suite.host_case("command_program_kernel_dump", [] { run_command_program_kernel_dump_checks(); }); suite.host_case("command_shape_hash", [] { run_command_shape_hash_checks(); }); + suite.host_case("graph_compute_mask", [] { run_graph_compute_mask_checks(); }); suite.host_case("graph_program_cache_uid_mismatch", [] { run_graph_program_cache_uid_mismatch_checks(); }); suite.host_case("graph_match_empty_expert_views", [] { run_graph_match_empty_expert_views_checks(); }); suite.host_case("graph_match_hash_collision", [] { run_graph_match_hash_collision_checks(); }); @@ -13933,12 +14148,81 @@ static void register_hrx_backend_host_cases(test_runner::Suite & suite) { }); } +static void run_initializer_host_input_ownership() { + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + auto * backend_context = static_cast(backend->context); + ggml_context * ctx = ggml_init({ 1024 * 1024, nullptr, false }); + REQUIRE(ctx != nullptr); + ggml::hrx::Graph graph; + std::array tensors; + std::array values; + for (size_t i = 0; i < tensors.size(); ++i) { + tensors[i] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 128); + values[i] = graph.values().get_or_add_tensor_value(tensors[i], ggml::hrx::ValueKind::External); + ggml::hrx::ValueBufferBinding binding; + binding.host_data = tensors[i]->data; + binding.length = binding.capacity = ggml_nbytes(tensors[i]); + binding.identity = i + 1; + REQUIRE(graph.values().bind_buffer(values[i], binding)); + } + auto scale = [&](size_t input, size_t output, const char * factor) { + ggml::hrx::Dispatch dispatch; + dispatch.kernel = ggml::hrx::make_kernel_specialization( + ggml::hrx::kernel_catalog_ref("loom_libs", "ggml_scale_bias_f32")); + dispatch.kernel.integer_parameters.emplace("element_count", 128); + dispatch.kernel.compile_parameters.emplace("ggml.scale.scale", factor); + dispatch.kernel.compile_parameters.emplace("ggml.scale.bias", "0"); + dispatch.bindings.push_back({ values[input], 0, ggml_nbytes(tensors[input]) }); + dispatch.bindings.push_back({ values[output], 0, ggml_nbytes(tensors[output]) }); + return dispatch; + }; + ggml::hrx::CommandPlan plan; + plan.initialization_dispatches.push_back(scale(0, 0, "2")); + plan.dispatches.push_back(scale(0, 1, "3")); + plan.dispatches.push_back(scale(2, 2, "5")); + const auto & corpus = ggml::hrx::get_qwen_kernel_corpus(); + const char * target = backend_context->device->architecture.c_str(); + const auto commands = ggml::hrx::build_command_program(graph, plan, corpus, target); + REQUIRE(commands.valid()); + const auto bindings = ggml::hrx::CommandProgramBindings::from_value_map(graph.values()); + REQUIRE(bindings.valid()); + const ggml::hrx::CommandProgramExecutionContext execution_context = { + backend_context->device->device, backend_context->stream, target, &corpus, + &backend_context->kernel_executables, &backend_context->transient_arena, + &backend_context->host_transfers, &backend_context->host_weights, + &backend_context->device->host_buffers, &backend_context->graph_replay_state, + }; + for (float input : { 1.0f, 3.0f }) { + std::fill_n(static_cast(tensors[0]->data), 128, input); + std::fill_n(static_cast(tensors[1]->data), 128, 0.0f); + std::fill_n(static_cast(tensors[2]->data), 128, 2.0f); + REQUIRE(backend_context->prepared_programs.execute( + execution_context, 84001, "initializer-host-input-ownership", commands, bindings)); + // A scheduler may replace the next batch's inputs before draining deferred host outputs. + std::fill_n(static_cast(tensors[0]->data), 128, 7.0f); + ggml_backend_synchronize(backend); + for (size_t i = 0; i < 128; ++i) { + REQUIRE(static_cast(tensors[0]->data)[i] == 7.0f); + REQUIRE(static_cast(tensors[1]->data)[i] == input * 6.0f); + REQUIRE(static_cast(tensors[2]->data)[i] == 10.0f); + } + } + REQUIRE(backend_context->prepared_programs.stats().builds == 1); + REQUIRE(backend_context->prepared_programs.stats().hits == 1); + ggml_backend_free(backend); + ggml_free(ctx); +} + static void register_hrx_backend_device_cases(test_runner::Suite & suite) { + suite.device_case("initializer_host_input_ownership", [] { run_initializer_host_input_ownership(); }); + suite.device_case("zero_output_device_support", [] { run_zero_output_device_support_checks(); }); suite.device_case("scale_f32_device_support", [] { run_scale_f32_device_support_checks(); }); suite.device_case("binary_q8_publication_numerics", [] { run_binary_q8_publication_numerics(); }); suite.device_case("qwen_expert_table_partition_prefill_512_execution", [] { run_qwen_expert_table_partition_prefill_512_execution(); }); + suite.device_case("inactive_set_rows_replay", [] { run_inactive_set_rows_replay(); }); suite.device_case("add_f32", [] { run_add_f32(); }); suite.device_case("rope_scale_f32_numerics", [] { run_rope_scale_f32_numerics(); }); suite.device_case("gather_add_rmsnorm_f32_numerics", [] { run_gather_add_rmsnorm_f32_numerics(); }); diff --git a/tests/hrx-loom-jit-test.cpp b/tests/hrx-loom-jit-test.cpp index dfd48045cded..c335d1a08c86 100644 --- a/tests/hrx-loom-jit-test.cpp +++ b/tests/hrx-loom-jit-test.cpp @@ -1126,6 +1126,7 @@ static void run_loom_jit_compile_checks() { { "llm.gated_delta_net.query_sequence_ratio", "1" }, { "llm.gated_delta_net.snapshot_stride", "49216" }, { "llm.gated_delta_net.l2_epsilon", "0.000001" }, + { "llm.gated_delta_net.additive_epsilon", "0" }, { "ggml.rmsnorm_gate_f32.rms_epsilon", "0.000001" }, { "ggml.rmsnorm_gate_f32.gate_op", "15" }, }); @@ -1471,6 +1472,7 @@ static void run_loom_jit_materialize_checks() { device, "ggml_rope_f32", { { "token_count", 3 }, { "input_span", 768 } }, { { "token_count", 7 }, { "input_span", 1792 } }, { { "ggml.rope_f32.head_size", "64" }, { "ggml.rope_f32.n_dims", "64" }, + { "ggml.rope_f32.n_offs", "0" }, { "ggml.rope_f32.head_count", "4" }, { "ggml.rope_f32.token_capacity", "8" }, { "ggml.rope_f32.input_stride1", "64" }, { "ggml.rope_f32.input_stride2", "256" }, { "ggml.rope_f32.mscale", "1" }, { "ggml.rope_f32.mode", "0" } }); diff --git a/tests/hrx-ops-test.cpp b/tests/hrx-ops-test.cpp index f72c53dc4d68..e8cb79bac47b 100644 --- a/tests/hrx-ops-test.cpp +++ b/tests/hrx-ops-test.cpp @@ -2627,6 +2627,7 @@ static void run_external_view_input_rmsnorm_case() { ggml_cgraph * graph = ggml_new_graph(ctx); REQUIRE(graph != nullptr); + output->flags |= GGML_TENSOR_FLAG_COMPUTE; ggml_graph_add_node(graph, output); require_kernel_subsequence(scheduled_kernel_sequence(graph), { "loom_libs:ggml_rmsnorm_f32" }); @@ -2675,6 +2676,7 @@ static void run_split_local_view_alias_import_case() { ggml_cgraph * graph = ggml_new_graph(ctx); REQUIRE(graph != nullptr); + flattened->flags |= GGML_TENSOR_FLAG_COMPUTE; ggml_graph_add_node(graph, flattened); ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); @@ -3467,7 +3469,7 @@ static void run_add_f32_cpu_reference_case() { ggml_backend_free(hrx_backend); } -static void run_gdn_snapshot_cpu_reference_case(int64_t tokens) { +static void run_gdn_snapshot_cpu_reference_case(int64_t tokens, bool rms_scale_norm = false) { ggml_backend_t cpu = init_cpu_backend(); ggml_backend_t hrx = ggml_backend_hrx_init(0); REQUIRE(hrx != nullptr); @@ -3515,8 +3517,13 @@ static void run_gdn_snapshot_cpu_reference_case(int64_t tokens) { f.state = ggml_new_tensor_2d(f.ctx, GGML_TYPE_F32, state_size, 1); auto * state = ggml_reshape_4d(f.ctx, f.state, width, width, heads, 1); f.cache = ggml_new_tensor_2d(f.ctx, GGML_TYPE_F32, state_size, tokens); - auto * gdn = ggml_gated_delta_net(f.ctx, ggml_l2_norm(f.ctx, f.q, 1.e-6f), - ggml_l2_norm(f.ctx, f.k, 1.e-6f), f.v, gate, beta, state, tokens); + const auto normalize = [&](ggml_tensor * input) { + if (rms_scale_norm) { + return ggml_scale(f.ctx, ggml_rms_norm(f.ctx, input, 1.e-6f / width), 1.0f / std::sqrt(float(width))); + } + return ggml_l2_norm(f.ctx, input, 1.e-6f); + }; + auto * gdn = ggml_gated_delta_net(f.ctx, normalize(f.q), normalize(f.k), f.v, gate, beta, state, tokens); f.attention = ggml_view_4d(f.ctx, gdn, width, heads, tokens, 1, width * sizeof(float), width * heads * sizeof(float), width * heads * tokens * sizeof(float), 0); auto * snapshots = ggml_view_3d(f.ctx, gdn, state_size, 1, tokens, @@ -3591,6 +3598,17 @@ static void run_gdn_snapshot_cpu_reference_case(int64_t tokens) { ggml_backend_tensor_set(h.state, hrx_state, 0, state_bytes); compute_and_compare(); } + if (rms_scale_norm) { + // Small Q/K rows distinguish additive epsilon from the old L2 clamp. + for (int64_t t = 0; t < tokens; ++t) { + for (size_t i = 0; i < 2 * row_width; ++i) { + raw[3 * t * row_width + i] *= t == 0 ? 0.0f : 1.e-4f; + } + } + set(c.raw, h.raw, raw); + set(c.state, h.state, state); + compute_and_compare(); + } for (auto * f : { &c, &h }) { ggml_backend_buffer_free(f->buffer); ggml_free(f->ctx); @@ -5240,7 +5258,9 @@ static void run_rope_set_rows_cpu_reference_case() { ggml_free(hrx_ctx); } - { + auto run_partial_rope_case = [&](int offset, int mode, int64_t partial_tokens, + int64_t partial_head_size = 128, int64_t partial_n_dims = 96, + int64_t partial_head_count = 3) { ggml_init_params params = {}; params.mem_size = 8 * 1024 * 1024; params.no_alloc = true; @@ -5249,22 +5269,20 @@ static void run_rope_set_rows_cpu_reference_case() { REQUIRE(cpu_ctx != nullptr); REQUIRE(hrx_ctx != nullptr); - constexpr int64_t partial_head_size = 128; - constexpr int64_t partial_n_dims = 96; - constexpr int64_t partial_head_count = 3; - constexpr int64_t partial_tokens = 4; ggml_tensor * cpu_input = ggml_new_tensor_3d(cpu_ctx, GGML_TYPE_F32, partial_head_size, partial_head_count, partial_tokens); ggml_tensor * cpu_pos = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_I32, partial_tokens); ggml_tensor * cpu_freq = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, partial_n_dims / 2); ggml_tensor * cpu_out = ggml_rope_ext(cpu_ctx, cpu_input, cpu_pos, cpu_freq, partial_n_dims, - GGML_ROPE_TYPE_NEOX, 0, 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + mode, 0, 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); ggml_tensor * hrx_input = ggml_new_tensor_3d(hrx_ctx, GGML_TYPE_F32, partial_head_size, partial_head_count, partial_tokens); ggml_tensor * hrx_pos = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_I32, partial_tokens); ggml_tensor * hrx_freq = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, partial_n_dims / 2); ggml_tensor * hrx_out = ggml_rope_ext(hrx_ctx, hrx_input, hrx_pos, hrx_freq, partial_n_dims, - GGML_ROPE_TYPE_NEOX, 0, 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + mode, 0, 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + ggml_rope_set_offset(cpu_out, offset); + ggml_rope_set_offset(hrx_out, offset); REQUIRE(cpu_out != nullptr); REQUIRE(hrx_out != nullptr); @@ -5283,7 +5301,7 @@ static void run_rope_set_rows_cpu_reference_case() { const std::vector input = make_pattern_f32(static_cast(partial_head_size * partial_head_count * partial_tokens), 47, 0.01f); - const std::vector pos = { 0, 1, 13, 29 }; + const std::vector pos(static_cast(partial_tokens), 29); std::vector freq(static_cast(partial_n_dims / 2)); for (size_t i = 0; i < freq.size(); ++i) { freq[i] = 1.0f + 0.01f * static_cast(i); @@ -5303,8 +5321,18 @@ static void run_rope_set_rows_cpu_reference_case() { ggml_backend_buffer_free(hrx_buffer); ggml_free(cpu_ctx); ggml_free(hrx_ctx); + }; + for (int mode : { GGML_ROPE_TYPE_NORMAL, GGML_ROPE_TYPE_NEOX }) { + for (int offset : { 0, 16, 32 }) { + run_partial_rope_case(offset, mode, 1); + run_partial_rope_case(offset, mode, 4); + } } + // MiniCPM3 rotates only the trailing 32 channels of each 96-channel query head. + run_partial_rope_case(64, GGML_ROPE_TYPE_NEOX, 1, 96, 32, 40); + run_partial_rope_case(64, GGML_ROPE_TYPE_NEOX, 512, 96, 32, 40); + auto run_view_set_rows_case = [&](int64_t view_token_count) { ggml_init_params params = {}; params.mem_size = 16 * 1024 * 1024; @@ -6160,6 +6188,8 @@ static void register_basic_ops_cases(Suite & suite) { for (int64_t tokens = 2; tokens <= 5; ++tokens) { suite.device_case("gdn_snapshot.cpu_reference." + std::to_string(tokens), [tokens] { run_gdn_snapshot_cpu_reference_case(tokens); }); + suite.device_case("gdn_snapshot.rms_scale.cpu_reference." + std::to_string(tokens), + [tokens] { run_gdn_snapshot_cpu_reference_case(tokens, true); }); } for (const ggml_type type : { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K }) { suite.device_case("ssm_conv_prefill_recurrent." + type_name(type) + ".plain",