diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat.cpp index d3cd14387e2c..145f7d347508 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat.cpp @@ -728,12 +728,25 @@ static bool match_mul_mat_swiglu_dispatch(const DispatchMatchContext & context, return false; } // engine #2137 (dispatch-q4k-packed-tail.h): the WMMA body reading both weights in the packed row layout - const bool packed_tail = use_wmma && match.gate_format == CommonMulMatWeightFormat::Q4K && - match.up_format == CommonMulMatWeightFormat::Q4K && - q4k_packed_tail_weight_ok(*match.gate_weight, match.input_size, match.output_size) && - q4k_packed_tail_weight_ok(*match.up_weight, match.input_size, match.output_size) && - q4k_packed_tail_enabled(); - const bool row_layout = pack_q4 || packed_tail; + // each weight on its own: Q4_K or Q6_K in its packed layout, anything else stored + const auto tail_format = [&](const Value & weight, CommonMulMatWeightFormat format) { + if (!use_wmma) { + return format; + } + if (format == CommonMulMatWeightFormat::Q4K && + q4k_packed_tail_weight_ok(weight, match.input_size, match.output_size) && q4k_packed_tail_enabled()) { + return CommonMulMatWeightFormat::Q4KRow64; + } + if (format == CommonMulMatWeightFormat::Q6K && + q6k_packed_rows_weight_ok(weight, match.input_size, match.output_size) && q6k_packed_rows_enabled()) { + return CommonMulMatWeightFormat::Q6KRow64; + } + return format; + }; + const CommonMulMatWeightFormat gate_bound = + pack_q4 ? CommonMulMatWeightFormat::Q4KRow64 : tail_format(*match.gate_weight, match.gate_format); + const CommonMulMatWeightFormat up_bound = + pack_q4 ? CommonMulMatWeightFormat::Q4KRow64 : tail_format(*match.up_weight, match.up_format); dispatch.kernel = make_kernel_specialization(publish_q8 ? kMulMatSwiGLUQ4Q8OutputKernel : use_q8 ? kMulMatSwiGLUQ4Q8LowTokenDotKernel : (use_direct_dot ? kMulMatSwiGLUF32F32LowTokenDotKernel : kMulMatSwiGLUF32F32WmmaKernel)); @@ -746,24 +759,25 @@ static bool match_mul_mat_swiglu_dispatch(const DispatchMatchContext & context, common_to_config_value(match.output_size)); dispatch.kernel.compile_parameters.emplace( "ggml.mul_mat_swiglu.gate_weight_format", - common_to_config_value(common_mul_mat_format_config_value(row_layout ? CommonMulMatWeightFormat::Q4KRow64 : - match.gate_format))); + common_to_config_value(common_mul_mat_format_config_value(gate_bound))); dispatch.kernel.compile_parameters.emplace( "ggml.mul_mat_swiglu.up_weight_format", - common_to_config_value(common_mul_mat_format_config_value(row_layout ? CommonMulMatWeightFormat::Q4KRow64 : - match.up_format))); + common_to_config_value(common_mul_mat_format_config_value(up_bound))); dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_swiglu.op", common_to_config_value(static_cast(binary_kind_config_value(match.op)))); dispatch.bindings.push_back(activation); - if (row_layout) { - for (const Value * weight : { match.gate_weight, match.up_weight }) { - dispatch.bindings.push_back({ weight->id, 0, weight->byte_count, kQ4KPackedK256Row64Layout, weight->type, - match.input_size, match.output_size, weight->byte_count }); + const auto bind_weight = [&](const Value & weight, CommonMulMatWeightFormat bound) { + const char * layout = bound == CommonMulMatWeightFormat::Q4KRow64 ? kQ4KPackedK256Row64Layout : + bound == CommonMulMatWeightFormat::Q6KRow64 ? kQ6KPackedK256Row64ScaleRowLayout : nullptr; + if (layout != nullptr) { + dispatch.bindings.push_back({ weight.id, 0, weight.byte_count, layout, weight.type, match.input_size, + match.output_size, weight.byte_count }); + } else { + dispatch.bindings.push_back({ weight.id, 0, weight.byte_count }); } - } else { - dispatch.bindings.push_back({ match.gate_weight->id, 0, match.gate_weight->byte_count }); - dispatch.bindings.push_back({ match.up_weight->id, 0, match.up_weight->byte_count }); - } + }; + bind_weight(*match.gate_weight, gate_bound); + bind_weight(*match.up_weight, up_bound); dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); if (publish_q8) { diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat.cpp index 6505bb7f1f32..5233efe375da 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat.cpp @@ -189,6 +189,17 @@ static bool q4k_packed_tail_applies(const DispatchMatchContext & context, !weight_has_native_layout_reader(context.graph, weight) && q4k_packed_tail_enabled(); } +static bool q6k_packed_tail_applies(const DispatchMatchContext & context, + const Value & weight, + CommonMulMatWeightFormat format, + int64_t token_count, + int64_t input_size, + int64_t output_size) { + return format == CommonMulMatWeightFormat::Q6K && token_count > 5 && token_count <= kMulMatPrefillMaxTokens && + q6k_packed_rows_weight_ok(weight, input_size, output_size) && + !weight_has_native_layout_reader(context.graph, weight) && q6k_packed_rows_enabled(); +} + static bool has_normalized_rope_consumer(const Graph & graph, const Value & value) { if (!graph.has_index()) { return false; @@ -676,9 +687,14 @@ static bool build_q6_k_prefill_wave32_dispatch(const CommonMulMatMatch & match, int input_scale_log2, KernelCatalogRef kernel, size_t root_index, - DispatchMatch & dispatch_match) { + DispatchMatch & dispatch_match, + bool packed_rows = false) { Dispatch dispatch; dispatch.kernel = make_kernel_specialization(kernel); + if (packed_rows) { + // engine #2137: the kernel reads the packed row layout decode reads (dispatch-q4k-packed-tail.h) + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_q6_k_packed.weight_layout", "2"); + } dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_q6_k_packed.input_size", common_to_config_value(match.input_size)); @@ -689,7 +705,13 @@ static bool build_q6_k_prefill_wave32_dispatch(const CommonMulMatMatch & match, common_to_config_value(match.token_count)); common_set_f16_prefill_activation_scale(dispatch, input_scale_log2); dispatch.bindings.push_back({ input, 0, input_bytes }); - dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count }); + if (packed_rows) { + dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count, kQ6KPackedK256Row64ScaleRowLayout, + match.weight->type, match.input_size, match.output_size, + match.weight->byte_count }); + } else { + dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count }); + } dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); dispatch_match.covered_nodes.push_back(root_index); @@ -719,11 +741,23 @@ static bool match_q6_k_prefill_wave32_dispatch(const DispatchMatchContext & cont if (!prepared) { return false; } + // engine #2137: with GGML_HRX_Q6K_PACKED_ROWS=1 the f16 kernel reads the packed rows decode reads. It takes + // precedence over the older GGML_HRX_Q6K_PACKED_PREFILL, whose kernel is bound to the stored layout. + // The from_packed kernel reads the stored rows on its wide tile (tokens a multiple of 256, outputs of 128) and + // 17,536-byte int8 tiles on every other tile, while this dispatch binds the stored layout. Off the wide tile it + // read past the weight: test-backend-ops MUL_MAT(q6_K, m=2048, n=128, k=4096, weight=resident) with + // GGML_HRX_Q6K_PACKED_PREFILL=1 ended in "HRX buffer download failed" and a bus error. So the switch selects that + // kernel only where it reads what is bound. + const bool from_packed = q6k_packed_prefill_enabled() && match.token_count % 256 == 0 && + match.output_size % 128 == 0; + const bool packed_rows = q6k_packed_rows_weight_ok(*match.weight, match.input_size, match.output_size) && + !weight_has_native_layout_reader(context.graph, *match.weight) && q6k_packed_rows_enabled(); return build_q6_k_prefill_wave32_dispatch(match, activation.value, activation.length, activation_scale_log2, - q6k_packed_prefill_enabled() ? kMulMatQ6KF16WmmaPrefillFromPackedKernel - : kMulMatQ6KF16WmmaPrefillWave32Kernel, + packed_rows ? kMulMatQ6KF16WmmaPrefillWave32Kernel : + from_packed ? kMulMatQ6KF16WmmaPrefillFromPackedKernel + : kMulMatQ6KF16WmmaPrefillWave32Kernel, context.root_index, - dispatch_match); + dispatch_match, packed_rows); } static bool match_symmetric_i4_prefill_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { @@ -1171,6 +1205,11 @@ static bool build_mul_mat_dispatch(const DispatchMatchContext & context, match.output_size)) { pack_q4 = true; } + if (!pack_q4 && !pack_q6 && !split_k && + q6k_packed_tail_applies(context, *match.weight, match.weight_format, match.token_count, match.input_size, + match.output_size)) { + pack_q6 = true; + } const CommonMulMatWeightFormat format = pack_q4 ? CommonMulMatWeightFormat::Q4KRow64 : pack_q6 ? CommonMulMatWeightFormat::Q6KRow64 : match.weight_format; dispatch.kernel.compile_parameters.emplace( @@ -1295,6 +1334,11 @@ static bool match_mul_mat_postops_dispatch(const DispatchMatchContext & context, match.output_size)) { pack_q4 = true; // engine #2137 } + if (!pack_q4 && !pack_q6 && + q6k_packed_tail_applies(context, *match.weight, match.weight_format, match.token_count, match.input_size, + match.output_size)) { + pack_q6 = true; // engine #2137 + } const CommonMulMatWeightFormat format = pack_q4 ? CommonMulMatWeightFormat::Q4KRow64 : pack_q6 ? CommonMulMatWeightFormat::Q6KRow64 : match.weight_format; dispatch.kernel.compile_parameters.emplace( diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-q4k-packed-tail.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-q4k-packed-tail.h index 7b878268452a..a26eece88b9c 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-q4k-packed-tail.h +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-q4k-packed-tail.h @@ -53,9 +53,8 @@ inline bool q4k_packed_tail_enabled() { // tensor in an ordinary buffer can be rewritten between computes, and the executor refuses to read one in a packed // layout ("its buffer is not flagged as weights"), which would fail a graph the stored layout runs (measured on // test-backend-ops' q4_K large-activation cases at 22 tokens). -inline bool q4k_packed_tail_weight_ok(const Value & weight, int64_t input_size, int64_t output_size) { - if (weight.type != GGML_TYPE_Q4_K || input_size % 256 != 0 || output_size % 64 != 0 || - weight.alias_source.value >= 0) { +inline bool packed_tail_weight_ok(const Value & weight, ggml_type type, int64_t input_size, int64_t output_size) { + if (weight.type != type || input_size % 256 != 0 || output_size % 64 != 0 || weight.alias_source.value >= 0) { return false; } const ggml_tensor * tensor = weight.tensor; @@ -64,4 +63,27 @@ inline bool q4k_packed_tail_weight_ok(const Value & weight, int64_t input_size, return buffer != nullptr && ggml_backend_buffer_get_usage(buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS; } +inline bool q4k_packed_tail_weight_ok(const Value & weight, int64_t input_size, int64_t output_size) { + return packed_tail_weight_ok(weight, GGML_TYPE_Q4_K, input_size, output_size); +} + +// Q6_K (GGML_HRX_Q6K_PACKED_ROWS=1): prefill reads the packed row layout decode reads, at every batch size. The tiled +// kernels take it as weight format 46 and the 128-aligned f16 prefill kernel as weight_layout 2. Stored stays the +// default. (GGML_HRX_Q6K_PACKED_PREFILL is an older switch on the aligned kernel alone; it binds the stored layout and +// saves no copy.) +inline bool q6k_packed_rows_enabled() { + static const bool enabled = [] { + const char * env = std::getenv("GGML_HRX_Q6K_PACKED_ROWS"); + const bool on = env != nullptr && env[0] == '1'; + GGML_LOG_ERROR("%s: engine #2137 q6_k prefill reads the %s layout (GGML_HRX_Q6K_PACKED_ROWS=%s)\n", __func__, + on ? "PACKED" : "stored", on ? "1" : "unset"); + return on; + }(); + return enabled; +} + +inline bool q6k_packed_rows_weight_ok(const Value & weight, int64_t input_size, int64_t output_size) { + return packed_tail_weight_ok(weight, GGML_TYPE_Q6_K, input_size, output_size); +} + } // namespace ggml::hrx 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 86dde65a6d05..c30912bed1b7 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 @@ -11080,6 +11080,92 @@ func.def inline @ggml_q6k_f32_vector4(%weight: buffer, %weight_row_byte_base: of func.return %values : vector<4xf32> } +// The same four values from HRX's packed row layout of Q6_K (q6k-packed-k256-row64-scalerow, host-memory.cpp +// materialize_q6_k_packed_k256_row64_scalerow): rows in groups of 64, and per group and 256-value block one 13,440-byte +// tile = 64 d (f16) | 64 x 2 x 8 scales (i8, index (lane*2 + s%2)*8 + s/2 for stored scale s) | per half three planes +// ql[0..32), ql[32..64), qh, each 64 rows x 32 bytes. Byte-identical codes, so the result equals ggml_q6k_f32_vector4. +// engine #2137: lets the tiled kernels share the copy decode reads. +func.def inline @ggml_q6k_packed_row64_f32_vector4(%weight: buffer, %row: index, %input_size: index, %q6_block: index, %q6_group: index, %packet: index) -> (vector<4xf32>) { + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %tile_bytes = index.constant 13440 : offset + %d_bytes = index.constant 2 : offset + %scale_plane_offset = index.constant 128 : offset + %raw_plane_offset = index.constant 1152 : offset + %byte = index.constant 1 : offset + %field_row_bytes = index.constant 32 : offset + %c4_i32v = vector.constant 4 : vector<1xi32> + %nibble_mask = vector.constant 252645135 : vector<1xi32> + %high_mask = vector.constant 50529027 : vector<1xi32> + %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 + %lane = index.rem %row, %c64 : index + %tile_index = index.madd %row_group, %blocks, %q6_block : index + %tile_base = index.scale %tile_index, %tile_bytes : index, offset -> offset + %group_in_half = index.rem %bounded_group, %c4 : index + %half = index.div %bounded_group, %c4 : index + %ql_side = index.rem %group_in_half, %c2 : index + %nibble = index.div %group_in_half, %c2 : index + %nibble_shift_index = index.mul %nibble, %c4 : index + %nibble_shift_i32 = index.cast %nibble_shift_index : index to i32 + %nibble_shift = vector.splat %nibble_shift_i32 : vector<1xi32> + %qh_shift_index = index.mul %group_in_half, %c2 : index + %qh_shift_i32 = index.cast %qh_shift_index : index to i32 + %qh_shift = vector.splat %qh_shift_i32 : vector<1xi32> + // d + %d_add = index.scale %lane, %d_bytes : index, offset -> offset + %d_byte_base = index.add %tile_base, %d_add : offset + %d_view = buffer.view %weight[%d_byte_base] : buffer -> view<1xf16> + // scale: stored index s = group*2 + packet/4, kept at (lane*2 + s%2)*8 + s/2 + %scale_half = index.div %bounded_packet, %c4 : index + %scale_lane = index.madd %lane, %c2, %scale_half : index + %scale_index = index.madd %scale_lane, %c8, %bounded_group : index + %scale_add = index.scale %scale_index, %byte : index, offset -> offset + %scale_plane_base = index.add %tile_base, %scale_plane_offset : offset + %scale_byte_base = index.add %scale_plane_base, %scale_add : offset + %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<1xi8> + // codes: plane (half*3 + field), field 0/1 = the two ql halves, 2 = qh; then the row, then the packet's 4 bytes + %raw_base = index.add %tile_base, %raw_plane_offset : offset + %half_planes = index.mul %half, %c3 : index + %ql_plane = index.add %half_planes, %ql_side : index + %qh_plane = index.add %half_planes, %c2 : index + %ql_row = index.madd %ql_plane, %c64, %lane : index + %qh_row = index.madd %qh_plane, %c64, %lane : index + %ql_row_add = index.scale %ql_row, %field_row_bytes : index, offset -> offset + %qh_row_add = index.scale %qh_row, %field_row_bytes : index, offset -> offset + %ql_row_base = index.add %raw_base, %ql_row_add : offset + %qh_row_base = index.add %raw_base, %qh_row_add : offset + %ql_view = buffer.view %weight[%ql_row_base] : buffer -> view<8xi32> + %qh_view = buffer.view %weight[%qh_row_base] : buffer -> view<8xi32> + %ql_word = vector.load %ql_view[%bounded_packet] : view<8xi32> -> vector<1xi32> + %qh_word = vector.load %qh_view[%bounded_packet] : view<8xi32> -> vector<1xi32> + %ql_shifted = vector.shrui %ql_word, %nibble_shift : vector<1xi32> + %ql = vector.andi %ql_shifted, %nibble_mask : vector<1xi32> + %qh_shifted = vector.shrui %qh_word, %qh_shift : vector<1xi32> + %qh_low = vector.andi %qh_shifted, %high_mask : vector<1xi32> + %qh = vector.shli %qh_low, %c4_i32v : vector<1xi32> + %code = vector.ori %ql, %qh : vector<1xi32> + %code_i8 = vector.bitcast %code : vector<1xi32> to vector<4xi8> + %code_f32 = vector.uitofp %code_i8 : vector<4xi8> to vector<4xf32> + %centered = vector.subf %code_f32, %c32_f32v : vector<4xf32> + %scale_i8 = view.load %scale_view[%c0] : view<1xi8> -> i8 + %d_f16 = view.load %d_view[%c0] : view<1xf16> -> f16 + %scale = scalar.sitofp %scale_i8 : i8 to f32 + %d = scalar.extf %d_f16 : f16 to f32 + %combined_scale = scalar.mulf %scale, %d : f32 + %combined_scale_vector = vector.splat %combined_scale : vector<4xf32> + %values = vector.mulf %centered, %combined_scale_vector : 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/mul_mat_f32_f32_wmma_core.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_f32_f32_wmma_core.loom index 5d3a83e6420d..f9ad35c92013 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_f32_f32_wmma_core.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_f32_f32_wmma_core.loom @@ -5,6 +5,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_packed_row64_f32_vector4(%weight: buffer, %row: index, %input_size: index, %q6_block: 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, %input_size: index, %quant_block: index, %quant_group: index, %packet: index, %k: index) -> (vector<4xf32>) @@ -133,8 +134,17 @@ func.def inline @ggml_mul_mat_f32_f32_wmma_tiled(%weight_format: index, %token_c %packed_f16 = vector.fptrunc %packed_f32 : vector<4xf32> to vector<4xf16> scf.yield %packed_f16 : vector<4xf16> } else { - %native = func.call @ggml_dequant_f16_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %weight_noalias, %row_byte_base, %bounded_input_size, %quant_block, %quant_group, %load_packet, %weight_k) : (vector<16xi8>, 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>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf16>) - scf.yield %native : vector<4xf16> + %packed_q6_format = index.constant 46 : index + %packed_q6_rows = index.cmp eq, %weight_format, %packed_q6_format : index + %stored_or_q6 = scf.if %packed_q6_rows -> (vector<4xf16>) { + %q6_f32 = func.call @ggml_q6k_packed_row64_f32_vector4(%weight_noalias, %channel, %bounded_input_size, %quant_block, %quant_group, %load_packet) : (buffer, index, index, index, index, index) -> (vector<4xf32>) + %q6_f16 = vector.fptrunc %q6_f32 : vector<4xf32> to vector<4xf16> + scf.yield %q6_f16 : vector<4xf16> + } else { + %native = func.call @ggml_dequant_f16_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %weight_noalias, %row_byte_base, %bounded_input_size, %quant_block, %quant_group, %load_packet, %weight_k) : (vector<16xi8>, 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>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf16>) + scf.yield %native : vector<4xf16> + } + scf.yield %stored_or_q6 : vector<4xf16> } scf.yield %decoded : vector<4xf16> } else { diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_quantized_f16_prefill.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_quantized_f16_prefill.loom index 396873ca8f68..3b864217290e 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_quantized_f16_prefill.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_quantized_f16_prefill.loom @@ -10,6 +10,7 @@ template.decl @ggml.binary_f32.apply_vector8(%op: index, %lhs: vector<8xf32>, %r func.decl @ggml_q4k_native_row64_f16_pair(%half_packet: i1, %paired: i1, %weight: buffer, %peer: buffer, %is_up: i1, %row_base: offset, %group0: index, %packet: index) -> (vector<16xf16>, vector<16xf16>) func.decl @ggml_q6k_f16_pair(%half_packet: i1, %weight: buffer, %weight_row_byte_base: offset, %q6_block: index, %q6_group: index, %packet: index) -> (vector<16xf16>, vector<16xf16>) +func.decl @ggml_q6k_packed_row64_f16_pair(%half_packet: i1, %weight: buffer, %row: index, %input_size: index, %q6_block: index, %q6_group: index, %packet: index) -> (vector<16xf16>, vector<16xf16>) func.def inline @ggml_f16_prefill_load_pair(%input: buffer, %token_count: index, %input_size: index, %token_origin: index, %k_origin: index) -> (vector<8xf16>, vector<8xf16>) { %c0_offset = index.constant 0 : offset @@ -270,8 +271,12 @@ func.def inline @ggml_f16_prefill_stage_mma(%paired: i1, %packed_input: i1, %bou func.def inline @ggml_f16_prefill_publish_weights(%weight_format: index, %paired: i1, %packed_input: i1, %short_weight_packets: i1, %active: i1, %bounded_input_size: index, %bounded_output_size: index, %channel_tile: index, %weight_epoch: index, %weight_noalias: buffer, %up_weight: buffer, %weight_stage: buffer, %stage_offset: offset) { + // 6 = Q6_K in the stored rows, 46 = Q6_K in the packed row layout decode reads (engine #2137); same values %q6_format = index.constant 6 : index - %q6_weights = index.cmp eq, %weight_format, %q6_format : index + %q6_packed_format = index.constant 46 : index + %q6_stored_weights = index.cmp eq, %weight_format, %q6_format : index + %q6_packed_weights = index.cmp eq, %weight_format, %q6_packed_format : index + %q6_weights = scalar.ori %q6_stored_weights, %q6_packed_weights : i1 %layout_true = scalar.constant true : i1 %staged_input = scalar.xori %packed_input, %layout_true : i1 %workgroup_size = kernel.workgroup.size : index @@ -344,10 +349,16 @@ func.def inline @ggml_f16_prefill_publish_weights(%weight_format: index, %paired %valid_channel = index.cmp ult, %channel, %bounded_output_size : index %weight_values0, %weight_values1 = scf.if %valid_channel -> (vector<16xf16>, vector<16xf16>) { %decoded0, %decoded1 = scf.if %q6_weights -> (vector<16xf16>, vector<16xf16>) { - %native_block_bytes = index.constant 210 : offset - %native_row_bytes = index.scale %quant_block_count, %native_block_bytes : index, offset -> offset - %row_byte_base = index.scale %channel, %native_row_bytes : index, offset -> offset - %decoded0, %decoded1 = func.call @ggml_q6k_f16_pair(%short_weight_packets, %weight_noalias, %row_byte_base, %quant_block, %weight_group, %load_packet) : (i1, buffer, offset, index, index, index) -> (vector<16xf16>, vector<16xf16>) + %decoded0, %decoded1 = scf.if %q6_packed_weights -> (vector<16xf16>, vector<16xf16>) { + %p0, %p1 = func.call @ggml_q6k_packed_row64_f16_pair(%short_weight_packets, %weight_noalias, %channel, %bounded_input_size, %quant_block, %weight_group, %load_packet) : (i1, buffer, index, index, index, index, index) -> (vector<16xf16>, vector<16xf16>) + scf.yield %p0, %p1 : vector<16xf16>, vector<16xf16> + } else { + %native_block_bytes = index.constant 210 : offset + %native_row_bytes = index.scale %quant_block_count, %native_block_bytes : index, offset -> offset + %row_byte_base = index.scale %channel, %native_row_bytes : index, offset -> offset + %s0, %s1 = func.call @ggml_q6k_f16_pair(%short_weight_packets, %weight_noalias, %row_byte_base, %quant_block, %weight_group, %load_packet) : (i1, buffer, offset, index, index, index) -> (vector<16xf16>, vector<16xf16>) + scf.yield %s0, %s1 : vector<16xf16>, vector<16xf16> + } scf.yield %decoded0, %decoded1 : vector<16xf16>, vector<16xf16> } else { %packed_record_bytes = index.constant 16 : offset @@ -526,8 +537,12 @@ func.def inline @ggml_prefill_conv4_strip(%deferred_state: i1, %width: index, %c } func.def inline @ggml_mul_mat_quantized_f16_prefill_wave32_body(%weight_format: index, %binary_op: index, %paired: i1, %packed_input: i1, %packed_output: i1, %token_count: index, %input_size0: index, %output_size0: index, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %up_weight: buffer, %output: buffer, %f16_output: buffer, %conv_epilogue: i1, %conv_state: buffer, %conv_filter: buffer, %conv_output: buffer, %conv_cache: buffer, %deferred_state: i1) { + // 6 = Q6_K in the stored rows, 46 = Q6_K in the packed row layout decode reads (engine #2137); same values %q6_format = index.constant 6 : index - %q6_weights = index.cmp eq, %weight_format, %q6_format : index + %q6_packed_format = index.constant 46 : index + %q6_stored_weights = index.cmp eq, %weight_format, %q6_format : index + %q6_packed_weights = index.cmp eq, %weight_format, %q6_packed_format : index + %q6_weights = scalar.ori %q6_stored_weights, %q6_packed_weights : i1 %layout_true = scalar.constant true : i1 %unpaired = scalar.xori %paired, %layout_true : i1 %packed_single = scalar.andi %unpaired, %packed_input : i1 diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/q6_k_f16.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/q6_k_f16.loom index 6c9bccfdff33..4698b99a2204 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/q6_k_f16.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/q6_k_f16.loom @@ -180,3 +180,133 @@ func.def public inline @ggml_q6k_f16_pair(%half_packet: i1, %weight: buffer, %we %values1 = vector.fptrunc %values_f32_1 : vector<16xf32> to vector<16xf16> func.return %values0, %values1 : vector<16xf16>, vector<16xf16> } + +// The same pair from HRX's packed row layout of Q6_K (q6k-packed-k256-row64-scalerow, the copy decode reads; layout in +// host-memory.cpp materialize_q6_k_packed_k256_row64_scalerow and ggml_q6k_packed_row64_f32_vector4). Byte-identical +// codes, so the values equal ggml_q6k_f16_pair's. engine #2137: lets aligned prefill share that copy. +func.def public inline @ggml_q6k_packed_row64_f16_pair(%half_packet: i1, %weight: buffer, %row: index, %input_size: index, %q6_block: index, %q6_group: index, %packet: index) -> (vector<16xf16>, vector<16xf16>) { + %padding = vector.constant 0 : vector<2xi32> + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c4_i32v = vector.constant 4 : vector<4xi32> + %nibble_mask = vector.constant 252645135 : vector<4xi32> + %high_mask = vector.constant 50529027 : vector<4xi32> + %c32_f32v = vector.constant 32.0 : vector<16xf32> + %bounded_group = index.assume %q6_group [range(%q6_group, 0, 6), multiple_of(%q6_group, 2)] : index + %bounded_packet = index.assume %packet [range(%packet, 0, 6), multiple_of(%packet, 2)] : index + %c3 = index.constant 3 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %tile_bytes = index.constant 13440 : offset + %d_bytes = index.constant 2 : offset + %byte = index.constant 1 : offset + %scale_plane_offset = index.constant 128 : offset + %raw_plane_offset = index.constant 1152 : offset + %field_row_bytes = index.constant 32 : offset + %group_in_half = index.rem %bounded_group, %c4 : index + %half = index.div %bounded_group, %c4 : index + %blocks = index.div %input_size, %c256 : index + %row_group = index.div %row, %c64 : index + %lane = index.rem %row, %c64 : index + %tile_index = index.madd %row_group, %blocks, %q6_block : index + %tile_base = index.scale %tile_index, %tile_bytes : index, offset -> offset + %d_add = index.scale %lane, %d_bytes : index, offset -> offset + %d_byte_base = index.add %tile_base, %d_add : offset + %d_view = buffer.view %weight[%d_byte_base] : buffer -> view<1xf16> + // codes: plane (half*3 + field) of 64 rows x 32 bytes, field 0/1 = the two ql halves (this group and the next), 2 = qh + %raw_base = index.add %tile_base, %raw_plane_offset : offset + %half_planes = index.mul %half, %c3 : index + %ql_plane1 = index.add %half_planes, %c1 : index + %qh_plane = index.add %half_planes, %c2 : index + %ql_row0 = index.madd %half_planes, %c64, %lane : index + %ql_row1 = index.madd %ql_plane1, %c64, %lane : index + %qh_row = index.madd %qh_plane, %c64, %lane : index + %ql_row_add0 = index.scale %ql_row0, %field_row_bytes : index, offset -> offset + %ql_row_add1 = index.scale %ql_row1, %field_row_bytes : index, offset -> offset + %qh_row_add = index.scale %qh_row, %field_row_bytes : index, offset -> offset + %ql_row_base0 = index.add %raw_base, %ql_row_add0 : offset + %ql_row_base1 = index.add %raw_base, %ql_row_add1 : offset + %qh_row_base = index.add %raw_base, %qh_row_add : offset + %ql_view0 = buffer.view %weight[%ql_row_base0] : buffer -> view<8xi32> + %ql_view1 = buffer.view %weight[%ql_row_base1] : buffer -> view<8xi32> + %qh_view = buffer.view %weight[%qh_row_base] : buffer -> view<8xi32> + %nibble = index.div %group_in_half, %c2 : index + %nibble_shift_index = index.mul %nibble, %c4 : index + %nibble_shift_i32 = index.cast %nibble_shift_index : index to i32 + %nibble_shift = vector.splat %nibble_shift_i32 : vector<4xi32> + %qh_shift_index0 = index.mul %group_in_half, %c2 : index + %qh_shift_index1 = index.add %qh_shift_index0, %c2 : index + %qh_shift_i32_0 = index.cast %qh_shift_index0 : index to i32 + %qh_shift_i32_1 = index.cast %qh_shift_index1 : index to i32 + %qh_shift0 = vector.splat %qh_shift_i32_0 : vector<4xi32> + %qh_shift1 = vector.splat %qh_shift_i32_1 : vector<4xi32> + // scales: stored index s = group*2 + packet/4 sits at (lane*2 + s%2)*8 + s/2, so this group and the next are adjacent bytes + %scale_packet_half = index.div %bounded_packet, %c4 : index + %scale_lane = index.madd %lane, %c2, %scale_packet_half : index + %scale_index = index.madd %scale_lane, %c8, %bounded_group : index + %scale_add = index.scale %scale_index, %byte : index, offset -> offset + %scale_plane_base = index.add %tile_base, %scale_plane_offset : offset + %scale_byte_base = index.add %scale_plane_base, %scale_add : offset + %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<2xi8> + %ql_word0 = scf.if %half_packet -> (vector<4xi32>) { + %words = vector.load %ql_view0[%bounded_packet] : view<8xi32> -> vector<2xi32> + %padded = vector.concat<0> %words, %padding : vector<2xi32>, vector<2xi32> -> vector<4xi32> + scf.yield %padded : vector<4xi32> + } else { + %words = vector.load %ql_view0[%bounded_packet] : view<8xi32> -> vector<4xi32> + scf.yield %words : vector<4xi32> + } + %ql_word1 = scf.if %half_packet -> (vector<4xi32>) { + %words = vector.load %ql_view1[%bounded_packet] : view<8xi32> -> vector<2xi32> + %padded = vector.concat<0> %words, %padding : vector<2xi32>, vector<2xi32> -> vector<4xi32> + scf.yield %padded : vector<4xi32> + } else { + %words = vector.load %ql_view1[%bounded_packet] : view<8xi32> -> vector<4xi32> + scf.yield %words : vector<4xi32> + } + %qh_word = scf.if %half_packet -> (vector<4xi32>) { + %words = vector.load %qh_view[%bounded_packet] : view<8xi32> -> vector<2xi32> + %padded = vector.concat<0> %words, %padding : vector<2xi32>, vector<2xi32> -> vector<4xi32> + scf.yield %padded : vector<4xi32> + } else { + %words = vector.load %qh_view[%bounded_packet] : view<8xi32> -> vector<4xi32> + scf.yield %words : vector<4xi32> + } + %scale_i8_0 = view.load %scale_view[%c0] : view<2xi8> -> i8 + %scale_i8_1 = view.load %scale_view[%c1] : view<2xi8> -> i8 + %d_f16 = view.load %d_view[%c0] : view<1xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %ql_shifted0 = vector.shrui %ql_word0, %nibble_shift : vector<4xi32> + %ql0 = vector.andi %ql_shifted0, %nibble_mask : vector<4xi32> + %qh_shifted0 = vector.shrui %qh_word, %qh_shift0 : vector<4xi32> + %qh_low0 = vector.andi %qh_shifted0, %high_mask : vector<4xi32> + %qh0 = vector.shli %qh_low0, %c4_i32v : vector<4xi32> + %code0 = vector.ori %ql0, %qh0 : vector<4xi32> + %code_i8_0 = vector.bitcast %code0 : vector<4xi32> to vector<16xi8> + %code_f32_0 = vector.uitofp %code_i8_0 : vector<16xi8> to vector<16xf32> + %centered0 = vector.subf %code_f32_0, %c32_f32v : vector<16xf32> + %scale0 = scalar.sitofp %scale_i8_0 : i8 to f32 + %combined_scale0 = scalar.mulf %scale0, %d : f32 + %combined_scale_vector0 = vector.splat %combined_scale0 : vector<16xf32> + %values_f32_0 = vector.mulf %centered0, %combined_scale_vector0 : vector<16xf32> + %values0 = vector.fptrunc %values_f32_0 : vector<16xf32> to vector<16xf16> + %ql_shifted1 = vector.shrui %ql_word1, %nibble_shift : vector<4xi32> + %ql1 = vector.andi %ql_shifted1, %nibble_mask : vector<4xi32> + %qh_shifted1 = vector.shrui %qh_word, %qh_shift1 : vector<4xi32> + %qh_low1 = vector.andi %qh_shifted1, %high_mask : vector<4xi32> + %qh1 = vector.shli %qh_low1, %c4_i32v : vector<4xi32> + %code1 = vector.ori %ql1, %qh1 : vector<4xi32> + %code_i8_1 = vector.bitcast %code1 : vector<4xi32> to vector<16xi8> + %code_f32_1 = vector.uitofp %code_i8_1 : vector<16xi8> to vector<16xf32> + %centered1 = vector.subf %code_f32_1, %c32_f32v : vector<16xf32> + %scale1 = scalar.sitofp %scale_i8_1 : i8 to f32 + %combined_scale1 = scalar.mulf %scale1, %d : f32 + %combined_scale_vector1 = vector.splat %combined_scale1 : vector<16xf32> + %values_f32_1 = vector.mulf %centered1, %combined_scale_vector1 : vector<16xf32> + %values1 = vector.fptrunc %values_f32_1 : vector<16xf32> to vector<16xf16> + func.return %values0, %values1 : vector<16xf16>, vector<16xf16> +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_q6_k_packed_f16_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_q6_k_packed_f16_wmma.loom index 4109a00c59c5..0b70b7bcb3b7 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_q6_k_packed_f16_wmma.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_q6_k_packed_f16_wmma.loom @@ -10,6 +10,11 @@ config.decl @ggml.mul_mat_q6_k_packed.weight_offset : %value: index where [range config.decl @ggml.mul_mat_q6_k_packed.token_capacity : %value: index where [range(%value, 128, 2048), multiple_of(%value, 128)] +// engine #2137: which resident layout the f16 prefill kernel reads the weight in. 0 = the stored Q6_K rows (default); +// 2 = the packed row layout decode reads (q6k-packed-k256-row64-scalerow), so prefill and decode share one device copy. +// (1 is the body's int8-tile layout, selected by the from_packed kernel, not by this parameter.) +config.def @ggml.mul_mat_q6_k_packed.weight_layout = 0 : index + // Defined (default 0) with the shared f16 prefill body in motifs/mul_mat_quantized_f16_prefill.loom. config.decl @ggml.f16_prefill.activation_scale_log2 : %value: index where [range(%value, 0, 6)] @@ -34,6 +39,7 @@ template.decl @ggml.mul_mat_q6_k_packed.prefill_wave32.body(%input_is_f16: i1, % amdgpu.target @ggml_mul_mat_q6_k_prefill_gfx11_wave32 {subgroup_size = 32} func.decl @ggml_q6k_f16_pair(%half_packet: i1, %weight: buffer, %weight_row_byte_base: offset, %q6_block: index, %q6_group: index, %packet: index) -> (vector<16xf16>, vector<16xf16>) +func.decl @ggml_q6k_packed_row64_f16_pair(%half_packet: i1, %weight: buffer, %row: index, %input_size: index, %q6_block: index, %q6_group: index, %packet: index) -> (vector<16xf16>, vector<16xf16>) template.def<@ggml.mul_mat_q6_k_packed.publish_f16_wave32> device @ggml_mul_mat_q6_k_packed_publish_f16_wave32(%result: vector<16xf16>, %token_offset: index, %token_tile_base: index, %bounded_token_count: index, %bounded_output_size: index, %channel_tile_base: index, %subgroup_channel_add: index, %result_stage: buffer, %output_noalias: buffer, %output_accumulation: index) { %lane = kernel.subgroup.lane.id : index @@ -1430,10 +1436,18 @@ template.def<@ggml.mul_mat_q6_k_packed.prefill_wave32.body> device @ggml_mul_mat %v1 = vector.fptrunc %v1f : vector<16xf32> to vector<16xf16> scf.yield %v0, %v1 : vector<16xf16>, vector<16xf16> } else { - %channel_byte_index = index.mul %channel, %weight_row_bytes : index - %row_byte_base = index.cast %channel_byte_index : index to offset %half_packet = scalar.constant false : i1 - %d0, %d1 = func.call @ggml_q6k_f16_pair(%half_packet, %weight_noalias, %row_byte_base, %quant_block, %quant_group, %load_packet) : (i1, buffer, offset, index, index, index) -> (vector<16xf16>, vector<16xf16>) + %c2_layout = index.constant 2 : index + %is_packed_rows = index.cmp eq, %weight_layout, %c2_layout : index + %d0, %d1 = scf.if %is_packed_rows -> (vector<16xf16>, vector<16xf16>) { + %p0, %p1 = func.call @ggml_q6k_packed_row64_f16_pair(%half_packet, %weight_noalias, %channel, %bounded_input_size, %quant_block, %quant_group, %load_packet) : (i1, buffer, index, index, index, index, index) -> (vector<16xf16>, vector<16xf16>) + scf.yield %p0, %p1 : vector<16xf16>, vector<16xf16> + } else { + %channel_byte_index = index.mul %channel, %weight_row_bytes : index + %row_byte_base = index.cast %channel_byte_index : index to offset + %s0, %s1 = func.call @ggml_q6k_f16_pair(%half_packet, %weight_noalias, %row_byte_base, %quant_block, %quant_group, %load_packet) : (i1, buffer, offset, index, index, index) -> (vector<16xf16>, vector<16xf16>) + scf.yield %s0, %s1 : vector<16xf16>, vector<16xf16> + } scf.yield %d0, %d1 : vector<16xf16>, vector<16xf16> } scf.yield %decoded0, %decoded1 : vector<16xf16>, vector<16xf16> @@ -1674,7 +1688,12 @@ kernel.def target(@ggml_mul_mat_q6_k_prefill_gfx11_wave32) export("ggml_mul_mat_ %c512 = index.constant 512 : index %wide_tile = index.cmp uge, %workgroup_size, %c512 : index scf.if %wide_tile { - %weight_format = index.constant 6 : index + %resident_layout = config.get @ggml.mul_mat_q6_k_packed.weight_layout : index + %packed_rows_layout = index.constant 2 : index + %reads_packed_rows = index.cmp eq, %resident_layout, %packed_rows_layout : index + %stored_format = index.constant 6 : index + %packed_rows_format = index.constant 46 : index + %weight_format = scf.select %reads_packed_rows, %packed_rows_format, %stored_format : index %paired = scalar.constant false : i1 %prefill_binary_op = index.constant 0 : index func.call @ggml_mul_mat_quantized_f16_prefill_wave32(%weight_format, %prefill_binary_op, %paired, %packed_input, %paired, %bounded_token_count, %input_size, %output_size, %channel_tile, %token_tile, %input, %weight, %weight, %output, %output) : (index, index, i1, i1, i1, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer) @@ -1682,7 +1701,7 @@ kernel.def target(@ggml_mul_mat_q6_k_prefill_gfx11_wave32) export("ggml_mul_mat_ %input_is_f16 = scalar.constant true : i1 // engine #856: 0 = read the loader's stored Q6_K rows (today's behaviour); 1 = read the // repacked int8 tiles the decode kernel already uses, so both share one resident copy. - %weight_layout = index.constant 0 : index + %weight_layout = config.get @ggml.mul_mat_q6_k_packed.weight_layout : index template.apply<@ggml.mul_mat_q6_k_packed.prefill_wave32.body>(%input_is_f16, %weight_layout, %bounded_token_count, %input_size, %output_size, %channel_tile, %token_tile, %input, %weight, %output) : (i1, index, index, index, index, index, index, buffer, buffer, buffer) } kernel.return diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_swiglu_f32_f32_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_swiglu_f32_f32_wmma.loom index e6bc464edbec..e81435ae8cfc 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_swiglu_f32_f32_wmma.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_swiglu_f32_f32_wmma.loom @@ -8,6 +8,7 @@ func.decl @ggml_dot_u8_s8_vector8_f32(%weight: vector<2xi32>, %activation: vecto 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, %input_size: index, %quant_block: index, %quant_group: index, %packet: index, %k: index) -> (vector<4xf32>) 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_packed_row64_f32_vector4(%weight: buffer, %row: index, %input_size: index, %q6_block: index, %q6_group: index, %packet: index) -> (vector<4xf32>) template.decl @ggml.unary_f32.apply_vector4(%op: index, %values: vector<4xf32>) -> (vector<4xf32>) @@ -211,8 +212,17 @@ template.def<@ggml.mul_mat_swiglu.body> device @ggml_mul_mat_swiglu_f32_f32_wmma %packed_f16 = vector.fptrunc %packed_f32 : vector<4xf32> to vector<4xf16> scf.yield %packed_f16 : vector<4xf16> } else { - %native = func.call @ggml_dequant_f16_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %gate_weight_format, %gate_weight_noalias, %row_byte_base, %bounded_input_size, %quant_block, %quant_group, %load_packet, %weight_k) : (vector<16xi8>, 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>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf16>) - scf.yield %native : vector<4xf16> + %packed_q6_format = index.constant 46 : index + %packed_q6_rows = index.cmp eq, %gate_weight_format, %packed_q6_format : index + %stored_or_q6 = scf.if %packed_q6_rows -> (vector<4xf16>) { + %q6_f32 = func.call @ggml_q6k_packed_row64_f32_vector4(%gate_weight_noalias, %channel, %bounded_input_size, %quant_block, %quant_group, %load_packet) : (buffer, index, index, index, index, index) -> (vector<4xf32>) + %q6_f16 = vector.fptrunc %q6_f32 : vector<4xf32> to vector<4xf16> + scf.yield %q6_f16 : vector<4xf16> + } else { + %native = func.call @ggml_dequant_f16_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %gate_weight_format, %gate_weight_noalias, %row_byte_base, %bounded_input_size, %quant_block, %quant_group, %load_packet, %weight_k) : (vector<16xi8>, 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>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf16>) + scf.yield %native : vector<4xf16> + } + scf.yield %stored_or_q6 : vector<4xf16> } scf.yield %decoded : vector<4xf16> } else { @@ -229,8 +239,17 @@ template.def<@ggml.mul_mat_swiglu.body> device @ggml_mul_mat_swiglu_f32_f32_wmma %packed_f16 = vector.fptrunc %packed_f32 : vector<4xf32> to vector<4xf16> scf.yield %packed_f16 : vector<4xf16> } else { - %native = func.call @ggml_dequant_f16_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %up_weight_format, %up_weight_noalias, %row_byte_base, %bounded_input_size, %quant_block, %quant_group, %load_packet, %weight_k) : (vector<16xi8>, 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>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf16>) - scf.yield %native : vector<4xf16> + %packed_q6_format = index.constant 46 : index + %packed_q6_rows = index.cmp eq, %up_weight_format, %packed_q6_format : index + %stored_or_q6 = scf.if %packed_q6_rows -> (vector<4xf16>) { + %q6_f32 = func.call @ggml_q6k_packed_row64_f32_vector4(%up_weight_noalias, %channel, %bounded_input_size, %quant_block, %quant_group, %load_packet) : (buffer, index, index, index, index, index) -> (vector<4xf32>) + %q6_f16 = vector.fptrunc %q6_f32 : vector<4xf32> to vector<4xf16> + scf.yield %q6_f16 : vector<4xf16> + } else { + %native = func.call @ggml_dequant_f16_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %up_weight_format, %up_weight_noalias, %row_byte_base, %bounded_input_size, %quant_block, %quant_group, %load_packet, %weight_k) : (vector<16xi8>, 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>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf16>) + scf.yield %native : vector<4xf16> + } + scf.yield %stored_or_q6 : vector<4xf16> } scf.yield %decoded : vector<4xf16> } else { diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index d85fbbae5bee..73f491326f77 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -9529,6 +9529,19 @@ static std::vector> make_test_cases_eval() { } test_cases.emplace_back(new test_mul_mat_resident_weight(GGML_TYPE_Q4_K, 9728, 164, 2560)); test_cases.emplace_back(new test_mul_mat_resident_weight(GGML_TYPE_Q4_K, 2560, 164, 9728)); + // and q6_K, off the 128 grid its prefill kernel takes: stored by default, packed with GGML_HRX_Q6K_PACKED_ROWS=1 + for (int64_t n : { 22, 100, 164, 300 }) { + test_cases.emplace_back(new test_mul_mat_resident_weight(GGML_TYPE_Q6_K, 2048, n, 4096)); + test_cases.emplace_back(new test_mul_mat_resident_weight(GGML_TYPE_Q6_K, 2048, n, 512)); + } + test_cases.emplace_back(new test_mul_mat_resident_weight(GGML_TYPE_Q6_K, 2560, 164, 2560)); + // q6_K on its prefill kernel's 128 grid but off the 256 one (128, 384 tokens), and an output that is not a multiple + // of 128 at 256 tokens: the tile shapes of that kernel other than the wide one + for (int64_t n : { 128, 384 }) { + test_cases.emplace_back(new test_mul_mat_resident_weight(GGML_TYPE_Q6_K, 2048, n, 4096)); + test_cases.emplace_back(new test_mul_mat_resident_weight(GGML_TYPE_Q6_K, 1024, n, 2560)); + } + test_cases.emplace_back(new test_mul_mat_resident_weight(GGML_TYPE_Q6_K, 2112, 256, 2560)); test_cases.emplace_back(new test_mul_mat_resident_weight(GGML_TYPE_Q6_K, 2048, 512, 4096)); test_cases.emplace_back(new test_mul_mat_resident_weight(GGML_TYPE_Q6_K, 2048, 8, 4096)); test_cases.emplace_back(new test_mul_mat_resident_weight(GGML_TYPE_Q6_K, 4096, 8, 512));