From fdc03e0dc0c65972421517d122369c8d9530c06a Mon Sep 17 00:00:00 2001 From: AaronStGeorge Date: Thu, 5 Nov 2026 20:57:44 +0000 Subject: [PATCH] Revert "ggml-hrx: support Gemma routed IDs and expert scales (#180)" This reverts commit 62eb8dc18998e1c5a044cf3cf35f8f2967f2caa3. --- .../common/dispatch-get-rows.cpp | 66 --------- .../qwen/dispatch-moe-router.cpp | 6 +- .../kernels/loom-libs/manifest.json | 48 ------- .../motifs/routing/router_topk_row.loom | 4 +- .../ops/rows/get_rows_scalar_3d_f32.loom | 54 -------- tests/hrx-backend-test.cpp | 67 ++-------- tests/hrx-ops-test.cpp | 125 ++---------------- 7 files changed, 26 insertions(+), 344 deletions(-) delete mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/get_rows_scalar_3d_f32.loom diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-get-rows.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-get-rows.cpp index fd7fbfb481cd..022d1f05279e 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-get-rows.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-get-rows.cpp @@ -19,8 +19,6 @@ namespace ggml::hrx { namespace { static constexpr KernelCatalogRef kGetRowsF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_get_rows_f32"); -static constexpr KernelCatalogRef kGetRowsScalar3dF32Kernel = - GGML_HRX_KERNEL_REF("loom_libs", "ggml_get_rows_scalar_3d_f32"); static constexpr KernelCatalogRef kGetRowsF32NextKernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_get_rows_f32_next"); static constexpr KernelCatalogRef kGetRowsScaleF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_get_rows_scale_f32"); @@ -155,62 +153,6 @@ static void add_primary_bindings(Dispatch & dispatch, const GetRowsMatch & match dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); } -static bool match_get_rows_scalar_3d_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { - const GraphNode * node = context.root_node; - const bool is_get_rows = node != nullptr && node->op == GGML_OP_GET_ROWS; - const bool has_two_inputs = node != nullptr && node->inputs.size() == 2; - if (!is_get_rows || !has_two_inputs) { - return false; - } - - const Value * scales = graph_value(context.graph, node->inputs[0]); - const Value * ids = graph_value(context.graph, node->inputs[1]); - const Value * output = graph_value(context.graph, node->output); - const bool has_values = scales != nullptr && ids != nullptr && output != nullptr; - if (!has_values) { - return false; - } - - const int64_t token_count = output->ne[2]; - const bool has_supported_tokens = is_supported_token_count(token_count); - const bool has_scale_shape = scales->ne[0] == 1 && scales->ne[1] == 128 && - scales->ne[2] == token_count && scales->ne[3] == 1; - const bool has_id_shape = ids->ne[0] == 8 && ids->ne[1] == token_count && - ids->ne[2] == 1 && ids->ne[3] == 1; - const bool has_output_shape = output->ne[0] == 1 && output->ne[1] == 8 && - output->ne[3] == 1; - const bool has_types = scales->type == GGML_TYPE_F32 && ids->type == GGML_TYPE_I32 && - output->type == GGML_TYPE_F32; - const bool has_contiguous_values = scales->contiguous && output->contiguous; - const bool has_id_columns = ids->nb[0] == sizeof(int32_t) && ids->nb[1] % sizeof(int32_t) == 0; - const bool has_supported_shape = has_supported_tokens && has_scale_shape && has_id_shape && - has_output_shape && has_types && has_contiguous_values && has_id_columns; - if (!has_supported_shape) { - return false; - } - - const size_t route_stride = ids->nb[1] / sizeof(int32_t); - const bool has_supported_stride = route_stride >= 8 && route_stride <= 512; - const size_t expected_id_bytes = (static_cast(token_count) - 1) * ids->nb[1] + 8 * sizeof(int32_t); - const bool has_exact_id_span = ids->byte_count == expected_id_bytes; - if (!has_supported_stride || !has_exact_id_span) { - return false; - } - - Dispatch dispatch; - dispatch.kernel = make_kernel_specialization(kGetRowsScalar3dF32Kernel); - dispatch.kernel.workload_specialization = WorkloadSpecialization::Dynamic; - dispatch.kernel.integer_parameters.emplace("token_count", token_count); - dispatch.kernel.integer_parameters.emplace("route_stride", static_cast(route_stride)); - dispatch.bindings.push_back({ ids->id, 0, ids->byte_count }); - dispatch.bindings.push_back({ scales->id, 0, scales->byte_count }); - dispatch.bindings.push_back({ output->id, 0, output->byte_count }); - - match.covered_nodes.push_back(context.root_index); - match.dispatches.push_back(std::move(dispatch)); - return true; -} - static GetRowsMatch match_get_rows_f32(const Graph & graph, const GraphNode * node) { GetRowsMatch match; if (node == nullptr || node->op != GGML_OP_GET_ROWS || node->inputs.size() != 2) { @@ -530,14 +472,6 @@ static bool match_get_rows_rmsnorm_binary_dispatch(const DispatchMatchContext & } // namespace void register_get_rows_dispatches(DispatchRegistryBuilder & registry) { - registry.add({ - "common.get_rows.scalar_3d_f32", - GGML_OP_GET_ROWS, - DispatchMatchKind::SingleOp, - 125, - DispatchSource::Common, - match_get_rows_scalar_3d_f32_dispatch, - }); registry.add({ "common.get_rows_scale.f32", GGML_OP_GET_ROWS, diff --git a/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-moe-router.cpp b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-moe-router.cpp index 440f706220f6..44de4f943b78 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-moe-router.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-moe-router.cpp @@ -536,7 +536,8 @@ static bool match_moe_router_projection_top8_fused_decode_dispatch(const Dispatc dispatch.kernel.compile_parameters.emplace("qwen3_moe.workload.token_capacity", to_config_value(match.top8.token_count)); - const size_t route_id_length = match.top8.route_ids->byte_count; + const size_t route_id_length = + static_cast(match.top8.token_count * match.top8.route_stride) * sizeof(int32_t); dispatch.bindings.push_back({ match.input->id, 0, match.input->byte_count }); dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count }); dispatch.bindings.push_back({ match.logits->id, 0, match.logits->byte_count }); @@ -575,7 +576,8 @@ static bool match_moe_router_top8_dispatch(const DispatchMatchContext & context, dispatch.kernel.compile_parameters.emplace("qwen3_moe.workload.token_capacity", to_config_value(router_match.token_count)); - const size_t route_id_length = router_match.route_ids->byte_count; + const size_t route_id_length = + static_cast(router_match.token_count * router_match.route_stride) * sizeof(int32_t); dispatch.bindings.push_back({ router_match.logits->id, 0, router_match.logits->byte_count }); dispatch.bindings.push_back({ router_match.route_ids->id, 0, route_id_length }); dispatch.bindings.push_back({ router_match.route_weights->id, 0, router_match.route_weights->byte_count }); diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/manifest.json b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/manifest.json index 5c3a44febca5..4e1acc579461 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/manifest.json +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/manifest.json @@ -314,9 +314,6 @@ { "path": "ops/rows/get_rows_f32.loom" }, - { - "path": "ops/rows/get_rows_scalar_3d_f32.loom" - }, { "path": "ops/rows/get_rows_scale_f32.loom" }, @@ -6663,51 +6660,6 @@ "motifs/quantize_q8_1_x4.loom" ] }, - { - "name": "ggml_get_rows_scalar_3d_f32", - "family": "loom_libs", - "symbol": "ggml_get_rows_scalar_3d_f32", - "source": "ops/rows/get_rows_scalar_3d_f32.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "route_stride", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "route_stride", - "type": "index" - } - ], - "bindings": [ - "route_ids", - "scales", - "output" - ], - "binding_access": [ - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "ops/rows/get_rows_scalar_3d_f32.loom" - ], - "library_sources": [] - }, - "compile_dependencies": [] - }, { "name": "ggml_get_rows_f32", "family": "loom_libs", diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/router_topk_row.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/router_topk_row.loom index 0b6551572ed1..b028425c98e4 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/router_topk_row.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/router_topk_row.loom @@ -28,9 +28,7 @@ template.def<@qwen3_moe.router.top8.row> device requires [#target.subgroup.size< %experts_per_lane0 = index.div %expert_count, %wave_size : index %experts_per_lane = index.assume %experts_per_lane0 [range(%experts_per_lane0, 1, 16)] : index %lane_expert_base = index.mul %lane, %experts_per_lane : index - %last_token = index.sub %launch_token_count, %c1 : index - %last_token_base = index.mul %last_token, %bounded_route_id_stride : index - %route_id_storage_count = index.add %last_token_base, %route_count : index + %route_id_storage_count = index.mul %launch_token_count, %bounded_route_id_stride : index %route_weight_storage_count = index.mul %launch_token_count, %route_count : index %logits_noalias, %route_ids_noalias, %route_weights_noalias = buffer.assume.noalias %logits, %route_ids, %route_weights : buffer, buffer, buffer %logits_view = buffer.view %logits_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%expert_count]xf32> diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/get_rows_scalar_3d_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/get_rows_scalar_3d_f32.loom deleted file mode 100644 index 0c0133b55079..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/get_rows_scalar_3d_f32.loom +++ /dev/null @@ -1,54 +0,0 @@ -// Copyright 2026 The HRX Authors -// SPDX-License-Identifier: Apache-2.0 - -// Gathers one F32 expert scale for each route and token. The IDs may be a -// narrow view of wider argsort rows; only the selected IDs are in bounds. -amdgpu.target @ggml_get_rows_scalar_3d_gfx11_wave64 {subgroup_size = 64} - -kernel.def target(@ggml_get_rows_scalar_3d_gfx11_wave64) @ggml_get_rows_scalar_3d_f32(%token_count: index, %route_stride: index) { - %c1 = index.constant 1 : index - %c64 = index.constant 64 : index - kernel.launch.config workgroups(%token_count, %c1, %c1) workgroup_size(%c64, %c1, %c1) : index -} launch(%token_count: index, %route_stride: index, %route_ids: buffer, %scales: buffer, %output: buffer) where [range(%token_count, 1, 2048), range(%route_stride, 8, 512)] { - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index - %bounded_route_stride = index.assume %route_stride [range(%route_stride, 8, 512)] : index - %token0 = kernel.workgroup.id : index - %route0 = kernel.workitem.id : index - %c1 = index.constant 1 : index - %c8 = index.constant 8 : index - %c0_i32 = scalar.constant 0 : i32 - %c128_i32 = scalar.constant 128 : i32 - %c0_f32 = scalar.constant 0.0 : f32 - %c0_offset = index.constant 0 : offset - %valid_token = index.cmp ult, %token0, %bounded_token_count : index - %valid_route = index.cmp ult, %route0, %c8 : index - %publish = scalar.andi %valid_token, %valid_route : i1 - %last_token = index.sub %bounded_token_count, %c1 : index - %last_row_base = index.mul %last_token, %bounded_route_stride : index - %route_id_span = index.add %last_row_base, %c8 : index - %ids_view = buffer.view %route_ids[%c0_offset] : buffer -> view<[%route_id_span]xi32> - %scales_view = buffer.view %scales[%c0_offset] : buffer -> view<[%bounded_token_count]x128xf32> - %output_view = buffer.view %output[%c0_offset] : buffer -> view<[%bounded_token_count]x8xf32> - scf.if %publish { - %token = index.assume %token0 [lt(%token0, %bounded_token_count)] : index - %route = index.assume %route0 [range(%route0, 0, 7)] : index - %row_base = index.mul %token, %bounded_route_stride : index - %id_index = index.add %row_base, %route : index - %bounded_id_index = index.assume %id_index [lt(%id_index, %route_id_span)] : index - %expert_i32 = view.load %ids_view[%bounded_id_index] : view<[%route_id_span]xi32> -> i32 - %expert_nonnegative = scalar.cmpi sge, %expert_i32, %c0_i32 : i32 - %expert_below_limit = scalar.cmpi slt, %expert_i32, %c128_i32 : i32 - %valid_expert = scalar.andi %expert_nonnegative, %expert_below_limit : i1 - // Invalid IDs have no CPU GET_ROWS value; publish zero without reading a scale row. - %scale = scf.if %valid_expert -> (f32) { - %expert0 = index.cast %expert_i32 : i32 to index - %expert = index.assume %expert0 [range(%expert0, 0, 127)] : index - %selected_scale = view.load %scales_view[%token, %expert] : view<[%bounded_token_count]x128xf32> -> f32 - scf.yield %selected_scale : f32 - } else { - scf.yield %c0_f32 : f32 - } - view.store %scale, %output_view[%token, %route] : f32, view<[%bounded_token_count]x8xf32> - } - kernel.return -} diff --git a/tests/hrx-backend-test.cpp b/tests/hrx-backend-test.cpp index 530b449d7406..a90b8e9ea567 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -4457,54 +4457,6 @@ static void schedule_token_embedding_command(ggml_context * ctx, REQUIRE(commands.commands.front().bindings[2].name == "output"); } -static void run_scalar_3d_get_rows_dispatch_checks() { - auto check_case = [](int64_t expert_count, int64_t route_count, int64_t token_count, - int64_t route_stride, bool expected_match) { - ggml_init_params params = {}; - params.mem_size = 1024 * 1024; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - - ggml_tensor * scales = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, expert_count, token_count); - ggml_tensor * id_storage = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, route_stride, token_count); - REQUIRE(scales != nullptr); - REQUIRE(id_storage != nullptr); - ggml_tensor * ids = route_stride == route_count ? id_storage : - ggml_view_2d(ctx, id_storage, route_count, token_count, id_storage->nb[1], 0); - REQUIRE(ids != nullptr); - ggml_tensor * output = ggml_get_rows(ctx, scales, ids); - REQUIRE(output != nullptr); - - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, output); - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - - ggml::hrx::DispatchScheduler scheduler; - const bool scheduled = scheduler.schedule_graph(imported.graph, test_dispatch_target()); - REQUIRE(scheduled == expected_match); - if (expected_match) { - REQUIRE(scheduler.plan().dispatches.size() == 1); - const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); - REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_get_rows_scalar_3d_f32"); - REQUIRE(dispatch.bindings.size() == 3); - REQUIRE(dispatch.bindings[0].length == ggml_nbytes(ids)); - REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == token_count); - REQUIRE(dispatch.kernel.integer_parameters.at("route_stride") == route_stride); - } - ggml_free(ctx); - }; - - check_case(128, 8, 1, 8, true); - check_case(128, 8, 2, 128, true); - check_case(128, 8, 8, 128, true); - check_case(64, 8, 2, 128, false); - check_case(128, 4, 2, 128, false); - check_case(128, 8, 2, 513, false); -} - static void run_qwen_token_embedding_dispatch_checks() { ggml_init_params params = {}; params.mem_size = 4 * 1024 * 1024; @@ -9360,9 +9312,7 @@ static void schedule_qwen_router_top8_command(ggml_context * ctx, REQUIRE(scheduler.plan().valid()); const int64_t token_count = output->ne[2]; - const size_t route_id_length = static_cast(token_count - 1) * route_ids->nb[1] + - static_cast(route_ids->ne[0]) * sizeof(int32_t); - REQUIRE(route_ids_value->byte_count == route_id_length); + const size_t route_id_length = static_cast(token_count) * route_ids->nb[1]; const size_t expert_table_bytes = qwen_expert_table_size(token_count, expected_expert_count); const size_t partition_table_bytes = qwen_partition_table_size(token_count, expected_route_count, expected_expert_count); @@ -9513,8 +9463,7 @@ static void schedule_qwen_router_top8_command(ggml_context * ctx, const ggml::hrx::TransientAllocation * route_ids_allocation = ggml::hrx::find_transient_allocation(commands.transients, route_ids_value->storage_root); REQUIRE(route_ids_allocation != nullptr); - REQUIRE(route_ids_allocation->size == static_cast(token_count * expected_expert_count) * sizeof(int32_t)); - REQUIRE(dispatch.bindings[1].length <= route_ids_allocation->size); + REQUIRE(route_ids_allocation->size == dispatch.bindings[1].length); const ggml::hrx::TransientAllocation * expert_table_allocation = ggml::hrx::find_transient_allocation(commands.transients, expert_table_transient.value); REQUIRE(expert_table_allocation != nullptr); @@ -11066,8 +11015,15 @@ static void run_qwen_router_top8_dispatch_checks() { ggml_context * ctx = ggml_init(params); REQUIRE(ctx != nullptr); - for (int64_t token_count : { 1, 2, 4, 8 }) { - ggml_tensor * logits = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenRouterExpertCount, token_count); + { + ggml_tensor * logits = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenRouterExpertCount, 4); + ggml_tensor * route_ids = nullptr; + REQUIRE(logits != nullptr); + ggml_tensor * output = build_qwen_router_top8_graph(ctx, logits, &route_ids); + schedule_qwen_router_top8_command(ctx, output, route_ids); + } + { + ggml_tensor * logits = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenRouterExpertCount, 1); ggml_tensor * route_ids = nullptr; REQUIRE(logits != nullptr); ggml_tensor * output = build_qwen_router_top8_graph(ctx, logits, &route_ids); @@ -14924,7 +14880,6 @@ static void register_hrx_backend_host_cases(test_runner::Suite & suite) { suite.host_case("rope_set_rows_dispatch", [] { run_rope_set_rows_dispatch_checks(); }); suite.host_case("graph_traversal", [] { run_graph_traversal_checks(); }); suite.host_case("qwen_token_embedding_dispatch", [] { run_qwen_token_embedding_dispatch_checks(); }); - suite.host_case("scalar_3d_get_rows_dispatch", [] { run_scalar_3d_get_rows_dispatch_checks(); }); suite.host_case("get_rows_rmsnorm_binary_dispatch", [] { run_get_rows_rmsnorm_binary_dispatch_checks(); }); suite.host_case("get_rows_scale_dispatch", [] { run_get_rows_scale_dispatch_checks(); }); suite.host_case("gather_add_dispatch", [] { run_gather_add_dispatch_checks(); }); diff --git a/tests/hrx-ops-test.cpp b/tests/hrx-ops-test.cpp index 491321e2f555..eceafc2165db 100644 --- a/tests/hrx-ops-test.cpp +++ b/tests/hrx-ops-test.cpp @@ -231,15 +231,8 @@ static std::vector router_projection_reference(const std::vector & return output; } -struct RouterTop8Reference { - std::vector ids; - std::vector weights; -}; - -static RouterTop8Reference router_top8_reference(const std::vector & logits, int64_t token_count) { - RouterTop8Reference reference; - reference.ids.resize(kQwenRouterRouteCount * token_count); - reference.weights.resize(kQwenRouterRouteCount * token_count); +static std::vector router_top8_weights_reference(const std::vector & logits, int64_t token_count) { + std::vector output(kQwenRouterRouteCount * token_count); for (int64_t token = 0; token < token_count; ++token) { bool used[kQwenRouterExpertCount] = {}; int64_t selected[kQwenRouterRouteCount] = {}; @@ -256,7 +249,6 @@ static RouterTop8Reference router_top8_reference(const std::vector & logi } selected[route] = best_expert; used[best_expert] = true; - reference.ids[token * kQwenRouterRouteCount + route] = static_cast(best_expert); } float max_selected = -std::numeric_limits::infinity(); @@ -266,14 +258,14 @@ static RouterTop8Reference router_top8_reference(const std::vector & logi float sum = 0.0f; for (int64_t route = 0; route < kQwenRouterRouteCount; ++route) { const float value = std::exp(logits[token * kQwenRouterExpertCount + selected[route]] - max_selected); - reference.weights[token * kQwenRouterRouteCount + route] = value; + output[token * kQwenRouterRouteCount + route] = value; sum += value; } for (int64_t route = 0; route < kQwenRouterRouteCount; ++route) { - reference.weights[token * kQwenRouterRouteCount + route] /= sum; + output[token * kQwenRouterRouteCount + route] /= sum; } } - return reference; + return output; } static std::vector flash_attention_reference(const std::vector & query, @@ -4506,72 +4498,6 @@ static void run_q4_prefill_source_eligibility_case() { ggml_backend_free(backend); } -static void run_scalar_3d_get_rows_case(int64_t token_count, int64_t route_stride, bool include_invalid_ids = false) { - ggml_backend_t backend = ggml_backend_hrx_init(0); - REQUIRE(backend != nullptr); - - ggml_init_params params = {}; - params.mem_size = 1024 * 1024; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - - ggml_tensor * scales = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, 128, token_count); - ggml_tensor * ids_storage = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, route_stride, token_count); - REQUIRE(scales != nullptr); - REQUIRE(ids_storage != nullptr); - ggml_tensor * ids = route_stride == 8 ? ids_storage : - ggml_view_2d(ctx, ids_storage, 8, token_count, ids_storage->nb[1], 0); - REQUIRE(ids != nullptr); - ggml_tensor * output = ggml_get_rows(ctx, scales, ids); - REQUIRE(output != nullptr); - ggml_set_output(output); - - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, output); - require_kernel_subsequence(scheduled_kernel_sequence(graph), { "loom_libs:ggml_get_rows_scalar_3d_f32" }); - - ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); - REQUIRE(buffer != nullptr); - std::vector scales_data(128 * token_count); - std::vector ids_data(route_stride * token_count, -1); - std::vector expected(8 * token_count); - for (int64_t token = 0; token < token_count; ++token) { - for (int64_t expert = 0; expert < 128; ++expert) { - scales_data[token * 128 + expert] = static_cast(token * 1024) + expert * 0.25f; - } - for (int64_t route = 0; route < 8; ++route) { - int32_t expert = route == 0 ? 0 : route == 7 ? 127 : - static_cast((token * 17 + route * 19) % 128); - if (include_invalid_ids) { - switch (route) { - case 1: expert = -1; break; - case 2: expert = 128; break; - case 3: expert = std::numeric_limits::min(); break; - case 4: expert = std::numeric_limits::max(); break; - } - } - ids_data[token * route_stride + route] = expert; - const bool valid_expert = expert >= 0 && expert < 128; - expected[token * 8 + route] = valid_expert ? scales_data[token * 128 + expert] : 0.0f; - } - } - - ggml_backend_tensor_set(scales, scales_data.data(), 0, scales_data.size() * sizeof(float)); - ggml_backend_tensor_set(ids_storage, ids_data.data(), 0, ids_data.size() * sizeof(int32_t)); - REQUIRE(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); - ggml_backend_synchronize(backend); - - std::vector actual(expected.size()); - ggml_backend_tensor_get(output, actual.data(), 0, actual.size() * sizeof(float)); - REQUIRE(actual == expected); - - ggml_backend_buffer_free(buffer); - ggml_free(ctx); - ggml_backend_free(backend); -} - static void run_router_top8_case(int64_t token_count) { ggml_backend_t backend = ggml_backend_hrx_init(0); REQUIRE(backend != nullptr); @@ -4584,9 +4510,7 @@ static void run_router_top8_case(int64_t token_count) { ggml_tensor * logits = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenRouterExpertCount, token_count); REQUIRE(logits != nullptr); - ggml_tensor * route_ids = nullptr; - ggml_tensor * output = build_qwen_router_top8_graph(ctx, logits, &route_ids); - ggml_set_output(route_ids); + ggml_tensor * output = build_qwen_router_top8_graph(ctx, logits); ggml_cgraph * graph = ggml_new_graph(ctx); REQUIRE(graph != nullptr); @@ -4598,36 +4522,20 @@ static void run_router_top8_case(int64_t token_count) { REQUIRE(buffer != nullptr); const std::vector logits_data = make_router_logits(token_count); - const RouterTop8Reference expected = router_top8_reference(logits_data, token_count); + const std::vector expected = router_top8_weights_reference(logits_data, token_count); ggml_backend_tensor_set(logits, logits_data.data(), 0, logits_data.size() * sizeof(float)); REQUIRE(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); ggml_backend_synchronize(backend); - std::vector actual(expected.weights.size()); + std::vector actual(expected.size()); ggml_backend_tensor_get(output, actual.data(), 0, actual.size() * sizeof(float)); for (size_t i = 0; i < actual.size(); ++i) { - const float diff = std::fabs(actual[i] - expected.weights[i]); + const float diff = std::fabs(actual[i] - expected[i]); REQUIRE(diff <= 1.0e-5f); } - const size_t row_stride = route_ids->nb[1] / sizeof(int32_t); - const size_t span = (static_cast(token_count) - 1) * row_stride + kQwenRouterRouteCount; - std::vector actual_ids(span); - ggml_backend_tensor_get(route_ids, actual_ids.data(), 0, span * sizeof(int32_t)); - for (int64_t token = 0; token < token_count; ++token) { - for (int64_t route = 0; route < kQwenRouterRouteCount; ++route) { - const size_t output_index = token * kQwenRouterRouteCount + route; - const size_t storage_index = token * row_stride + route; - if (actual_ids[storage_index] != expected.ids[output_index]) { - std::fprintf(stderr, "router IDs token=%" PRId64 " route=%" PRId64 " stride=%zu actual=%d expected=%d\n", - token, route, row_stride, actual_ids[storage_index], expected.ids[output_index]); - } - REQUIRE(actual_ids[storage_index] == expected.ids[output_index]); - } - } - ggml_backend_buffer_free(buffer); ggml_free(ctx); ggml_backend_free(backend); @@ -9156,20 +9064,7 @@ static void register_rmsnorm_and_scheduling_cases(Suite & suite) { GGML_GLU_OP_SWIGLU, 2048, true); }); suite.device_case("mixed_symmetric_source_eligibility", [] { run_mixed_symmetric_source_eligibility_case(); }); - for (int64_t tokens : { 1, 2, 8 }) { - for (int64_t stride : { 8, 128 }) { - suite.device_case("scalar_3d_get_rows.tokens" + std::to_string(tokens) + ".stride" + - std::to_string(stride), - [=] { run_scalar_3d_get_rows_case(tokens, stride); }); - } - } - for (int64_t stride : { 8, 128 }) { - suite.device_case("scalar_3d_get_rows.invalid_ids.stride" + std::to_string(stride), - [=] { run_scalar_3d_get_rows_case(2, stride, true); }); - } - for (int64_t tokens : { 1, 2, 4, 8 }) { - suite.device_case("router_top8.tokens" + std::to_string(tokens), [=] { run_router_top8_case(tokens); }); - } + suite.device_case("router_top8.tokens4", [] { run_router_top8_case(4); }); for (int64_t experts : { 128, 256, 512 }) { for (int64_t tokens : { 1, 2, 512 }) { suite.device_case("router_common_matmul.experts" + std::to_string(experts) + ".tokens" + std::to_string(tokens),