From d36ffbf878e950963147cfd6350b65149e23023c Mon Sep 17 00:00:00 2001 From: AaronStGeorge Date: Mon, 9 Nov 2026 18:25:28 +0000 Subject: [PATCH] hrx: reuse Q6 row64 tile coordinates in tiled prefill --- .../kernels/loom-libs/motifs/dequant.loom | 20 ++++++++++++------- ...iled_input_f32_dequant_accumulate_f32.loom | 7 +++++-- 2 files changed, 18 insertions(+), 9 deletions(-) diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/dequant.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/dequant.loom index 21441997e290..2bd49cd95c18 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/dequant.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/dequant.loom @@ -2707,11 +2707,9 @@ func.def inline @ggml_q6k_f32_vector4(%weight: buffer, %weight_row_byte_base: of func.return %values : vector<4xf32> } -func.def inline @ggml_q6k_native_row64_f32_vector4(%weight: buffer, %row: index, %input_size: index, %q6_block: index, %q6_group: index, %packet: index) -> (vector<4xf32>) { +func.def inline @ggml_q6k_row64_tile_lane_f32_vector4(%weight: buffer, %tile: index, %row_lane0: index, %q6_group: index, %packet: index) -> (vector<4xf32>) { %c2 = index.constant 2 : index %c4 = index.constant 4 : index - %c64 = index.constant 64 : index - %c256 = index.constant 256 : index %tile_bytes = index.constant 13440 : offset %scale_byte_add = index.constant 128 : offset %raw_byte_add = index.constant 1152 : offset @@ -2721,10 +2719,7 @@ func.def inline @ggml_q6k_native_row64_f32_vector4(%weight: buffer, %row: index, %c32_f32v = vector.constant 32.0 : vector<4xf32> %bounded_group = index.assume %q6_group [range(%q6_group, 0, 7)] : index %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index - %blocks = index.div %input_size, %c256 : index - %row_group = index.div %row, %c64 : index - %row_lane = index.rem %row, %c64 : index - %tile = index.madd %row_group, %blocks, %q6_block : index + %row_lane = index.assume %row_lane0 [range(%row_lane0, 0, 63)] : index %tile_base = index.scale %tile, %tile_bytes : index, offset -> offset %scale_byte_base = index.add %tile_base, %scale_byte_add : offset %raw_byte_base = index.add %tile_base, %raw_byte_add : offset @@ -2763,6 +2758,17 @@ func.def inline @ggml_q6k_native_row64_f32_vector4(%weight: buffer, %row: index, func.return %values : vector<4xf32> } +func.def inline @ggml_q6k_native_row64_f32_vector4(%weight: buffer, %row: index, %input_size: index, %q6_block: index, %q6_group: index, %packet: index) -> (vector<4xf32>) { + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %blocks = index.div %input_size, %c256 : index + %row_group = index.div %row, %c64 : index + %row_lane = index.rem %row, %c64 : index + %tile = index.madd %row_group, %blocks, %q6_block : index + %values = func.call @ggml_q6k_row64_tile_lane_f32_vector4(%weight, %tile, %row_lane, %q6_group, %packet) : (buffer, index, index, index, index) -> (vector<4xf32>) + func.return %values : vector<4xf32> +} + func.def inline @ggml_q8_0_f32_vector4(%weight: buffer, %row_byte_base: offset, %q8_block: index, %packet: index) -> (vector<4xf32>) { %c0 = index.constant 0 : index %c4 = index.constant 4 : index diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom index 2f906c065d9a..ae6759f06730 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom @@ -14,7 +14,7 @@ func.decl @ggml_dot_u8_s8_vector8_f32(%weight: vector<2xi32>, %activation: vecto config.def @ggml.mul_mat.activation_format = 0 : index func.decl @ggml_q4k_native_row64_f32_vector4(%weight: buffer, %row: index, %input_size: index, %q4_block: index, %q4_group: index, %packet: index) -> (vector<4xf32>) -func.decl @ggml_q6k_native_row64_f32_vector4(%weight: buffer, %row: index, %input_size: index, %q6_block: index, %q6_group: index, %packet: index) -> (vector<4xf32>) +func.decl @ggml_q6k_row64_tile_lane_f32_vector4(%weight: buffer, %tile: index, %row_lane: index, %q6_group: index, %packet: index) -> (vector<4xf32>) func.decl @ggml_dequant_f32_vector4(%iq4nl_table: vector<16xi8>, %iq3s_grid_low0: vector<32xf32>, %iq3s_grid_low1: vector<32xf32>, %iq3s_grid_low2: vector<32xf32>, %iq3s_grid_low3: vector<32xf32>, %iq3s_grid_low4: vector<32xf32>, %iq3s_grid_low5: vector<32xf32>, %iq3s_grid_low6: vector<32xf32>, %iq3s_grid_low7: vector<32xf32>, %iq3s_grid_low8: vector<32xf32>, %iq3s_grid_low9: vector<32xf32>, %iq3s_grid_low10: vector<32xf32>, %iq3s_grid_low11: vector<32xf32>, %iq3s_grid_low12: vector<32xf32>, %iq3s_grid_low13: vector<32xf32>, %iq3s_grid_low14: vector<32xf32>, %iq3s_grid_low15: vector<32xf32>, %iq3s_grid_high0: vector<32xf32>, %iq3s_grid_high1: vector<32xf32>, %iq3s_grid_high2: vector<32xf32>, %iq3s_grid_high3: vector<32xf32>, %iq3s_grid_high4: vector<32xf32>, %iq3s_grid_high5: vector<32xf32>, %iq3s_grid_high6: vector<32xf32>, %iq3s_grid_high7: vector<32xf32>, %iq3s_grid_high8: vector<32xf32>, %iq3s_grid_high9: vector<32xf32>, %iq3s_grid_high10: vector<32xf32>, %iq3s_grid_high11: vector<32xf32>, %iq3s_grid_high12: vector<32xf32>, %iq3s_grid_high13: vector<32xf32>, %iq3s_grid_high14: vector<32xf32>, %iq3s_grid_high15: vector<32xf32>, %weight_format: index, %weight: buffer, %row_byte_base: offset, %row: index, %input_size: index, %quant_block: index, %quant_group: index, %packet: index, %k: index) -> (vector<4xf32>) @@ -102,6 +102,7 @@ func.def inline @ggml_matmul_tiled_input_f32_dequant_accumulate_f32(%full_token_ %result_fragment_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16, %result_fragment_layout> %result_physical_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16> %channel_tile_base = index.mul %channel_tile, %c64 : index + %q6_blocks_per_row_group = index.div %bounded_input_size, %c256 : index %token_tile_base = index.mul %token_tile, %c32 : index %load_packet = index.rem %workitem, %c8 : index %load_k = index.mul %load_packet, %c4 : index @@ -125,6 +126,8 @@ func.def inline @ggml_matmul_tiled_input_f32_dequant_accumulate_f32(%full_token_ %weight_reuses_header = scalar.ori %weight_is_q5, %weight_is_q4_row64 : i1 %weight_is_iq4 = index.cmp eq, %weight_format, %reuse_iq4_xs : index %result00, %result01, %result10, %result11 = scf.for %quant_block = [%quant_block_begin to %quant_block_end step %c1](%block_acc00 = %init00 : vector<8xf16>, %block_acc01 = %init01 : vector<8xf16>, %block_acc10 = %init10 : vector<8xf16>, %block_acc11 = %init11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { + // All rows in this output tile share the same row64 weight tile. + %q6_tile = index.madd %channel_tile, %q6_blocks_per_row_group, %quant_block : index // Retain Q4 row64 or Q5_K headers and Q5 packet high bits for all eight K32 groups; payloads are still loaded per group. %weight_dm, %weight_scales, %weight_qh = scf.if %weight_reuses_header -> (vector<4x2xf16>, vector<4x3xi32>, vector<4x4xi8>) { %zero_dm = vector.constant 0.0 : vector<4x2xf16> @@ -229,7 +232,7 @@ func.def inline @ggml_matmul_tiled_input_f32_dequant_accumulate_f32(%full_token_ scf.yield %half : vector<4xf16> } else { %q6_or_generic = scf.if %weight_is_q6_row64 -> (vector<4xf16>) { - %wide = func.call @ggml_q6k_native_row64_f32_vector4(%weight_noalias, %channel, %bounded_input_size, %quant_block, %quant_group, %load_packet) : (buffer, index, index, index, index, index) -> (vector<4xf32>) + %wide = func.call @ggml_q6k_row64_tile_lane_f32_vector4(%weight_noalias, %q6_tile, %local_row, %quant_group, %load_packet) : (buffer, index, index, index, index) -> (vector<4xf32>) %half = vector.fptrunc %wide : vector<4xf32> to vector<4xf16> scf.yield %half : vector<4xf16> } else {