From 5ebe0c678e912fff6ae305b25a70d9f21914c7c3 Mon Sep 17 00:00:00 2001 From: Aaron St George Date: Thu, 5 Nov 2026 06:20:49 -0700 Subject: [PATCH] ggml-hrx: support Gemma routed IDs and expert scales (#180) Bind exact router ID view spans and gather scalar expert scales with runtime ID bounds checks. Validated on gfx1151 with focused operator/backend tests, full paired HRX/Vulkan perplexity, and generation. High absolute perplexity and CPU REPEAT fallback remain documented limitations. --- .../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, 344 insertions(+), 26 deletions(-) create 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 022d1f05279..fd7fbfb481c 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,6 +19,8 @@ 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"); @@ -153,6 +155,62 @@ 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) { @@ -472,6 +530,14 @@ 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 44de4f943b7..440f706220f 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,8 +536,7 @@ 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 = - static_cast(match.top8.token_count * match.top8.route_stride) * sizeof(int32_t); + const size_t route_id_length = match.top8.route_ids->byte_count; 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 }); @@ -576,8 +575,7 @@ 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 = - static_cast(router_match.token_count * router_match.route_stride) * sizeof(int32_t); + const size_t route_id_length = router_match.route_ids->byte_count; 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 4e1acc57946..5c3a44febca 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,6 +314,9 @@ { "path": "ops/rows/get_rows_f32.loom" }, + { + "path": "ops/rows/get_rows_scalar_3d_f32.loom" + }, { "path": "ops/rows/get_rows_scale_f32.loom" }, @@ -6660,6 +6663,51 @@ "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 b028425c98e..0b6551572ed 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,7 +28,9 @@ 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 - %route_id_storage_count = index.mul %launch_token_count, %bounded_route_id_stride : 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_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 new file mode 100644 index 00000000000..0c0133b5507 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/get_rows_scalar_3d_f32.loom @@ -0,0 +1,54 @@ +// 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 a90b8e9ea56..530b449d740 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -4457,6 +4457,54 @@ 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; @@ -9312,7 +9360,9 @@ 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) * route_ids->nb[1]; + 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 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); @@ -9463,7 +9513,8 @@ 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 == dispatch.bindings[1].length); + REQUIRE(route_ids_allocation->size == static_cast(token_count * expected_expert_count) * sizeof(int32_t)); + REQUIRE(dispatch.bindings[1].length <= route_ids_allocation->size); const ggml::hrx::TransientAllocation * expert_table_allocation = ggml::hrx::find_transient_allocation(commands.transients, expert_table_transient.value); REQUIRE(expert_table_allocation != nullptr); @@ -11015,15 +11066,8 @@ static void run_qwen_router_top8_dispatch_checks() { ggml_context * ctx = ggml_init(params); REQUIRE(ctx != nullptr); - { - 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); + 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 * route_ids = nullptr; REQUIRE(logits != nullptr); ggml_tensor * output = build_qwen_router_top8_graph(ctx, logits, &route_ids); @@ -14880,6 +14924,7 @@ 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 eceafc2165d..491321e2f55 100644 --- a/tests/hrx-ops-test.cpp +++ b/tests/hrx-ops-test.cpp @@ -231,8 +231,15 @@ static std::vector router_projection_reference(const std::vector & return output; } -static std::vector router_top8_weights_reference(const std::vector & logits, int64_t token_count) { - std::vector output(kQwenRouterRouteCount * token_count); +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); for (int64_t token = 0; token < token_count; ++token) { bool used[kQwenRouterExpertCount] = {}; int64_t selected[kQwenRouterRouteCount] = {}; @@ -249,6 +256,7 @@ static std::vector router_top8_weights_reference(const std::vector } selected[route] = best_expert; used[best_expert] = true; + reference.ids[token * kQwenRouterRouteCount + route] = static_cast(best_expert); } float max_selected = -std::numeric_limits::infinity(); @@ -258,14 +266,14 @@ static std::vector router_top8_weights_reference(const std::vector float sum = 0.0f; for (int64_t route = 0; route < kQwenRouterRouteCount; ++route) { const float value = std::exp(logits[token * kQwenRouterExpertCount + selected[route]] - max_selected); - output[token * kQwenRouterRouteCount + route] = value; + reference.weights[token * kQwenRouterRouteCount + route] = value; sum += value; } for (int64_t route = 0; route < kQwenRouterRouteCount; ++route) { - output[token * kQwenRouterRouteCount + route] /= sum; + reference.weights[token * kQwenRouterRouteCount + route] /= sum; } } - return output; + return reference; } static std::vector flash_attention_reference(const std::vector & query, @@ -4498,6 +4506,72 @@ 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); @@ -4510,7 +4584,9 @@ 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 * output = build_qwen_router_top8_graph(ctx, logits); + ggml_tensor * route_ids = nullptr; + ggml_tensor * output = build_qwen_router_top8_graph(ctx, logits, &route_ids); + ggml_set_output(route_ids); ggml_cgraph * graph = ggml_new_graph(ctx); REQUIRE(graph != nullptr); @@ -4522,20 +4598,36 @@ static void run_router_top8_case(int64_t token_count) { REQUIRE(buffer != nullptr); const std::vector logits_data = make_router_logits(token_count); - const std::vector expected = router_top8_weights_reference(logits_data, token_count); + const RouterTop8Reference expected = router_top8_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.size()); + std::vector actual(expected.weights.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[i]); + const float diff = std::fabs(actual[i] - expected.weights[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); @@ -9064,7 +9156,20 @@ 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(); }); - suite.device_case("router_top8.tokens4", [] { run_router_top8_case(4); }); + 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); }); + } 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),