Skip to content

ggml-hrx: MUL_MAT_ID at input sizes that are a multiple of 32 (gpt-oss experts on HRX) + decode input row stride #66

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

Merged
bong-water-water-bong merged 4 commits into from
Oct 2, 2026
Merged

ggml-hrx: MUL_MAT_ID at input sizes that are a multiple of 32 (gpt-oss experts on HRX) + decode input row stride #66

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
7370b6e
ggml-hrx: MUL_MAT_ID WMMA kernels take input sizes that are a multipl…
Oct 2, 2026
7215efd
ggml-hrx: decode kernels read input row t at t * input_size
Oct 2, 2026
a8792c3
tests: test-hrx-decode-stride, decode-kernel input rows at input size…
Oct 2, 2026
40740e9
tests: test-hrx-mul-mat-id-k32, MUL_MAT_ID at input sizes that are a …
Oct 2, 2026
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
3 changes: 1 addition & 2 deletions ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-id-common.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,8 +52,7 @@ inline bool common_mul_mat_id_same_shape(const Value & lhs, const Value & rhs) {
}

inline bool common_mul_mat_id_supported_dense_input_size(CommonMulMatWeightFormat format, int64_t input_size) {
// the MUL_MAT_ID kernels declare input_size mul(256), also for the 32-value block formats
return common_mul_mat_supported_dense_input_size(format, input_size) && input_size % 256 == 0;
return common_mul_mat_supported_dense_input_size(format, input_size);
}

inline bool common_mul_mat_id_supported_dense_output_size(int64_t output_size) {
Expand Down
23 changes: 17 additions & 6 deletions ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_id_f32_f32_wmma_core.loom
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,8 @@ template.decl @ggml.mul_mat_id_f32_f32_wmma.publish_vector4(%arg0: i1, %arg1: in

func.decl @ggml_dequant_weight_tile_bytes(%weight_format: index) -> (offset)

func.decl @ggml_dequant_weight_row_bytes(%weight_format: index, %hidden_size: index) -> (offset)

func.decl @ggml_iq4nl_table_i8() -> (vector<16xi8>)

func.decl @ggml_iq3s_grid_f32_halves() -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>)
Expand Down Expand Up @@ -87,9 +89,11 @@ template.def<@ggml.mul_mat_id_f32_f32_wmma.core> device @ggml_mul_mat_id_f32_f32
%has_more_assignment_partitions = index.cmp ugt, %assignment_partition_count, %bounded_expert_count : index
%launch_partition_count = scf.select %has_more_assignment_partitions, %assignment_partition_count, %bounded_expert_count : index
%assignment_table_byte_base = index.scale %bounded_expert_count, %c4_bytes : index, offset -> offset
%weight_block_bytes = func.call @ggml_dequant_weight_tile_bytes(%weight_format) : (index) -> (offset)
%quant_block_count = index.div %input_size, %c256 : index
%weight_row_bytes = index.scale %quant_block_count, %weight_block_bytes : index, offset -> offset
%c255 = index.constant 255 : index
%c0_f32x4 = vector.constant 0.0 : vector<4xf32>
%padded_input_size = index.add %input_size, %c255 : index
%quant_block_count = index.div %padded_input_size, %c256 : index
%weight_row_bytes = func.call @ggml_dequant_weight_row_bytes(%weight_format, %input_size) : (index, index) -> (offset)
%weight_expert_bytes = index.scale %bounded_output_size, %weight_row_bytes : index, offset -> offset
%input_noalias, %expert_table_noalias, %partition_table_noalias, %weight_noalias = buffer.assume.noalias %input, %expert_table, %partition_table, %weight : buffer, buffer, buffer, buffer
%input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_input_route_count]x[%input_size]xf32>
Expand Down Expand Up @@ -172,7 +176,10 @@ template.def<@ggml.mul_mat_id_f32_f32_wmma.core> device @ggml_mul_mat_id_f32_f32
%local_row = index.assume %local_row0 [range(%local_row0, 0, 63)] : index
%channel = index.add %channel_tile_base, %local_row : index
%valid_channel = index.cmp ult, %channel, %bounded_output_size : index
%weight_values = scf.if %valid_channel -> (vector<4xf16>) {
%weight_k_tail = index.add %k_origin, %load_k : index
%valid_k = index.cmp ult, %weight_k_tail, %input_size : index
%valid_weight = scalar.andi %valid_channel, %valid_k : i1
%weight_values = scf.if %valid_weight -> (vector<4xf16>) {
%channel_byte_add = index.scale %channel, %weight_row_bytes : index, offset -> offset
%row_byte_base = index.add %expert_byte_base, %channel_byte_add : offset
%weight_k = index.add %k_origin, %load_k : index
Expand All @@ -187,7 +194,10 @@ template.def<@ggml.mul_mat_id_f32_f32_wmma.core> device @ggml_mul_mat_id_f32_f32
%activation_row = index.assume %local_row [range(%local_row, 0, 31)] : index
%assignment_i32 = view.load %route_stage_view[%activation_row] : view<32xi32> -> i32
%valid_assignment = scalar.cmpi sge, %assignment_i32, %c0_i32 : i32
%activation_values = scf.if %valid_assignment -> (vector<4xf16>) {
%activation_k = index.add %k_origin, %load_k : index
%valid_input = index.cmp ult, %activation_k, %input_size : index
%valid_activation = scalar.andi %valid_assignment, %valid_input : i1
%activation_values = scf.if %valid_activation -> (vector<4xf16>) {
%assignment0 = index.cast %assignment_i32 : i32 to index
%assignment = index.assume %assignment0 [range(%assignment0, 0, 65535)] : index
%bounded_assignment, %bounded_assignment_count = index.assume %assignment, %assignment_count [lt(%assignment, %assignment_count)] : index, index
Expand All @@ -197,7 +207,8 @@ template.def<@ggml.mul_mat_id_f32_f32_wmma.core> device @ggml_mul_mat_id_f32_f32
%input_route0 = index.rem %route0, %bounded_input_route_count : index
%input_route = index.assume %input_route0 [range(%input_route0, 0, 31), lt(%input_route0, %bounded_input_route_count)] : index
%input_k = index.add %k_origin, %load_k : index
%loaded = vector.load %input_view[%token, %input_route, %input_k] : view<[%bounded_token_count]x[%bounded_input_route_count]x[%input_size]xf32> -> vector<4xf32>
%mask = vector.mask.range [%input_k to %input_size step %c1] : index -> vector<4xi1>
%loaded = vector.load.mask %input_view[%token, %input_route, %input_k], %mask, %c0_f32x4 : view<[%bounded_token_count]x[%bounded_input_route_count]x[%input_size]xf32>, vector<4xi1>, vector<4xf32>
%converted = vector.fptrunc %loaded : vector<4xf32> to vector<4xf16>
scf.yield %converted : vector<4xf16>
} else {
Expand Down
10 changes: 5 additions & 5 deletions ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_f32_f32_decode.loom
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,7 @@ func.def inline @ggml_mul_mat_decode_load_f32_block(%token_count: index, %input_
%itid = index.rem %bounded_lane, %c16 : index
%vector_half = index.div %itid, %c8 : index
%vector_index = index.rem %itid, %c8 : index
%input_view = buffer.view %input[%c0_offset] : buffer -> view<[%token_count]x[%launch_block_count]x256xf32>
%input_view = buffer.view %input[%c0_offset] : buffer -> view<[%token_count]x[%input_size]xf32>
%vector_half_base = index.mul %vector_half, %c128 : index
%vector_offset = index.mul %vector_index, %c4 : index
%input_index0 = index.add %vector_half_base, %vector_offset : index
Expand All @@ -100,10 +100,10 @@ func.def inline @ggml_mul_mat_decode_load_f32_block(%token_count: index, %input_
%mask1 = vector.mask.range [%k1 to %input_size step %c1] : index -> vector<4xi1>
%mask2 = vector.mask.range [%k2 to %input_size step %c1] : index -> vector<4xi1>
%mask3 = vector.mask.range [%k3 to %input_size step %c1] : index -> vector<4xi1>
%input0 = vector.load.mask %input_view[%token, %safe_block, %input_index0], %mask0, %c0_f32x4 : view<[%token_count]x[%launch_block_count]x256xf32>, vector<4xi1>, vector<4xf32>
%input1 = vector.load.mask %input_view[%token, %safe_block, %input_index1], %mask1, %c0_f32x4 : view<[%token_count]x[%launch_block_count]x256xf32>, vector<4xi1>, vector<4xf32>
%input2 = vector.load.mask %input_view[%token, %safe_block, %input_index2], %mask2, %c0_f32x4 : view<[%token_count]x[%launch_block_count]x256xf32>, vector<4xi1>, vector<4xf32>
%input3 = vector.load.mask %input_view[%token, %safe_block, %input_index3], %mask3, %c0_f32x4 : view<[%token_count]x[%launch_block_count]x256xf32>, vector<4xi1>, vector<4xf32>
%input0 = vector.load.mask %input_view[%token, %k0], %mask0, %c0_f32x4 : view<[%token_count]x[%input_size]xf32>, vector<4xi1>, vector<4xf32>
%input1 = vector.load.mask %input_view[%token, %k1], %mask1, %c0_f32x4 : view<[%token_count]x[%input_size]xf32>, vector<4xi1>, vector<4xf32>
%input2 = vector.load.mask %input_view[%token, %k2], %mask2, %c0_f32x4 : view<[%token_count]x[%input_size]xf32>, vector<4xi1>, vector<4xf32>
%input3 = vector.load.mask %input_view[%token, %k3], %mask3, %c0_f32x4 : view<[%token_count]x[%input_size]xf32>, vector<4xi1>, vector<4xf32>
func.return %input0, %input1, %input2, %input3 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>
}

Expand Down
2 changes: 1 addition & 1 deletion ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_id_f32_f32_wmma.loom
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ template.decl @ggml.mul_mat_id_f32_f32_wmma.publish_vector4(%publish_word: i1, %

amdgpu.target<gfx11-generic> @ggml_mul_mat_id_gfx11_wave64 {subgroup_size = 64}

config.decl @ggml.mul_mat_id.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)]
config.decl @ggml.mul_mat_id.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 32)]

config.decl @ggml.mul_mat_id.output_size : %value: index where [range(%value, 1, 262144)]

Expand Down
2 changes: 1 addition & 1 deletion ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_id_postops_f32_f32_wmma.loom
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ template.decl @ggml.mul_mat_id_f32_f32_wmma.store_vector4(%arg0: index, %arg1: i

amdgpu.target<gfx11-generic> @ggml_mul_mat_id_postops_gfx11_wave64 {subgroup_size = 64}

config.decl @ggml.mul_mat_id_postops.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)]
config.decl @ggml.mul_mat_id_postops.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 32)]

config.decl @ggml.mul_mat_id_postops.output_size : %value: index where [range(%value, 1, 262144)]

Expand Down
2 changes: 1 addition & 1 deletion ...hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_id_postops_next_rmsnorm_f32_f32_wmma.loom
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ template.decl @ggml.mul_mat_id_f32_f32_wmma.store_vector4(%arg0: index, %arg1: i

amdgpu.target<gfx11-generic> @ggml_mul_mat_id_postops_next_rmsnorm_gfx11_wave64 {subgroup_size = 64}

config.decl @ggml.mul_mat_id_postops.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)]
config.decl @ggml.mul_mat_id_postops.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 32)]

config.decl @ggml.mul_mat_id_postops.output_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)]

Expand Down
10 changes: 10 additions & 0 deletions tests/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -366,4 +366,14 @@ if (TARGET ggml-hrx)
target_link_libraries(test-hrx-mxfp4 PRIVATE ggml-hrx ggml hrx::hrx loomc::loomc)
target_include_directories(test-hrx-mxfp4 PRIVATE ../ggml/include ../ggml/src ../ggml/src/ggml-hrx)
add_test(NAME test-hrx-mxfp4 COMMAND test-hrx-mxfp4)

add_executable(test-hrx-decode-stride test-hrx-decode-stride.cpp)
target_link_libraries(test-hrx-decode-stride PRIVATE ggml-hrx ggml hrx::hrx loomc::loomc)
target_include_directories(test-hrx-decode-stride PRIVATE ../ggml/include ../ggml/src ../ggml/src/ggml-hrx)
add_test(NAME test-hrx-decode-stride COMMAND test-hrx-decode-stride)

add_executable(test-hrx-mul-mat-id-k32 test-hrx-mul-mat-id-k32.cpp)
target_link_libraries(test-hrx-mul-mat-id-k32 PRIVATE ggml-hrx ggml hrx::hrx loomc::loomc)
target_include_directories(test-hrx-mul-mat-id-k32 PRIVATE ../ggml/include ../ggml/src ../ggml/src/ggml-hrx)
add_test(NAME test-hrx-mul-mat-id-k32 COMMAND test-hrx-mul-mat-id-k32)
endif()
Loading
Loading