From c2c3b8789c82e90a6caf8c00f2054796902c1a50 Mon Sep 17 00:00:00 2001 From: AaronStGeorge Date: Fri, 30 Oct 2026 04:03:01 +0000 Subject: [PATCH 1/4] hrx: preserve caller inputs during metadata initialization Port the focused fix from 602f1204e onto the upstream-rebase branch. Keep initializer accesses for uploads and hazards, but copy back only writes made by main commands. Add the replay regression for deferred input replacement, outputs, in-place writes, and prepared-cache reuse. Assisted-by: Codex (cherry picked from commit a5bf921a985725866b23b18dc8592b48a95c8f90) --- ggml/src/ggml-hrx/dispatch/command-plan.h | 2 + .../runtime/command-program-executor.cpp | 11 +-- tests/hrx-backend-test.cpp | 68 +++++++++++++++++++ 3 files changed, 77 insertions(+), 4 deletions(-) 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/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/tests/hrx-backend-test.cpp b/tests/hrx-backend-test.cpp index 7720806cf54c..0b7a2cc9b0a3 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -13933,7 +13933,75 @@ 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(); }); From f784a08967d9f80181b95c0ccd62dc5930e9c32b Mon Sep 17 00:00:00 2001 From: AaronStGeorge Date: Fri, 30 Oct 2026 03:54:13 +0000 Subject: [PATCH 2/4] ggml-hrx: restore recurrent RMS-scale normalization fusion Recognize the rebased GDN Q/K normalization graph and preserve its additive epsilon semantics while retaining L2 clamp support. Cover CPU-reference numerics, snapshot/projection entry points, and fusion boundaries. Bead: lemonade-ws-7mt.1 Assisted-by: Codex (cherry picked from commit 48263e20267d8328ea8e79e403afe38f9b184e5d) --- .../llm/dispatch-gated-delta-net.cpp | 126 ++++++++++++++++-- .../gated_delta_net/scan_state_head128.loom | 12 +- .../loom/v2_draft_programs.tg_c1.json | 7 + .../loom/v2_mtp4_companion_all.tg_c5.json | 7 + .../loom/v2_mtp4_program4.tg_c4.json | 1 + tests/hrx-backend-test.cpp | 72 ++++++++-- tests/hrx-loom-jit-test.cpp | 1 + tests/hrx-ops-test.cpp | 24 +++- 8 files changed, 223 insertions(+), 27 deletions(-) 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/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/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 0b7a2cc9b0a3..6a7055f925de 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -7483,7 +7483,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 +7590,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 +7652,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 +7704,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 +7731,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; @@ -13653,7 +13704,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 +13787,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) { @@ -13886,6 +13937,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", diff --git a/tests/hrx-loom-jit-test.cpp b/tests/hrx-loom-jit-test.cpp index dfd48045cded..b8940cceea5f 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" }, }); diff --git a/tests/hrx-ops-test.cpp b/tests/hrx-ops-test.cpp index f72c53dc4d68..279b51aa3f5d 100644 --- a/tests/hrx-ops-test.cpp +++ b/tests/hrx-ops-test.cpp @@ -3467,7 +3467,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 +3515,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 +3596,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); @@ -6160,6 +6176,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", From 9497bf32ca04d873a95eb657e51a694d2fab1633 Mon Sep 17 00:00:00 2001 From: AaronStGeorge Date: Fri, 30 Oct 2026 04:14:36 +0000 Subject: [PATCH 3/4] hrx: preserve RoPE offsets for MiniCPM3 queries (lemonade-ws-7mt.3) Assisted-by: Codex (cherry picked from commit 9791711cfdd0f6264cd7b5837154bcf4b939fd2c) --- .../common/dispatch-rmsnorm.cpp | 4 ++ .../common/dispatch-rope-set-rows.cpp | 17 ++++- .../llm/dispatch-attention-qkv.cpp | 3 + .../dispatch-qwen-attention-postprocess.cpp | 3 + ggml/src/ggml-hrx/graph/op-params.cpp | 4 +- ggml/src/ggml-hrx/graph/op-params.h | 1 + .../loom-libs/motifs/rope/rope_f32.loom | 19 ++++-- .../kernels/loom-libs/ops/rope/rope_f32.loom | 1 + tests/hrx-backend-test.cpp | 66 +++++++++++++++++-- tests/hrx-loom-jit-test.cpp | 1 + tests/hrx-ops-test.cpp | 26 +++++--- 11 files changed, 123 insertions(+), 22 deletions(-) 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/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/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/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/tests/hrx-backend-test.cpp b/tests/hrx-backend-test.cpp index 6a7055f925de..e470cfd564a3 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -12968,7 +12968,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; @@ -12986,6 +12986,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); @@ -13035,12 +13036,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); + } } } } @@ -13884,11 +13887,60 @@ static void run_gdn_selected_rms_q8_dispatch_checks(bool rms_scale_norm = false) 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(); }); diff --git a/tests/hrx-loom-jit-test.cpp b/tests/hrx-loom-jit-test.cpp index b8940cceea5f..c335d1a08c86 100644 --- a/tests/hrx-loom-jit-test.cpp +++ b/tests/hrx-loom-jit-test.cpp @@ -1472,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 279b51aa3f5d..dc7feff3fad6 100644 --- a/tests/hrx-ops-test.cpp +++ b/tests/hrx-ops-test.cpp @@ -5256,7 +5256,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; @@ -5265,22 +5267,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); @@ -5299,7 +5299,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); @@ -5319,8 +5319,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; From b4f0c2e951b6a7decc7498992ff2ab163f5099fe Mon Sep 17 00:00:00 2001 From: AaronStGeorge Date: Wed, 28 Oct 2026 01:07:46 +0000 Subject: [PATCH 4/4] hrx: honor compute masks in graph import and replay caches Skip inactive branch storage and scheduling, invalidate cached programs when branch selection changes, and preserve compute masks in graph snapshots. Add host and device regressions for inactive SET_ROWS and stable-UID selection changes. Assisted-by: Codex (cherry picked from commit 760fe39af35d424745bc4b4f4cc293c406623065) --- .../ggml-hrx/dispatch/dispatch-scheduler.cpp | 6 + ggml/src/ggml-hrx/graph/graph-diagnostics.cpp | 10 +- ggml/src/ggml-hrx/graph/graph.cpp | 14 +++ ggml/src/ggml-hrx/graph/graph.h | 1 + .../ggml-hrx/runtime/graph-program-cache.cpp | 24 ++++ tests/hrx-backend-test.cpp | 112 +++++++++++++++++- tests/hrx-ops-test.cpp | 2 + 7 files changed, 164 insertions(+), 5 deletions(-) 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/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/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/tests/hrx-backend-test.cpp b/tests/hrx-backend-test.cpp index e470cfd564a3..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); @@ -12132,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(); @@ -14028,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(); }); @@ -14113,6 +14222,7 @@ static void register_hrx_backend_device_cases(test_runner::Suite & suite) { 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-ops-test.cpp b/tests/hrx-ops-test.cpp index dc7feff3fad6..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);