Skip to content

ggml-hrx: support Gemma routed IDs and expert scales #183

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
AaronStGeorge wants to merge 1 commit into
base: hrx-graph-develop-v2
Choose a base branch
from
Open

ggml-hrx: support Gemma routed IDs and expert scales #183

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
66 changes: 66 additions & 0 deletions ggml/src/ggml-hrx/dispatch_registration/common/dispatch-get-rows.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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");
Expand Down Expand Up @@ -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<size_t>(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<int64_t>(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) {
Expand Down Expand Up @@ -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,
Expand Down
6 changes: 2 additions & 4 deletions ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-moe-router.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<size_t>(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 });
Expand Down Expand Up @@ -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<size_t>(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 });
Expand Down
48 changes: 48 additions & 0 deletions ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/manifest.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
},
Expand Down Expand Up @@ -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",
Expand Down
4 changes: 3 additions & 1 deletion ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/router_topk_row.loom
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand Down
54 changes: 54 additions & 0 deletions ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/get_rows_scalar_3d_f32.loom
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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<gfx11-generic> @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<x> : index
%route0 = kernel.workitem.id<x> : 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
}
67 changes: 56 additions & 11 deletions tests/hrx-backend-test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<size_t>(token_count) * route_ids->nb[1];
const size_t route_id_length = static_cast<size_t>(token_count - 1) * route_ids->nb[1] +
static_cast<size_t>(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);
Expand Down Expand Up @@ -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<size_t>(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);
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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(); });
Expand Down
Loading