diff --git a/ggml/src/ggml-hrx/CMakeLists.txt b/ggml/src/ggml-hrx/CMakeLists.txt index 6e15ce13931e..1442f4b3b8d9 100644 --- a/ggml/src/ggml-hrx/CMakeLists.txt +++ b/ggml/src/ggml-hrx/CMakeLists.txt @@ -187,6 +187,7 @@ ggml_add_backend_library(ggml-hrx dispatch/transient-allocator.h dispatch_registration/common/dispatch-binary.cpp dispatch_registration/common/dispatch-binary.h + dispatch_registration/common/dispatch-binary-common.h dispatch_registration/common/dispatch-common.cpp dispatch_registration/common/dispatch-common.h dispatch_registration/common/dispatch-copy.cpp diff --git a/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.tg8.json b/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.tg8.json index 9f3d05ccfdcd..36d40a7130f8 100644 --- a/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.tg8.json +++ b/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.tg8.json @@ -130,30 +130,26 @@ }, "kernel": "qwen3_moe:qwen3_moe_attention_qkv_postprocess_fused_decode", "library_sources": [ - "qwen3_moe/attention_postprocess_f32_f16.loom", + "qwen3_moe/motifs/attention_qkv_postprocess_fused_decode.loom", + "../loom-libs/motifs/attention/rmsnorm_neox_cache_f16.loom", "qwen3_moe/model_config.loom", - "qwen3_moe/attention_qkv_quantized.loom", - "qwen3_moe/dense_linear_quantized_f16_wmma.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" + "../loom-libs/motifs/matmul/q4_q6_q8_1_x4_qkv.loom", + "../loom-libs/motifs/matmul/q4_q8_1_x4_decode.loom", + "../loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ], "primary_sources": [ "qwen3_moe/attention_qkv_postprocess_fused.loom" ], "sources": [ "qwen3_moe/attention_qkv_postprocess_fused.loom", - "qwen3_moe/attention_postprocess_f32_f16.loom", + "qwen3_moe/motifs/attention_qkv_postprocess_fused_decode.loom", + "../loom-libs/motifs/attention/rmsnorm_neox_cache_f16.loom", "qwen3_moe/model_config.loom", - "qwen3_moe/attention_qkv_quantized.loom", - "qwen3_moe/dense_linear_quantized_f16_wmma.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" + "../loom-libs/motifs/matmul/q4_q6_q8_1_x4_qkv.loom", + "../loom-libs/motifs/matmul/q4_q8_1_x4_decode.loom", + "../loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ], "symbol": "qwen3_moe_attention_qkv_postprocess_fused_decode", "workload_parameters": [ @@ -228,24 +224,22 @@ }, "kernel": "qwen3_moe:qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8", "library_sources": [ - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" + "qwen3_moe/motifs/attention_prepare_target.loom", + "../loom-libs/motifs/matmul/q4_q8_1_x4_dense.loom", + "../loom-libs/motifs/matmul/q4_q8_1_x4_decode.loom", + "../loom-libs/motifs/normalization/rmsnorm_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ], "primary_sources": [ - "qwen3_moe/dense_linear_quantized_f16_wmma.loom" + "qwen3_moe/motifs/dense_linear_q4k_next_q8.loom" ], "sources": [ - "qwen3_moe/dense_linear_quantized_f16_wmma.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" + "qwen3_moe/motifs/dense_linear_q4k_next_q8.loom", + "qwen3_moe/motifs/attention_prepare_target.loom", + "../loom-libs/motifs/matmul/q4_q8_1_x4_dense.loom", + "../loom-libs/motifs/matmul/q4_q8_1_x4_decode.loom", + "../loom-libs/motifs/normalization/rmsnorm_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ], "symbol": "qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8", "workload_parameters": [ @@ -272,8 +266,9 @@ "kernel": "qwen3_moe:qwen3_moe_router_projection_top8_fused_decode_f32", "library_sources": [ "qwen3_moe/model_config.loom", - "qwen3_moe/router_projection_f32.loom", - "qwen3_moe/router_top8_f32.loom" + "../loom-libs/motifs/routing/router_projection_storage.loom", + "../loom-libs/motifs/routing/router_topk_row.loom", + "../loom-libs/motifs/routing/router_projection_fused.loom" ], "primary_sources": [ "qwen3_moe/router_projection_top8_fused_f32.loom" @@ -281,8 +276,9 @@ "sources": [ "qwen3_moe/router_projection_top8_fused_f32.loom", "qwen3_moe/model_config.loom", - "qwen3_moe/router_projection_f32.loom", - "qwen3_moe/router_top8_f32.loom" + "../loom-libs/motifs/routing/router_projection_storage.loom", + "../loom-libs/motifs/routing/router_topk_row.loom", + "../loom-libs/motifs/routing/router_projection_fused.loom" ], "symbol": "qwen3_moe_router_projection_top8_fused_decode_f32", "workload_parameters": [ @@ -387,7 +383,7 @@ "motifs/matmul_id/vector_weighted.loom", "motifs/matmul_id/vector_q4k_q8.loom", "motifs/matmul/q6_f32_decode.loom", - "../qwen_moe/qwen3_moe/attention_prepare_quantized.loom" + "motifs/normalization/rmsnorm_q8_1_x4.loom" ], "primary_sources": [ "ops/matmul_id/vector_weighted.loom" @@ -397,7 +393,7 @@ "motifs/matmul_id/vector_weighted.loom", "motifs/matmul_id/vector_q4k_q8.loom", "motifs/matmul/q6_f32_decode.loom", - "../qwen_moe/qwen3_moe/attention_prepare_quantized.loom" + "motifs/normalization/rmsnorm_q8_1_x4.loom" ], "symbol": "ggml_mul_mat_id_vector_weighted_wave64_publish_q8", "workload_parameters": [ @@ -447,30 +443,26 @@ }, "kernel": "qwen3_moe:qwen3_moe_attention_qkv_postprocess_fused_decode", "library_sources": [ - "qwen3_moe/attention_postprocess_f32_f16.loom", + "qwen3_moe/motifs/attention_qkv_postprocess_fused_decode.loom", + "../loom-libs/motifs/attention/rmsnorm_neox_cache_f16.loom", "qwen3_moe/model_config.loom", - "qwen3_moe/attention_qkv_quantized.loom", - "qwen3_moe/dense_linear_quantized_f16_wmma.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" + "../loom-libs/motifs/matmul/q4_q6_q8_1_x4_qkv.loom", + "../loom-libs/motifs/matmul/q4_q8_1_x4_decode.loom", + "../loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ], "primary_sources": [ "qwen3_moe/attention_qkv_postprocess_fused.loom" ], "sources": [ "qwen3_moe/attention_qkv_postprocess_fused.loom", - "qwen3_moe/attention_postprocess_f32_f16.loom", + "qwen3_moe/motifs/attention_qkv_postprocess_fused_decode.loom", + "../loom-libs/motifs/attention/rmsnorm_neox_cache_f16.loom", "qwen3_moe/model_config.loom", - "qwen3_moe/attention_qkv_quantized.loom", - "qwen3_moe/dense_linear_quantized_f16_wmma.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" + "../loom-libs/motifs/matmul/q4_q6_q8_1_x4_qkv.loom", + "../loom-libs/motifs/matmul/q4_q8_1_x4_decode.loom", + "../loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ], "symbol": "qwen3_moe_attention_qkv_postprocess_fused_decode", "workload_parameters": [ @@ -575,7 +567,7 @@ "motifs/matmul_id/vector_weighted.loom", "motifs/matmul_id/vector_q4k_q8.loom", "motifs/matmul/q6_f32_decode.loom", - "../qwen_moe/qwen3_moe/attention_prepare_quantized.loom" + "motifs/normalization/rmsnorm_q8_1_x4.loom" ], "primary_sources": [ "ops/matmul_id/vector_weighted.loom" @@ -585,7 +577,7 @@ "motifs/matmul_id/vector_weighted.loom", "motifs/matmul_id/vector_q4k_q8.loom", "motifs/matmul/q6_f32_decode.loom", - "../qwen_moe/qwen3_moe/attention_prepare_quantized.loom" + "motifs/normalization/rmsnorm_q8_1_x4.loom" ], "symbol": "ggml_mul_mat_id_vector_weighted_wave32_publish_q8", "workload_parameters": [ @@ -959,16 +951,22 @@ }, "kernel": "qwen3_moe:qwen3_moe_rmsnorm_f32_quantize_q8_1_x4", "library_sources": [ + "qwen3_moe/motifs/attention_prepare_target.loom", + "../loom-libs/motifs/normalization/rmsnorm_q8_1_x4.loom", "qwen3_moe/model_config.loom", - "ggml/quantize_q8_1_x4.loom" + "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ], "primary_sources": [ "qwen3_moe/attention_prepare_quantized.loom" ], "sources": [ "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/motifs/attention_prepare_target.loom", + "../loom-libs/motifs/normalization/rmsnorm_q8_1_x4.loom", "qwen3_moe/model_config.loom", - "ggml/quantize_q8_1_x4.loom" + "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ], "symbol": "qwen3_moe_rmsnorm_f32_quantize_q8_1_x4", "workload_parameters": [ @@ -993,14 +991,20 @@ }, "kernel": "qwen3_moe:ggml_linear_q6k_q8_1_x4", "library_sources": [ - "ggml/quantize_q8_1_x4.loom" + "../loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom", + "../loom-libs/motifs/q6_k_f16.loom", + "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ], "primary_sources": [ "ggml/linear_q6k_q8_1_x4.loom" ], "sources": [ "ggml/linear_q6k_q8_1_x4.loom", - "ggml/quantize_q8_1_x4.loom" + "../loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom", + "../loom-libs/motifs/q6_k_f16.loom", + "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ], "symbol": "ggml_linear_q6k_q8_1_x4", "workload_parameters": [ diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-binary-common.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-binary-common.h new file mode 100644 index 000000000000..631747896fe3 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-binary-common.h @@ -0,0 +1,154 @@ +#pragma once + +#include "../dispatch-registry.h" +#include "graph/op-params.h" + +#include +#include + +namespace ggml::hrx { + +inline constexpr uint32_t common_binary_kind_bit(BinaryKind kind) { + return uint32_t{ 1 } << static_cast(kind); +} + +inline constexpr uint32_t kCommonArithmeticBinaryKinds = + common_binary_kind_bit(BinaryKind::Add) | common_binary_kind_bit(BinaryKind::Sub) | + common_binary_kind_bit(BinaryKind::Mul) | common_binary_kind_bit(BinaryKind::Div); + +inline constexpr uint32_t kCommonGatedBinaryKinds = + common_binary_kind_bit(BinaryKind::SwiGLU) | common_binary_kind_bit(BinaryKind::GeGLU) | + common_binary_kind_bit(BinaryKind::RegLU) | common_binary_kind_bit(BinaryKind::GeGLUErf) | + common_binary_kind_bit(BinaryKind::GeGLUQuick); + +inline bool common_binary_kind_allowed(BinaryKind kind, uint32_t allowed_kinds) { + return binary_kind_supported(kind) && (allowed_kinds & common_binary_kind_bit(kind)) != 0; +} + +inline bool common_binary_kind_requires_order(BinaryKind kind) { + return kind == BinaryKind::Sub || kind == BinaryKind::Div; +} + +inline bool common_binary_kind_matches_graph_op(BinaryKind kind, ggml_op op) { + switch (kind) { + case BinaryKind::Add: + return op == GGML_OP_ADD; + case BinaryKind::Sub: + return op == GGML_OP_SUB; + case BinaryKind::Mul: + return op == GGML_OP_MUL; + case BinaryKind::Div: + return op == GGML_OP_DIV; + case BinaryKind::SwiGLU: + case BinaryKind::GeGLU: + case BinaryKind::RegLU: + case BinaryKind::GeGLUErf: + case BinaryKind::GeGLUQuick: + return op == GGML_OP_GLU; + } + return false; +} + +inline bool common_arithmetic_kind_from_graph_op(ggml_op op, BinaryKind & kind) { + switch (op) { + case GGML_OP_ADD: + kind = BinaryKind::Add; + return true; + case GGML_OP_SUB: + kind = BinaryKind::Sub; + return true; + case GGML_OP_MUL: + kind = BinaryKind::Mul; + return true; + case GGML_OP_DIV: + kind = BinaryKind::Div; + return true; + default: + return false; + } +} + +inline bool common_fused_binary_kind_from_params(const OpParams & params, BinaryKind & kind) { + const BinaryParams * binary_params = op_params_as(params); + if (binary_params != nullptr && binary_kind_supported(binary_params->op)) { + kind = binary_params->op; + return true; + } + + const GluParams * glu_params = op_params_as(params); + if (glu_params == nullptr) { + return false; + } + + switch (glu_params->op) { + case GGML_GLU_OP_REGLU: + kind = BinaryKind::RegLU; + return true; + case GGML_GLU_OP_SWIGLU: + kind = BinaryKind::SwiGLU; + return true; + case GGML_GLU_OP_GEGLU: + kind = BinaryKind::GeGLU; + return true; + case GGML_GLU_OP_GEGLU_ERF: + kind = BinaryKind::GeGLUErf; + return true; + case GGML_GLU_OP_GEGLU_QUICK: + kind = BinaryKind::GeGLUQuick; + return true; + default: + return false; + } +} + +struct CommonBinaryMatch { + const GraphNode * node = nullptr; + ValueId operand = ValueId(-1); + BinaryKind kind = BinaryKind::Add; + bool producer_is_lhs = true; + + bool matched() const { return node != nullptr && operand.value >= 0; } +}; + +inline CommonBinaryMatch common_match_binary_consumer(const GraphNode * node, + ValueId producer, + uint32_t allowed_kinds) { + CommonBinaryMatch match; + if (node == nullptr || node->inputs.size() != 2) { + return match; + } + + BinaryKind kind; + const bool has_kind = + common_fused_binary_kind_from_params(node->params, kind) || + (std::holds_alternative(node->params) && common_arithmetic_kind_from_graph_op(node->op, kind)); + if (!has_kind || !common_binary_kind_matches_graph_op(kind, node->op) || + !common_binary_kind_allowed(kind, allowed_kinds)) { + return match; + } + + const bool producer_is_lhs = node->inputs[0] == producer; + const bool producer_is_rhs = node->inputs[1] == producer; + if (producer_is_lhs == producer_is_rhs) { + return match; + } + + match.node = node; + match.operand = producer_is_lhs ? node->inputs[1] : node->inputs[0]; + match.kind = kind; + match.producer_is_lhs = producer_is_lhs; + return match; +} + +inline void common_set_binary_compile_parameters(KernelSpecialization & kernel, + const char * op_key, + const char * lhs_key, + BinaryKind kind, + bool producer_is_lhs) { + kernel.compile_parameters.emplace(op_key, std::to_string(binary_kind_config_value(kind))); + if (lhs_key != nullptr) { + kernel.compile_parameters.emplace(lhs_key, producer_is_lhs ? "1" : "0"); + } +} + +} // namespace ggml::hrx 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 6a3912f4d1bf..496938e8fb19 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 @@ -21,10 +21,6 @@ static constexpr KernelCatalogRef kMulMatTiledPairF32BinaryPublishF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32"); static constexpr KernelCatalogRef kMulMatTiledPairF32BinaryPublishF32K16Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32_k16"); -static constexpr KernelCatalogRef kMulMatTiledPairF32BinaryBiasPublishF32Kernel = - GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_pair_input_f32_binary_bias_publish_f32"); -static constexpr KernelCatalogRef kMulMatTiledPairF32BinaryAddPublishF32Kernel = - GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_pair_input_f32_binary_residual_publish_f32"); static constexpr KernelCatalogRef kMulMatTiledPairF32BinaryBiasAddPublishF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_pair_input_f32_binary_bias_residual_publish_f32"); static constexpr KernelCatalogRef kMulMatSwiGLUF32F32DecodeWave64Kernel = @@ -287,8 +283,8 @@ static MulMatSwiGLUMatch match_mul_mat_swiglu(const DispatchMatchContext & conte // gate/up halves and lowered through the split-weight gated matmul kernel. static PackedMulMatGluMatch match_packed_mul_mat_glu(const DispatchMatchContext & context) { PackedMulMatGluMatch match; - const CommonMulMatMatch root = - common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatTiledPairF32BinaryPublishF32Kernel, false); + const CommonMulMatMatch root = common_match_mul_mat_any_format(context.graph, context.root_node, + kMulMatTiledPairF32BinaryPublishF32Kernel, false); if (!root.matched() || !context.graph.has_index() || root.weight->alias_source.value >= 0) { return {}; } @@ -607,22 +603,21 @@ static bool match_mul_mat_swiglu_q4_q8_prefill_dispatch(const DispatchMatchConte DispatchMatch & dispatch_match) { const MulMatSwiGLUMatch match = match_mul_mat_swiglu(context); if (!match.matched() || match.gate_format != CommonMulMatWeightFormat::Q4K || - match.up_format != CommonMulMatWeightFormat::Q4K || match.token_count < 256 || - match.token_count > 2048 || match.token_count % 256 != 0 || match.input_size % 256 != 0 || - match.output_size % 64 != 0 || match.gate_weight->alias_source.value >= 0 || - match.up_weight->alias_source.value >= 0 || + match.up_format != CommonMulMatWeightFormat::Q4K || match.token_count < 256 || match.token_count > 2048 || + match.token_count % 256 != 0 || match.input_size % 256 != 0 || match.output_size % 64 != 0 || + match.gate_weight->alias_source.value >= 0 || match.up_weight->alias_source.value >= 0 || !distinct_storage(context.graph, *match.gate_weight, *match.up_weight) || !distinct_storage(context.graph, *match.gate_output, *match.up_output)) { return false; } // Keep in sync with ggml_swiglu_use_packed_f16 in the shared operation. - const bool packed_input = match.token_count % 512 == 0 && - (match.token_count / 512) * (match.output_size / 64) >= 64; + const bool packed_input = + match.token_count % 512 == 0 && (match.token_count / 512) * (match.output_size / 64) >= 64; DispatchBinding activation; const bool prepared = packed_input ? - common_prepare_k16_major_f16_input(context, *match.input, match.input_size, match.token_count, - dispatch_match, activation) : + common_prepare_k16_major_f16_input(context, *match.input, match.input_size, + match.token_count, dispatch_match, activation) : common_prepare_f16_input(context, *match.input, match.input_size, match.token_count, dispatch_match, activation); if (!prepared) { @@ -640,8 +635,8 @@ static bool match_mul_mat_swiglu_q4_q8_prefill_dispatch(const DispatchMatchConte Dispatch dispatch; dispatch.kernel = make_kernel_specialization(kMulMatSwiGLUQ4Q8PrefillKernel); dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_swiglu.f16_output_layout", packed_output ? "1" : "0"); - dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_swiglu.op", - common_to_config_value(static_cast(binary_kind_config_value(match.op)))); + dispatch.kernel.compile_parameters.emplace( + "ggml.mul_mat_swiglu.op", common_to_config_value(static_cast(binary_kind_config_value(match.op)))); dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_swiglu.input_size", common_to_config_value(match.input_size)); @@ -661,7 +656,8 @@ static bool match_mul_mat_swiglu_q4_q8_prefill_dispatch(const DispatchMatchConte dispatch_match.transients.push_back({ f16_output, f16_name, f16_bytes, 256 }); dispatch.bindings.push_back({ f16_output, 0, f16_bytes }); Status status; - const bool recorded = packed_output ? + const bool recorded = + packed_output ? dispatch_match.metadata.append_generated_resource( { match.output->id, GeneratedResourceRole::F16K16Major, f16_output, f16_bytes, {} }, status) : dispatch_match.metadata.append_alternate_value( @@ -686,8 +682,7 @@ static bool supports_swiglu_q8_output_shape(const MulMatSwiGLUMatch & match) { } static bool has_qualified_swiglu_q8_consumer(const Graph & graph, const MulMatSwiGLUMatch & match) { - if (!supports_swiglu_q8_output_shape(match) || - !distinct_storage(graph, *match.input, *match.gate_weight) || + if (!supports_swiglu_q8_output_shape(match) || !distinct_storage(graph, *match.input, *match.gate_weight) || !distinct_storage(graph, *match.input, *match.up_weight) || !distinct_storage(graph, *match.input, *match.output) || !distinct_storage(graph, *match.gate_weight, *match.up_weight) || @@ -736,8 +731,8 @@ static bool has_qualified_swiglu_k16_consumer(const Graph & graph, const MulMatS common_match_mul_mat_any_format(graph, consumer, kMulMatTiledPairF32BinaryPublishF32Kernel, false); if (projection.matched() && projection.input->id == match.output->id && projection.weight->alias_source.value < 0 && - common_mul_mat_uses_k16_major_f16(projection.weight_format, projection.input_size, - projection.output_size, projection.token_count, true)) { + common_mul_mat_uses_k16_major_f16(projection.weight_format, projection.input_size, projection.output_size, + projection.token_count, true)) { return true; } } @@ -800,13 +795,7 @@ static MulMatSwiGLUPostOpsMatch match_mul_mat_swiglu_postops(const DispatchMatch } match.residual_output = current; - if (match.has_bias && match.has_residual) { match.kernel = kMulMatTiledPairF32BinaryBiasAddPublishF32Kernel; - } else if (match.has_residual) { - match.kernel = kMulMatTiledPairF32BinaryAddPublishF32Kernel; - } else if (match.has_bias) { - match.kernel = kMulMatTiledPairF32BinaryBiasPublishF32Kernel; - } return match; } @@ -836,18 +825,28 @@ static bool build_tiled_pair_mul_mat_swiglu_dispatch(const DispatchMatchContext dispatch.kernel.compile_parameters.emplace( "ggml.mul_mat_swiglu.up_weight_format", common_to_config_value(common_mul_mat_format_config_value(match.up_format))); - dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_swiglu.op", - common_to_config_value(static_cast(binary_kind_config_value(match.op)))); + dispatch.kernel.compile_parameters.emplace( + "ggml.mul_mat_swiglu.op", common_to_config_value(static_cast(binary_kind_config_value(match.op)))); + if (postops != nullptr) { + dispatch.kernel.compile_parameters.emplace( + "ggml.mul_mat_swiglu.epilogue", + common_to_config_value((postops->has_bias ? int64_t{ 1 } : int64_t{ 0 }) | + (postops->has_residual ? int64_t{ 2 } : int64_t{ 0 }))); + } dispatch.bindings.push_back({ match.input->id, 0, match.input->byte_count }); 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 }); - if (postops != nullptr && postops->has_bias) { - dispatch.bindings.push_back({ postops->bias->id, 0, postops->bias->byte_count }); - } - if (postops != nullptr && postops->has_residual) { - dispatch.bindings.push_back({ postops->residual_input->id, 0, postops->residual_input->byte_count }); - } const Value * output = postops != nullptr ? postops->residual_output : match.output; + if (postops != nullptr) { + const Value * placeholder = postops->has_bias ? postops->bias : postops->residual_input; + dispatch.bindings.push_back(postops->has_bias ? + DispatchBinding{ postops->bias->id, 0, postops->bias->byte_count } : + DispatchBinding{ placeholder->id, 0, placeholder->byte_count }); + dispatch.bindings.push_back( + postops->has_residual ? + DispatchBinding{ postops->residual_input->id, 0, postops->residual_input->byte_count } : + DispatchBinding{ placeholder->id, 0, placeholder->byte_count }); + } dispatch.bindings.push_back({ output->id, 0, output->byte_count }); if (publish_k16) { @@ -913,9 +912,9 @@ static bool match_mul_mat_swiglu_dispatch(const DispatchMatchContext & context, const bool use_q8 = pack_q4; const bool publish_q8 = use_q8 && has_qualified_swiglu_q8_consumer(context.graph, match); DispatchBinding activation = { match.input->id, 0, match.input->byte_count }; - if (use_q8 && !common_prepare_q8_1_x4_input(context, *match.input, match.input_size, match.token_count, - dispatch_match, activation, - CommonQ8ActivationPolicy::AllowStandaloneQuantize)) { + if (use_q8 && + !common_prepare_q8_1_x4_input(context, *match.input, match.input_size, match.token_count, dispatch_match, + activation, CommonQ8ActivationPolicy::AllowStandaloneQuantize)) { return false; } @@ -925,7 +924,8 @@ static bool match_mul_mat_swiglu_dispatch(const DispatchMatchContext & context, if (use_wmma && match.token_count < 2) { return false; } - dispatch.kernel = make_kernel_specialization(publish_q8 ? kMulMatSwiGLUQ4Q8OutputKernel : + dispatch.kernel = make_kernel_specialization( + publish_q8 ? kMulMatSwiGLUQ4Q8OutputKernel : use_q8 ? kMulMatSwiGLUQ4Q8LowTokenDotKernel : (use_direct_dot ? kMulMatSwiGLUF32F32LowTokenDotKernel : kMulMatSwiGLUF32F32WmmaKernel)); dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); @@ -935,16 +935,14 @@ static bool match_mul_mat_swiglu_dispatch(const DispatchMatchContext & context, common_to_config_value(match.input_size)); dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_swiglu.output_size", 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( + pack_q4 ? CommonMulMatWeightFormat::Q4KRow64 : match.gate_format))); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_swiglu.up_weight_format", + common_to_config_value(common_mul_mat_format_config_value( + pack_q4 ? CommonMulMatWeightFormat::Q4KRow64 : match.up_format))); dispatch.kernel.compile_parameters.emplace( - "ggml.mul_mat_swiglu.gate_weight_format", - common_to_config_value(common_mul_mat_format_config_value(pack_q4 ? CommonMulMatWeightFormat::Q4KRow64 : - match.gate_format))); - dispatch.kernel.compile_parameters.emplace( - "ggml.mul_mat_swiglu.up_weight_format", - common_to_config_value(common_mul_mat_format_config_value(pack_q4 ? CommonMulMatWeightFormat::Q4KRow64 : - match.up_format))); - dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_swiglu.op", - common_to_config_value(static_cast(binary_kind_config_value(match.op)))); + "ggml.mul_mat_swiglu.op", common_to_config_value(static_cast(binary_kind_config_value(match.op)))); dispatch.bindings.push_back(activation); if (pack_q4) { for (const Value * weight : { match.gate_weight, match.up_weight }) { diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gather-add.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gather-add.cpp index a1992ba6eb29..d7a49803fa75 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gather-add.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gather-add.cpp @@ -1,15 +1,24 @@ #include "dispatch-gather-add.h" +#include "dispatch-binary-common.h" +#include "dispatch-mul-mat-common.h" #include "ggml.h" +#include "graph/op-params.h" #include "kernel-corpus/kernel-corpus-catalog-verify.h" +#include +#include #include +#include +#include #include namespace ggml::hrx { namespace { static constexpr KernelCatalogRef kGatherAddF32Kernel = GGML_HRX_KERNEL_REF("hrx", "ggml_gather_add_f32"); +static constexpr KernelCatalogRef kGatherAddRmsNormBinaryF32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_gather_add_rmsnorm_binary_f32"); static const Value * graph_value(const Graph & graph, ValueId id) { return graph.values().find(id); @@ -66,6 +75,7 @@ struct GatherAddMatch { int64_t source_token_count = 0; int64_t output_token_count = 0; int64_t hidden_size = 0; + BinaryKind op = BinaryKind::Add; bool matched() const { return first_get_rows != nullptr && second_get_rows != nullptr && add_node != nullptr && @@ -73,13 +83,30 @@ struct GatherAddMatch { } }; -static const GraphNode * find_single_add_consumer(const Graph & graph, const GraphNode & get_rows) { +struct GatherAddRmsNormBinaryMatch { + GatherAddMatch gather_add; + const GraphNode *rms_node = nullptr; + const GraphNode *binary_node = nullptr; + const Value * weight = nullptr; + const Value * normalized_output = nullptr; + size_t rms_node_index = 0; + size_t binary_node_index = 0; + float epsilon = 0.0f; + + bool matched() const { + return gather_add.matched() && rms_node != nullptr && binary_node != nullptr && weight != nullptr && + normalized_output != nullptr; + } +}; + +static const GraphNode * find_single_binary_consumer(const Graph & graph, const GraphNode & get_rows) { const std::vector & consumers = graph.index().consumers(get_rows.output); if (consumers.size() != 1) { return nullptr; } - const GraphNode * add = consumers.front(); - return add != nullptr && add->op == GGML_OP_ADD && add->inputs.size() == 2 ? add : nullptr; + const GraphNode * binary = consumers.front(); + return common_match_binary_consumer(binary, get_rows.output, kCommonArithmeticBinaryKinds).matched() ? binary : + nullptr; } static const GraphNode * peer_get_rows_input(const Graph & graph, @@ -104,7 +131,7 @@ static GatherAddMatch match_gather_add_f32(const Graph & graph, const GraphNode return match; } - const GraphNode * add_node = find_single_add_consumer(graph, *node); + const GraphNode * add_node = find_single_binary_consumer(graph, *node); const GraphNode * peer = add_node != nullptr ? peer_get_rows_input(graph, *add_node, *node) : nullptr; if (add_node == nullptr || peer == nullptr) { return {}; @@ -116,18 +143,22 @@ static GatherAddMatch match_gather_add_f32(const Graph & graph, const GraphNode return {}; } - const Value * first_source = graph_value(graph, node->inputs[0]); - const Value * first_ids = graph_value(graph, node->inputs[1]); - const Value * first_output = graph_value(graph, node->output); - const Value * second_source = graph_value(graph, peer->inputs[0]); - const Value * second_ids = graph_value(graph, peer->inputs[1]); - const Value * second_output = graph_value(graph, peer->output); + const GraphNode * first_get_rows = add_node->inputs[0] == node->output ? node : peer; + const GraphNode * second_get_rows = first_get_rows == node ? peer : node; + const size_t first_get_rows_index = first_get_rows == node ? node_index : peer_index; + const size_t second_get_rows_index = first_get_rows == node ? peer_index : node_index; + const Value * first_source = graph_value(graph, first_get_rows->inputs[0]); + const Value * first_ids = graph_value(graph, first_get_rows->inputs[1]); + const Value * first_output = graph_value(graph, first_get_rows->output); + const Value * second_source = graph_value(graph, second_get_rows->inputs[0]); + const Value * second_ids = graph_value(graph, second_get_rows->inputs[1]); + const Value * second_output = graph_value(graph, second_get_rows->output); const Value * output = graph_value(graph, add_node->output); if (first_source == nullptr || first_ids == nullptr || first_output == nullptr || second_source == nullptr || second_ids == nullptr || second_output == nullptr || output == nullptr) { return {}; } - if (node->inputs[1] != peer->inputs[1]) { + if (first_get_rows->inputs[1] != second_get_rows->inputs[1]) { return {}; } if (!is_2d_f32(*first_source) || !is_2d_f32(*second_source) || !is_2d_f32(*first_output) || @@ -151,24 +182,193 @@ static GatherAddMatch match_gather_add_f32(const Graph & graph, const GraphNode return {}; } - match.first_get_rows = node; - match.second_get_rows = peer; + const CommonBinaryMatch binary = + common_match_binary_consumer(add_node, first_get_rows->output, kCommonArithmeticBinaryKinds); + if (!binary.matched() || !binary.producer_is_lhs) { + return {}; + } + + match.first_get_rows = first_get_rows; + match.second_get_rows = second_get_rows; match.add_node = add_node; match.first_source = first_source; match.second_source = second_source; match.row_ids = first_ids; match.output = output; - match.first_get_rows_index = node_index; - match.second_get_rows_index = peer_index; + match.first_get_rows_index = first_get_rows_index; + match.second_get_rows_index = second_get_rows_index; match.add_node_index = add_index; match.source_token_count = source_token_count; match.output_token_count = output_token_count; match.hidden_size = hidden_size; + match.op = binary.kind; + return match; +} + +static bool is_packed_q8_consumer(const Graph & graph, const GraphNode * consumer, const Value & input) { + if (consumer == nullptr || consumer->op != GGML_OP_MUL_MAT || consumer->inputs.size() != 2 || + consumer->inputs[1] != input.id) { + return false; + } + const Value * weight = graph_value(graph, consumer->inputs[0]); + const Value * output = graph_value(graph, consumer->output); + if (weight == nullptr || output == nullptr || input.type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32 || + !input.contiguous || !weight->contiguous || !output->contiguous) { + return false; + } + const bool prefill = (weight->type == GGML_TYPE_Q5_K || weight->type == GGML_TYPE_IQ4_XS) && input.ne[1] >= 256 && + input.ne[1] <= 2048 && input.ne[1] % 256 == 0; + const bool decode = input.ne[1] >= 1 && input.ne[1] <= 5 && weight->alias_source.value < 0 && + (weight->type == GGML_TYPE_Q4_K || + (weight->type == GGML_TYPE_Q6_K && !graph.index().consumers(output->id).empty())); + return (prefill || decode) && input.ne[0] >= 256 && input.ne[0] <= 32768 && input.ne[0] % 256 == 0 && + input.ne[2] == 1 && input.ne[3] == 1 && weight->ne[0] == input.ne[0] && weight->ne[1] >= 64 && + weight->ne[1] <= 262144 && weight->ne[1] % 64 == 0 && weight->ne[2] == 1 && weight->ne[3] == 1 && + output->ne[0] == weight->ne[1] && output->ne[1] == input.ne[1] && output->ne[2] == 1 && output->ne[3] == 1; +} + +static bool has_packed_q8_consumer(const Graph & graph, const Value & value) { + for (const GraphNode * consumer : graph.index().consumers(value.id)) { + if (is_packed_q8_consumer(graph, consumer, value)) { + return true; + } + } + return false; +} + +template static bool pairwise_distinct_storage_roots(const std::array & values) { + for (size_t lhs = 0; lhs < values.size(); ++lhs) { + for (size_t rhs = lhs + 1; rhs < values.size(); ++rhs) { + if (values[lhs]->storage_root == values[rhs]->storage_root) { + return false; + } + } + } + return true; +} + +static GatherAddRmsNormBinaryMatch match_gather_add_rmsnorm_binary_f32(const Graph & graph, + const GraphNode * node, + size_t node_index) { + GatherAddRmsNormBinaryMatch match; + match.gather_add = match_gather_add_f32(graph, node, node_index); + if (!match.gather_add.matched()) { + return {}; + } + + const GraphNode * rms_node = nullptr; + for (const GraphNode * consumer : graph.index().consumers(match.gather_add.output->id)) { + if (consumer != nullptr && consumer->op == GGML_OP_RMS_NORM && consumer->inputs.size() == 1) { + if (rms_node != nullptr) { + return {}; + } + rms_node = consumer; + } + } + size_t rms_node_index = 0; + if (rms_node == nullptr || !graph.index().node_index(rms_node, rms_node_index)) { + return {}; + } + const RmsNormParams * rms_params = op_params_as(rms_node->params); + const Value * rms_output = graph_value(graph, rms_node->output); + if (rms_params == nullptr || !std::isfinite(rms_params->eps) || rms_params->eps <= 0.0f || rms_output == nullptr || + rms_output->type != GGML_TYPE_F32 || !rms_output->contiguous || + !same_shape(*match.gather_add.output, *rms_output)) { + return {}; + } + const std::vector & rms_consumers = graph.index().consumers(rms_output->id); + if (rms_consumers.size() != 1 || rms_consumers.front() == nullptr || rms_consumers.front()->op != GGML_OP_MUL || + rms_consumers.front()->inputs.size() != 2) { + return {}; + } + const GraphNode * binary_node = rms_consumers.front(); + size_t binary_node_index = 0; + if (!graph.index().node_index(binary_node, binary_node_index)) { + return {}; + } + const ValueId weight_id = binary_node->inputs[0] == rms_output->id ? binary_node->inputs[1] : + binary_node->inputs[1] == rms_output->id ? binary_node->inputs[0] : + ValueId{}; + const Value * weight = graph_value(graph, weight_id); + const Value * normalized_output = graph_value(graph, binary_node->output); + if (weight == nullptr || normalized_output == nullptr || weight->type != GGML_TYPE_F32 || + normalized_output->type != GGML_TYPE_F32 || !weight->contiguous || !normalized_output->contiguous || + weight->ne[0] != match.gather_add.hidden_size || weight->ne[1] != 1 || weight->ne[2] != 1 || + weight->ne[3] != 1 || !same_shape(*match.gather_add.output, *normalized_output) || + has_packed_q8_consumer(graph, *normalized_output) || + common_has_symmetric_i4_lowrow_consumer(graph, *normalized_output) || + !pairwise_distinct_storage_roots(std::array{ + match.gather_add.first_source, match.gather_add.second_source, match.gather_add.row_ids, + match.gather_add.output, weight, normalized_output })) { + return {}; + } + + match.rms_node = rms_node; + match.binary_node = binary_node; + match.weight = weight; + match.normalized_output = normalized_output; + match.rms_node_index = rms_node_index; + match.binary_node_index = binary_node_index; + match.epsilon = rms_params->eps; return match; } +static std::string to_config_value(float value) { + std::ostringstream out; + out.precision(9); + out << value; + return out.str(); +} + } // namespace +static bool match_gather_add_rmsnorm_binary_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { + const GatherAddRmsNormBinaryMatch fused = + match_gather_add_rmsnorm_binary_f32(context.graph, context.root_node, context.root_index); + if (!fused.matched() || fused.gather_add.first_get_rows_index >= context.covered_nodes.size() || + fused.gather_add.second_get_rows_index >= context.covered_nodes.size() || + fused.gather_add.add_node_index >= context.covered_nodes.size() || + fused.rms_node_index >= context.covered_nodes.size() || + fused.binary_node_index >= context.covered_nodes.size() || + context.covered_nodes[fused.gather_add.first_get_rows_index] || + context.covered_nodes[fused.gather_add.second_get_rows_index] || + context.covered_nodes[fused.gather_add.add_node_index] || context.covered_nodes[fused.rms_node_index] || + context.covered_nodes[fused.binary_node_index]) { + return false; + } + if (!value_is_available(context.graph, fused.gather_add.first_source->id, context.covered_nodes) || + !value_is_available(context.graph, fused.gather_add.second_source->id, context.covered_nodes) || + !value_is_available(context.graph, fused.gather_add.row_ids->id, context.covered_nodes) || + !value_is_available(context.graph, fused.weight->id, context.covered_nodes)) { + return false; + } + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kGatherAddRmsNormBinaryF32Kernel); + dispatch.kernel.integer_parameters.emplace("source_token_count", fused.gather_add.source_token_count); + dispatch.kernel.integer_parameters.emplace("output_token_count", fused.gather_add.output_token_count); + dispatch.kernel.compile_parameters.emplace("ggml.gather_add_rmsnorm_binary_f32.hidden_size", + std::to_string(fused.gather_add.hidden_size)); + dispatch.kernel.compile_parameters.emplace("ggml.gather_add_rmsnorm_binary_f32.rms_epsilon", + to_config_value(fused.epsilon)); + common_set_binary_compile_parameters(dispatch.kernel, "ggml.gather_add_rmsnorm_binary_f32.gather_op", nullptr, + fused.gather_add.op, true); + dispatch.bindings.push_back({ fused.gather_add.first_source->id, 0, fused.gather_add.first_source->byte_count }); + dispatch.bindings.push_back({ fused.gather_add.second_source->id, 0, fused.gather_add.second_source->byte_count }); + dispatch.bindings.push_back({ fused.gather_add.row_ids->id, 0, fused.gather_add.row_ids->byte_count }); + dispatch.bindings.push_back({ fused.gather_add.output->id, 0, fused.gather_add.output->byte_count }); + dispatch.bindings.push_back({ fused.weight->id, 0, fused.weight->byte_count }); + dispatch.bindings.push_back({ fused.normalized_output->id, 0, fused.normalized_output->byte_count }); + + match.covered_nodes.push_back(fused.gather_add.first_get_rows_index); + match.covered_nodes.push_back(fused.gather_add.second_get_rows_index); + match.covered_nodes.push_back(fused.gather_add.add_node_index); + match.covered_nodes.push_back(fused.rms_node_index); + match.covered_nodes.push_back(fused.binary_node_index); + match.dispatches.push_back(std::move(dispatch)); + return true; +} + static bool match_gather_add_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { const GatherAddMatch gather_add = match_gather_add_f32(context.graph, context.root_node, context.root_index); if (!gather_add.matched() || gather_add.first_get_rows_index >= context.covered_nodes.size() || @@ -179,7 +379,8 @@ static bool match_gather_add_f32_dispatch(const DispatchMatchContext & context, return false; } if (!value_is_available(context.graph, gather_add.first_source->id, context.covered_nodes) || - !value_is_available(context.graph, gather_add.second_source->id, context.covered_nodes)) { + !value_is_available(context.graph, gather_add.second_source->id, context.covered_nodes) || + !value_is_available(context.graph, gather_add.row_ids->id, context.covered_nodes)) { return false; } @@ -188,6 +389,8 @@ static bool match_gather_add_f32_dispatch(const DispatchMatchContext & context, dispatch.kernel.integer_parameters.emplace("source_token_count", gather_add.source_token_count); dispatch.kernel.integer_parameters.emplace("output_token_count", gather_add.output_token_count); dispatch.kernel.integer_parameters.emplace("hidden_size", gather_add.hidden_size); + common_set_binary_compile_parameters(dispatch.kernel, "ggml.gather_add_f32.binary_op", nullptr, gather_add.op, + true); dispatch.bindings.push_back({ gather_add.first_source->id, 0, gather_add.first_source->byte_count }); dispatch.bindings.push_back({ gather_add.second_source->id, 0, gather_add.second_source->byte_count }); dispatch.bindings.push_back({ gather_add.row_ids->id, 0, gather_add.row_ids->byte_count }); @@ -201,6 +404,14 @@ static bool match_gather_add_f32_dispatch(const DispatchMatchContext & context, } void register_gather_add_dispatch(DispatchRegistryBuilder & registry) { + registry.add({ + "common.gather_add_rmsnorm_binary_f32", + GGML_OP_GET_ROWS, + DispatchMatchKind::Fused, + 1100, + DispatchSource::Common, + match_gather_add_rmsnorm_binary_f32_dispatch, + }); registry.add({ "common.gather_add_f32", GGML_OP_GET_ROWS, diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-get-rows.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-get-rows.cpp index e125427665fa..73b4ca960656 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-get-rows.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-get-rows.cpp @@ -6,6 +6,7 @@ #include "kernel-corpus/kernel-corpus-catalog-verify.h" #include +#include #include #include #include @@ -18,6 +19,8 @@ namespace { static constexpr KernelCatalogRef kGetRowsF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_get_rows_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"); static constexpr KernelCatalogRef kGetRowsRmsNormBinaryQ8_1X4F16Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_get_rows_rmsnorm_binary_q8_1_x4_f16"); static constexpr int64_t kMaximumHiddenElements = int64_t{ 1 } << 30; @@ -171,6 +174,19 @@ struct GetRowsRmsNormMatch { } }; +struct GetRowsScaleMatch { + GetRowsMatch rows; + const GraphNode * scale_node = nullptr; + const Value * output = nullptr; + size_t scale_index = 0; + float scale = 0.0f; + float bias = 0.0f; + + bool matched() const { + return rows.matched() && scale_node != nullptr && output != nullptr; + } +}; + static void add_common_compile_parameters(Dispatch & dispatch, const GetRowsMatch & match) { dispatch.kernel.compile_parameters.emplace("ggml.get_rows_f32.token_capacity", to_config_value(match.token_count)); dispatch.kernel.compile_parameters.emplace("ggml.get_rows_f32.hidden_capacity", to_config_value(match.hidden_size)); @@ -230,6 +246,53 @@ static GetRowsMatch match_get_rows_f32(const Graph & graph, const GraphNode * no return match; } +static GetRowsScaleMatch match_get_rows_scale_f32(const DispatchMatchContext & context) { + GetRowsScaleMatch match; + match.rows = match_get_rows_f32(context.graph, context.root_node); + if (!match.rows.matched() || !context.graph.has_index()) { + return {}; + } + + const bool supported_weight = match.rows.weight->type == GGML_TYPE_F32 || + (match.rows.weight->type == GGML_TYPE_Q5_1 && match.rows.token_count > 1); + if (!supported_weight) { + return {}; + } + + const std::vector & consumers = context.graph.index().consumers(match.rows.output->id); + if (consumers.size() != 1 || consumers.front() == nullptr || consumers.front()->op != GGML_OP_SCALE || + consumers.front()->inputs.size() != 1 || consumers.front()->inputs[0] != match.rows.output->id) { + return {}; + } + + const GraphNode * scale_node = consumers.front(); + const ScaleParams * params = op_params_as(scale_node->params); + const Value * output = graph_value(context.graph, scale_node->output); + if (params == nullptr || !std::isfinite(params->scale) || !std::isfinite(params->bias) || output == nullptr || + output->type != GGML_TYPE_F32 || !output->contiguous || output->alias_source.value >= 0 || + output->ne != match.rows.output->ne || output->byte_count != match.rows.output->byte_count) { + return {}; + } + + const std::array buffers = { match.rows.ids, match.rows.weight, match.rows.output, output }; + for (size_t lhs = 0; lhs < buffers.size(); ++lhs) { + for (size_t rhs = lhs + 1; rhs < buffers.size(); ++rhs) { + if (buffers[lhs]->storage_root == buffers[rhs]->storage_root) { + return {}; + } + } + } + if (!context.graph.index().node_index(scale_node, match.scale_index)) { + return {}; + } + + match.scale_node = scale_node; + match.output = output; + match.scale = params->scale; + match.bias = params->bias; + return match; +} + static GetRowsRmsNormMatch match_get_rows_rmsnorm_binary(const DispatchMatchContext & context) { GetRowsRmsNormMatch match; match.rows = match_get_rows_f32(context.graph, context.root_node); @@ -339,6 +402,31 @@ static bool match_get_rows_f32_dispatch(const DispatchMatchContext & context, Di return true; } +static bool match_get_rows_scale_f32_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + const GetRowsScaleMatch match = match_get_rows_scale_f32(context); + if (!match.matched() || match.scale_index >= context.covered_nodes.size() || + context.covered_nodes[context.root_index] || context.covered_nodes[match.scale_index]) { + return false; + } + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kGetRowsScaleF32Kernel); + add_common_integer_parameters(dispatch, match.rows); + add_common_compile_parameters(dispatch, match.rows); + dispatch.kernel.compile_parameters.emplace("ggml.get_rows_scale_f32.scale", to_config_value(match.scale)); + dispatch.kernel.compile_parameters.emplace("ggml.get_rows_scale_f32.bias", to_config_value(match.bias)); + dispatch.bindings.push_back({ match.rows.ids->id, 0, match.rows.ids->byte_count }); + dispatch.bindings.push_back({ match.rows.weight->id, 0, match.rows.weight->byte_count }); + dispatch.bindings.push_back({ match.rows.output->id, 0, match.rows.output->byte_count }); + dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + + dispatch_match.covered_nodes.push_back(context.root_index); + dispatch_match.covered_nodes.push_back(match.scale_index); + dispatch_match.dispatches.push_back(std::move(dispatch)); + return true; +} + static bool match_get_rows_f32_next_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { const GetRowsMatch match = match_get_rows_f32(context.graph, context.root_node); if (!match.matched()) { @@ -438,6 +526,14 @@ static bool match_get_rows_rmsnorm_binary_dispatch(const DispatchMatchContext & } // namespace void register_get_rows_dispatches(DispatchRegistryBuilder & registry) { + registry.add({ + "common.get_rows_scale.f32", + GGML_OP_GET_ROWS, + DispatchMatchKind::Fused, + 300, + DispatchSource::Common, + match_get_rows_scale_f32_dispatch, + }); registry.add({ "common.get_rows_rmsnorm_binary.q8_1_x4_f16", GGML_OP_GET_ROWS, diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-common.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-common.h index 3356ed5b79fd..25a024d7b01c 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-common.h +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-common.h @@ -1,6 +1,7 @@ #pragma once #include "../dispatch-registry.h" +#include "dispatch-binary-common.h" #include "dispatch-mul-mat-weight-format.h" #include "ggml.h" #include "graph/graph-matcher.h" @@ -146,8 +147,8 @@ inline bool common_mul_mat_uses_k16_major_f16(CommonMulMatWeightFormat format, int64_t output_size, int64_t token_count, bool packed_producer = false) { - if (token_count < 512 || token_count > 2048 || token_count % 512 != 0 || - input_size % 256 != 0 || output_size % 64 != 0) { + if (token_count < 512 || token_count > 2048 || token_count % 512 != 0 || input_size % 256 != 0 || + output_size % 64 != 0) { return false; } if (format == CommonMulMatWeightFormat::Q4K) { @@ -381,39 +382,6 @@ inline bool common_is_swiglu_params(const OpParams & params) { return glu_params != nullptr && glu_params->op == GGML_GLU_OP_SWIGLU; } -inline bool common_fused_binary_kind_from_params(const OpParams & params, BinaryKind & kind) { - const BinaryParams * binary_params = op_params_as(params); - if (binary_params != nullptr && binary_kind_supported(binary_params->op)) { - kind = binary_params->op; - return true; - } - - const GluParams * glu_params = op_params_as(params); - if (glu_params == nullptr) { - return false; - } - - switch (glu_params->op) { - case GGML_GLU_OP_REGLU: - kind = BinaryKind::RegLU; - return true; - case GGML_GLU_OP_SWIGLU: - kind = BinaryKind::SwiGLU; - return true; - case GGML_GLU_OP_GEGLU: - kind = BinaryKind::GeGLU; - return true; - case GGML_GLU_OP_GEGLU_ERF: - kind = BinaryKind::GeGLUErf; - return true; - case GGML_GLU_OP_GEGLU_QUICK: - kind = BinaryKind::GeGLUQuick; - return true; - default: - return false; - } -} - inline CommonMulMatMatch common_match_mul_mat_any_format(const Graph & graph, const GraphNode * node, KernelCatalogRef kernel, @@ -525,8 +493,8 @@ inline bool common_prepare_k16_major_f16_input(const DispatchMatchContext & cont copy.bindings.push_back({ packed, 0, bytes }); match.dispatches.push_back(std::move(copy)); Status status; - if (!match.metadata.append_generated_resource( - { input.id, GeneratedResourceRole::F16K16Major, packed, bytes, {} }, status)) { + if (!match.metadata.append_generated_resource({ input.id, GeneratedResourceRole::F16K16Major, packed, bytes, {} }, + status)) { match.status.append(status); return false; } 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 b7e16b6e2b90..5a35e811c092 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 @@ -5,8 +5,8 @@ #include "graph/graph-matcher.h" #include "kernel-corpus/kernel-corpus-catalog-verify.h" -#include #include +#include #include #include #include @@ -22,10 +22,6 @@ static constexpr KernelCatalogRef kMulMatTiledF32F16AlternateKernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_input_f32_publish_f32_f16_alternate"); static constexpr KernelCatalogRef kMulMatSkinnyF32F32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_skinny_input_f32_publish_f32"); -static constexpr KernelCatalogRef kMulMatSkinnyBiasF32F32Kernel = - GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_skinny_input_f32_bias_publish_f32"); -static constexpr KernelCatalogRef kMulMatSkinnyAddF32F32Kernel = - GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_skinny_input_f32_residual_publish_f32"); static constexpr KernelCatalogRef kMulMatSkinnyBiasAddF32F32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32"); static constexpr KernelCatalogRef kMulMatF32F32DecodeWave64Kernel = @@ -34,18 +30,10 @@ static constexpr KernelCatalogRef kMulMatAddF32F32DecodeWave64Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_add_f32_f32_decode_wave64"); static constexpr KernelCatalogRef kMulMatVectorF32F32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_vector_f32_f32"); -static constexpr KernelCatalogRef kMulMatVectorBiasF32F32Kernel = - GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_vector_bias_f32_f32"); -static constexpr KernelCatalogRef kMulMatVectorAddF32F32Kernel = - GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_vector_residual_f32_f32"); static constexpr KernelCatalogRef kMulMatVectorBiasAddF32F32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_vector_bias_residual_f32_f32"); static constexpr KernelCatalogRef kQuantizeQ8_1X4F32Kernel = GGML_HRX_KERNEL_REF("qwen3_moe", "ggml_quantize_q8_1_x4_f32"); -static constexpr KernelCatalogRef kMulMatTiledBiasF32F32Kernel = - GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_input_f32_bias_publish_f32"); -static constexpr KernelCatalogRef kMulMatTiledAddF32F32Kernel = - GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_input_f32_residual_publish_f32"); static constexpr KernelCatalogRef kMulMatTiledBiasAddF32F32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32"); static constexpr KernelCatalogRef kQuantizeF32SymmetricI4K64PlaneKernel = @@ -181,7 +169,6 @@ static bool common_mul_mat_postops_is_tiled_route(const MulMatPostOpsMatch & mat static bool common_mul_mat_uses_factored_tiled_kernel(KernelCatalogRef kernel) { return kernel.id == kMulMatTiledF32F32Kernel.id || kernel.id == kMulMatTiledF32F16AlternateKernel.id || - kernel.id == kMulMatTiledBiasF32F32Kernel.id || kernel.id == kMulMatTiledAddF32F32Kernel.id || kernel.id == kMulMatTiledBiasAddF32F32Kernel.id; } @@ -384,8 +371,9 @@ static bool has_embedded_external_concat_projection_ancestor(const Graph & graph } static bool is_qwen_endpoint_rmsnorm_projection(const Graph & graph, const CommonMulMatMatch & match) { - if (!graph.has_index() || match.weight == nullptr || match.input == nullptr || match.weight->type != GGML_TYPE_Q6_K || - match.token_count != 1 || match.output_size != kQwenVectorVocabularyCount) { + if (!graph.has_index() || match.weight == nullptr || match.input == nullptr || + match.weight->type != GGML_TYPE_Q6_K || match.token_count != 1 || + match.output_size != kQwenVectorVocabularyCount) { return false; } @@ -523,8 +511,8 @@ static bool match_q6_vector_selected_outputs_dispatch(const DispatchMatchContext if (!match.matched() || !context.graph.has_index() || match.weight->type != GGML_TYPE_Q6_K || match.token_count != 1 || match.weight->alias_source.value >= 0 || match.input_size % 256 != 0 || match.input_size > 8192 || match.output_size < 65536 || match.output_size % 64 != 0 || - match.output_size > kMulMatVectorSelectedOutputsMaxOutputSize || - match.output_size < 16 * match.input_size || !context.graph.index().consumers(match.output->id).empty() || + match.output_size > kMulMatVectorSelectedOutputsMaxOutputSize || match.output_size < 16 * match.input_size || + !context.graph.index().consumers(match.output->id).empty() || !has_embedded_external_concat_projection_ancestor(context.graph, match.input->id, match.input_size)) { return false; } @@ -573,17 +561,13 @@ static bool match_q6_vector_selected_outputs_dispatch(const DispatchMatchContext dispatch_match.transients.push_back({ selected_groups, "common.mul_mat.vector_q6_selected_outputs.groups", static_cast(selected_group_count) * sizeof(int32_t), 256 }); - dispatch_match.transients.push_back( - { partial_values, "common.mul_mat.vector_q6_selected_outputs.partial_values", + dispatch_match.transients.push_back({ partial_values, "common.mul_mat.vector_q6_selected_outputs.partial_values", partition_entry_count * sizeof(float), 256 }); - dispatch_match.transients.push_back( - { partial_ids, "common.mul_mat.vector_q6_selected_outputs.partial_ids", + dispatch_match.transients.push_back({ partial_ids, "common.mul_mat.vector_q6_selected_outputs.partial_ids", partition_entry_count * sizeof(int32_t), 256 }); - dispatch_match.transients.push_back( - { selected_outputs, "common.mul_mat.vector_q6_selected_outputs.indices", + dispatch_match.transients.push_back({ selected_outputs, "common.mul_mat.vector_q6_selected_outputs.indices", selected_output_count * sizeof(int32_t), 256 }); - dispatch_match.transients.push_back( - { selected_values, "common.mul_mat.vector_q6_selected_outputs.values", + dispatch_match.transients.push_back({ selected_values, "common.mul_mat.vector_q6_selected_outputs.values", selected_output_count * sizeof(float), 256 }); Dispatch select_groups; @@ -654,18 +638,17 @@ static bool match_q6_vector_selected_outputs_dispatch(const DispatchMatchContext selected_output_pass.kernel.compile_parameters.emplace("ggml.matmul.vector.output_size", common_to_config_value(match.output_size)); selected_output_pass.kernel.compile_parameters.emplace("ggml.matmul.vector.output_accumulation", "0"); - selected_output_pass.kernel.compile_parameters.emplace("ggml.matmul.vector.weight_offset", - common_to_config_value( - static_cast(symmetric_i2_bytes))); + selected_output_pass.kernel.compile_parameters.emplace( + "ggml.matmul.vector.weight_offset", common_to_config_value(static_cast(symmetric_i2_bytes))); selected_output_pass.bindings.push_back({ match.input->id, 0, match.input->byte_count }); selected_output_pass.bindings.push_back({ match.weight->id, 0, materialized_weight_bytes, kQ6KSymmetricI2PackedK256Row64ScaleRowLayout, match.weight->type, match.input_size, match.output_size, match.weight->byte_count }); - selected_output_pass.bindings.push_back( - { selected_outputs, static_cast(selected_output_offset) * sizeof(int32_t), + selected_output_pass.bindings.push_back({ selected_outputs, + static_cast(selected_output_offset) * sizeof(int32_t), static_cast(selected_output_chunk_size) * sizeof(int32_t) }); - selected_output_pass.bindings.push_back( - { selected_values, static_cast(selected_output_offset) * sizeof(float), + selected_output_pass.bindings.push_back({ selected_values, + static_cast(selected_output_offset) * sizeof(float), static_cast(selected_output_chunk_size) * sizeof(float) }); selected_output_pass.bindings.push_back({ match.output->id, 0, match.output->byte_count }); dispatch_match.dispatches.push_back(std::move(selected_output_pass)); @@ -697,9 +680,8 @@ static bool build_q6_vector_dispatch(const CommonMulMatMatch & match, common_to_config_value(static_cast(storage))); dispatch.bindings.push_back({ match.input->id, 0, match.input->byte_count }); if (storage == Q6VectorStorage::I8) { - dispatch.bindings.push_back({ match.weight->id, 0, weight_bytes, kQ6KI8K32Row64Layout, - match.weight->type, match.input_size, match.output_size, - match.weight->byte_count }); + dispatch.bindings.push_back({ match.weight->id, 0, weight_bytes, kQ6KI8K32Row64Layout, match.weight->type, + match.input_size, match.output_size, match.weight->byte_count }); } else { dispatch.bindings.push_back({ match.weight->id, 0, weight_bytes, kQ6KPackedK256Row64ScaleRowLayout, match.weight->type, match.input_size, match.output_size, @@ -731,9 +713,9 @@ static bool match_q6_vector_dispatch(const DispatchMatchContext & context, Dispa if (is_qwen_endpoint_rmsnorm_projection(context.graph, match)) { return false; } - if (find_alternate_value(context.graph, context.plan, match.input->id, GGML_TYPE_Q8_1, - static_cast(match.token_count) * ggml_row_size(GGML_TYPE_Q8_1, - match.input_size)) != nullptr) { + if (find_alternate_value( + context.graph, context.plan, match.input->id, GGML_TYPE_Q8_1, + static_cast(match.token_count) * ggml_row_size(GGML_TYPE_Q8_1, match.input_size)) != nullptr) { return false; } @@ -765,30 +747,28 @@ static bool build_q6_k_prefill_wave32_dispatch(const CommonMulMatMatch & match, return true; } -static bool match_q6_k_prefill_wave32_dispatch(const DispatchMatchContext & context, - DispatchMatch & dispatch_match) { - const CommonMulMatMatch match = common_match_mul_mat_any_format( - context.graph, context.root_node, kMulMatQ6KF32WmmaPrefillWave32Kernel, false); +static bool match_q6_k_prefill_wave32_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + const CommonMulMatMatch match = + common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatQ6KF32WmmaPrefillWave32Kernel, false); if (!match.matched() || match.weight->type != GGML_TYPE_Q6_K || match.weight->alias_source.value >= 0 || match.token_count < 128 || match.token_count > 2048 || match.token_count % 128 != 0 || match.input_size % 256 != 0 || match.output_size % 64 != 0) { return false; } - const bool packed_input = common_mul_mat_uses_k16_major_f16( - match.weight_format, match.input_size, match.output_size, match.token_count); + const bool packed_input = + common_mul_mat_uses_k16_major_f16(match.weight_format, match.input_size, match.output_size, match.token_count); DispatchBinding activation; const bool prepared = packed_input ? - common_prepare_k16_major_f16_input(context, *match.input, match.input_size, match.token_count, - dispatch_match, activation) : + common_prepare_k16_major_f16_input(context, *match.input, match.input_size, + match.token_count, dispatch_match, activation) : common_prepare_f16_input(context, *match.input, match.input_size, match.token_count, dispatch_match, activation); if (!prepared) { return false; } return build_q6_k_prefill_wave32_dispatch(match, activation.value, activation.length, - kMulMatQ6KF16WmmaPrefillWave32Kernel, context.root_index, - dispatch_match); + kMulMatQ6KF16WmmaPrefillWave32Kernel, context.root_index, dispatch_match); } static bool match_symmetric_i4_prefill_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { @@ -919,7 +899,6 @@ static bool match_q5_k_symmetric_i8_prefill_dispatch(const DispatchMatchContext return true; } - static bool match_packed_q8_1_x4_prefill_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { const CommonMulMatMatch match = common_match_mul_mat_any_format(context.graph, context.root_node, @@ -941,22 +920,23 @@ static bool match_packed_q8_1_x4_prefill_dispatch(const DispatchMatchContext & c const bool use_f16 = match.weight->type == GGML_TYPE_Q4_K && match.weight->alias_source.value < 0 && match.output_size >= match.input_size / 4; DispatchBinding activation; - const bool packed_input = use_f16 && common_mul_mat_uses_k16_major_f16( - match.weight_format, match.input_size, match.output_size, match.token_count, - context.plan.metadata.find_generated_resource(match.input->id, GeneratedResourceRole::F16K16Major) != nullptr); + const bool packed_input = + use_f16 && + common_mul_mat_uses_k16_major_f16(match.weight_format, match.input_size, match.output_size, match.token_count, + context.plan.metadata.find_generated_resource( + match.input->id, GeneratedResourceRole::F16K16Major) != nullptr); if (use_f16) { const bool prepared = packed_input ? - common_prepare_k16_major_f16_input(context, *match.input, match.input_size, match.token_count, - dispatch_match, activation) : + common_prepare_k16_major_f16_input(context, *match.input, match.input_size, + match.token_count, dispatch_match, activation) : common_prepare_f16_input(context, *match.input, match.input_size, match.token_count, dispatch_match, activation); if (!prepared) { return false; } - } else if (!common_prepare_q8_1_x4_input(context, *match.input, match.input_size, match.token_count, - dispatch_match, activation, - match.weight->type == GGML_TYPE_Q4_K ? - CommonQ8ActivationPolicy::AllowStandaloneQuantize : + } else if (!common_prepare_q8_1_x4_input( + context, *match.input, match.input_size, match.token_count, dispatch_match, activation, + match.weight->type == GGML_TYPE_Q4_K ? CommonQ8ActivationPolicy::AllowStandaloneQuantize : CommonQ8ActivationPolicy::ExistingAlternateOnly)) { return false; } @@ -965,17 +945,15 @@ static bool match_packed_q8_1_x4_prefill_dispatch(const DispatchMatchContext & c dispatch.kernel = make_kernel_specialization(use_f16 ? kMulMatQ4KF16WmmaPrefillWave32Kernel : kMulMatQ5KIQ4XSQ8_1X4WmmaToken256Kernel); dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); - dispatch.kernel.compile_parameters.emplace(use_f16 ? "ggml.mul_mat.input_size" : - "ggml.mul_mat_q8_1_x4.input_size", + dispatch.kernel.compile_parameters.emplace(use_f16 ? "ggml.mul_mat.input_size" : "ggml.mul_mat_q8_1_x4.input_size", common_to_config_value(match.input_size)); - dispatch.kernel.compile_parameters.emplace(use_f16 ? "ggml.mul_mat.output_size" : - "ggml.mul_mat_q8_1_x4.output_size", + dispatch.kernel.compile_parameters.emplace( + use_f16 ? "ggml.mul_mat.output_size" : "ggml.mul_mat_q8_1_x4.output_size", common_to_config_value(match.output_size)); - dispatch.kernel.compile_parameters.emplace(use_f16 ? "ggml.workload.token_capacity" : - "ggml.mul_mat_q8_1_x4.token_capacity", + dispatch.kernel.compile_parameters.emplace( + use_f16 ? "ggml.workload.token_capacity" : "ggml.mul_mat_q8_1_x4.token_capacity", common_to_config_value(match.token_count)); - const bool pack_q4 = match.weight_format == CommonMulMatWeightFormat::Q4K && - match.weight->alias_source.value < 0; + const bool pack_q4 = match.weight_format == CommonMulMatWeightFormat::Q4K && match.weight->alias_source.value < 0; if (use_f16) { dispatch.kernel.compile_parameters.emplace("ggml.mul_mat.f16_input_layout", packed_input ? "1" : "0"); } else { @@ -999,8 +977,7 @@ static bool match_packed_q8_1_x4_prefill_dispatch(const DispatchMatchContext & c return true; } -static bool match_q4_k_q8_1_x4_prefill_dispatch(const DispatchMatchContext & context, - DispatchMatch & dispatch_match) { +static bool match_q4_k_q8_1_x4_prefill_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { if (context.root_node == nullptr || context.root_node->inputs.empty()) { return false; } @@ -1024,11 +1001,11 @@ static MulMatPostOpsMatch match_mul_mat_postops(const DispatchMatchContext & con return match; } - const bool packed_row_contraction = root.token_count <= 5 && root.weight->alias_source.value < 0 && - (root.weight_format == CommonMulMatWeightFormat::Q4K || - root.weight_format == CommonMulMatWeightFormat::Q6K) && - root.input_size >= 4096 && root.output_size >= 4096 && - root.output_size <= root.input_size && root.output_size % 64 == 0; + const bool packed_row_contraction = + root.token_count <= 5 && root.weight->alias_source.value < 0 && + (root.weight_format == CommonMulMatWeightFormat::Q4K || root.weight_format == CommonMulMatWeightFormat::Q6K) && + root.input_size >= 4096 && root.output_size >= 4096 && root.output_size <= root.input_size && + root.output_size % 64 == 0; const bool batch1_requires_q8_activation = root.token_count == 1 && !packed_row_contraction; const Value * current = root.output; @@ -1041,7 +1018,8 @@ static MulMatPostOpsMatch match_mul_mat_postops(const DispatchMatchContext & con current = reshaped; } } - const int add_limit = match.layout_node != nullptr || (root.token_count == 1 && !allow_batch1_bias_residual) ? 1 : 2; + const int add_limit = + match.layout_node != nullptr || (root.token_count == 1 && !allow_batch1_bias_residual) ? 1 : 2; for (int add_index = 0; add_index < add_limit; ++add_index) { const GraphNode * add_node = common_find_only_consumer_with_op(context.graph, current->id, GGML_OP_ADD); if (add_node == nullptr || !common_binary_node_is_add(*add_node)) { @@ -1112,12 +1090,9 @@ static MulMatPostOpsMatch match_mul_mat_postops(const DispatchMatchContext & con return match; } -static MulMatPostOpsMatch match_mul_mat_postops(const DispatchMatchContext & context, - bool allow_batch1_bias_residual) { - return match_mul_mat_postops(context, kMulMatF32F32WmmaKernel, - kMulMatVectorBiasF32F32Kernel, - kMulMatVectorAddF32F32Kernel, - kMulMatVectorBiasAddF32F32Kernel, +static MulMatPostOpsMatch match_mul_mat_postops(const DispatchMatchContext & context, bool allow_batch1_bias_residual) { + return match_mul_mat_postops(context, kMulMatF32F32WmmaKernel, kMulMatVectorBiasAddF32F32Kernel, + kMulMatVectorBiasAddF32F32Kernel, kMulMatVectorBiasAddF32F32Kernel, allow_batch1_bias_residual); } @@ -1160,7 +1135,8 @@ static bool try_match_fused_unary(const DispatchMatchContext & context, CommonMu } // namespace -static bool build_mul_mat_dispatch(const DispatchMatchContext & context, +static bool build_mul_mat_dispatch( + const DispatchMatchContext & context, const CommonMulMatMatch & match, DispatchMatch & dispatch_match, CommonQ8ActivationPolicy q8_policy = CommonQ8ActivationPolicy::ExistingAlternateOnly) { @@ -1169,8 +1145,8 @@ static bool build_mul_mat_dispatch(const DispatchMatchContext & context, const bool publish_q8_alternate = match.kernel.id == kMulMatTiledF32F32Kernel.id && common_mul_mat_has_q8_alternate_consumer(context.graph, match); Dispatch dispatch; - dispatch.kernel = make_kernel_specialization(publish_f16_alternate ? kMulMatTiledF32F16AlternateKernel : - match.kernel); + dispatch.kernel = + make_kernel_specialization(publish_f16_alternate ? kMulMatTiledF32F16AlternateKernel : match.kernel); dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); dispatch.kernel.compile_parameters.emplace("ggml.workload.token_capacity", common_to_config_value(match.token_count)); @@ -1179,8 +1155,7 @@ static bool build_mul_mat_dispatch(const DispatchMatchContext & context, dispatch.kernel.compile_parameters.emplace("ggml.mul_mat.output_accumulation", "0"); dispatch.kernel.compile_parameters.emplace("ggml.mul_mat.output_unary_op", std::to_string(unary_kind_config_value(match.output_unary_op))); - const bool pack_rows = match.token_count <= 5 && - common_mul_mat_has_skinny_reduction_shape(match.input_size) && + const bool pack_rows = match.token_count <= 5 && common_mul_mat_has_skinny_reduction_shape(match.input_size) && match.output_size % 64 == 0 && match.weight->alias_source.value < 0; const bool can_pack_q4 = pack_rows && match.weight_format == CommonMulMatWeightFormat::Q4K; const bool can_pack_q6 = pack_rows && match.weight_format == CommonMulMatWeightFormat::Q6K; @@ -1196,8 +1171,8 @@ static bool build_mul_mat_dispatch(const DispatchMatchContext & context, activation = { f16_activation->alternate_value, 0, f16_activation_bytes }; dispatch.kernel.compile_parameters.emplace("ggml.mul_mat.activation_format", std::to_string(GGML_TYPE_F16)); } else if (can_pack_q4 || can_pack_q6) { - if (common_prepare_q8_1_x4_input(context, *match.input, match.input_size, match.token_count, - dispatch_match, activation, q8_policy)) { + if (common_prepare_q8_1_x4_input(context, *match.input, match.input_size, match.token_count, dispatch_match, + activation, q8_policy)) { pack_q4 = can_pack_q4; pack_q6 = can_pack_q6; dispatch.kernel.compile_parameters.emplace("ggml.mul_mat.activation_format", @@ -1207,23 +1182,23 @@ static bool build_mul_mat_dispatch(const DispatchMatchContext & context, } } const CommonMulMatWeightFormat format = pack_q4 ? CommonMulMatWeightFormat::Q4KRow64 : - pack_q6 ? CommonMulMatWeightFormat::Q6KRow64 : match.weight_format; - dispatch.kernel.compile_parameters.emplace( - "ggml.mul_mat.weight_format", + pack_q6 ? CommonMulMatWeightFormat::Q6KRow64 : + match.weight_format; + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat.weight_format", common_to_config_value(common_mul_mat_format_config_value(format))); dispatch.bindings.push_back(activation); if (pack_q4 || pack_q6) { const char * layout = pack_q4 ? kQ4KPackedK256Row64Layout : kQ6KPackedK256Row64ScaleRowLayout; - dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count, layout, - match.weight->type, match.input_size, match.output_size, - match.weight->byte_count }); + dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count, layout, 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 }); if (publish_f16_alternate) { const size_t bytes = match.output->byte_count / 2; - const ValueId f16_output(context.next_plan_value.value + static_cast(dispatch_match.transients.size())); + const ValueId f16_output(context.next_plan_value.value + + static_cast(dispatch_match.transients.size())); constexpr const char * name = "common.mul_mat.tiled.f16"; dispatch_match.transients.push_back({ f16_output, name, bytes, 256 }); Status status; @@ -1323,15 +1298,13 @@ static int64_t vector_max_output_size(CommonMulMatWeightFormat format) { } static bool vector_mul_mat_route_supports(const CommonMulMatMatch & match) { - return match.matched() && match.token_count == 1 && match.input_size >= 256 && - match.input_size <= 32768 && match.input_size % 32 == 0 && - match.output_size <= vector_max_output_size(match.weight_format); + return match.matched() && match.token_count == 1 && match.input_size >= 256 && match.input_size <= 32768 && + match.input_size % 32 == 0 && match.output_size <= vector_max_output_size(match.weight_format); } static bool vector_postops_route_supports(const MulMatPostOpsMatch & match) { - return match.matched() && match.token_count == 1 && match.input_size >= 256 && - match.input_size <= 32768 && match.input_size % 32 == 0 && - match.output_size <= vector_max_output_size(match.weight_format); + return match.matched() && match.token_count == 1 && match.input_size >= 256 && match.input_size <= 32768 && + match.input_size % 32 == 0 && match.output_size <= vector_max_output_size(match.weight_format); } static size_t vector_publish_byte_count(ggml_type type, int64_t token_count, int64_t output_size) { @@ -1355,11 +1328,10 @@ static bool is_vector_packed_q8_consumer(const Graph & graph, const GraphNode * } const bool supported_weight = weight->type == GGML_TYPE_Q4_K || weight->type == GGML_TYPE_Q6_K; - return supported_weight && input.ne[1] >= 1 && input.ne[1] <= 5 && input.ne[0] >= 256 && - input.ne[0] <= 32768 && input.ne[0] % 256 == 0 && input.ne[2] == 1 && input.ne[3] == 1 && - weight->ne[0] == input.ne[0] && weight->ne[1] >= 64 && weight->ne[1] <= 262144 && - weight->ne[1] % 64 == 0 && weight->ne[2] == 1 && weight->ne[3] == 1 && - output->ne[0] == weight->ne[1] && output->ne[1] == input.ne[1] && output->ne[2] == 1 && + return supported_weight && input.ne[1] >= 1 && input.ne[1] <= 5 && input.ne[0] >= 256 && input.ne[0] <= 32768 && + input.ne[0] % 256 == 0 && input.ne[2] == 1 && input.ne[3] == 1 && weight->ne[0] == input.ne[0] && + weight->ne[1] >= 64 && weight->ne[1] <= 262144 && weight->ne[1] % 64 == 0 && weight->ne[2] == 1 && + weight->ne[3] == 1 && output->ne[0] == weight->ne[1] && output->ne[1] == input.ne[1] && output->ne[2] == 1 && output->ne[3] == 1; } @@ -1393,9 +1365,8 @@ static bool is_vector_f16_consumer(const Graph & graph, const GraphNode * consum return false; } return input.ne[0] == weight->ne[0] && output->ne[0] == weight->ne[1] && output->ne[1] == input.ne[1] && - input.ne[2] == 1 && input.ne[3] == 1 && weight->ne[2] == 1 && weight->ne[3] == 1 && - output->ne[2] == 1 && output->ne[3] == 1 && - common_mul_mat_uses_k16_major_f16(format, input.ne[0], output->ne[0], input.ne[1]); + input.ne[2] == 1 && input.ne[3] == 1 && weight->ne[2] == 1 && weight->ne[3] == 1 && output->ne[2] == 1 && + output->ne[3] == 1 && common_mul_mat_uses_k16_major_f16(format, input.ne[0], output->ne[0], input.ne[1]); } static bool has_vector_f16_consumer(const Graph & graph, const Value & value) { @@ -1418,11 +1389,13 @@ static VectorPublishPlan make_vector_publish_plan(const DispatchMatchContext & c VectorPublishPlan plan; if (has_vector_packed_q8_consumer(context.graph, output)) { plan.publish_q8 = true; - plan.q8_output = ValueId(context.next_plan_value.value + static_cast(dispatch_match.transients.size())); + plan.q8_output = + ValueId(context.next_plan_value.value + static_cast(dispatch_match.transients.size())); plan.q8_byte_count = vector_publish_byte_count(GGML_TYPE_Q8_1, token_count, output_size); } else if (has_vector_f16_consumer(context.graph, output)) { plan.fused_format = VectorPublishFormat::F16; - plan.fused_output = ValueId(context.next_plan_value.value + static_cast(dispatch_match.transients.size())); + plan.fused_output = + ValueId(context.next_plan_value.value + static_cast(dispatch_match.transients.size())); plan.fused_byte_count = vector_publish_byte_count(GGML_TYPE_F16, token_count, output_size); } return plan; @@ -1543,8 +1516,7 @@ static bool build_vector_mul_mat_dispatch(const DispatchMatchContext & context, } dispatch.kernel.compile_parameters.emplace( - "ggml.matmul.vector.weight_format", - common_to_config_value(common_mul_mat_format_config_value(weight_format))); + "ggml.matmul.vector.weight_format", common_to_config_value(common_mul_mat_format_config_value(weight_format))); if (uses_q8_activation) { dispatch.kernel.compile_parameters.emplace("ggml.matmul.vector.activation_format", std::to_string(GGML_TYPE_Q8_1)); @@ -1552,9 +1524,8 @@ static bool build_vector_mul_mat_dispatch(const DispatchMatchContext & context, dispatch.bindings.push_back(activation); if (weight_layout != nullptr) { - dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count, weight_layout, - match.weight->type, match.input_size, match.output_size, - match.weight->byte_count }); + dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count, weight_layout, 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 }); } @@ -1582,19 +1553,12 @@ static bool build_vector_mul_mat_postops_dispatch(const DispatchMatchContext & c return false; } - KernelCatalogRef kernel = {}; - if (match.has_bias && match.has_residual) { - kernel = kMulMatVectorBiasAddF32F32Kernel; - } else if (match.has_bias) { - kernel = kMulMatVectorBiasF32F32Kernel; - } else if (match.has_residual) { - kernel = kMulMatVectorAddF32F32Kernel; - } else { + if (!match.has_bias && !match.has_residual) { return false; } Dispatch dispatch; - dispatch.kernel = make_kernel_specialization(kernel); + dispatch.kernel = make_kernel_specialization(kMulMatVectorBiasAddF32F32Kernel); dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); dispatch.kernel.compile_parameters.emplace("ggml.workload.token_capacity", common_to_config_value(match.token_count)); @@ -1604,7 +1568,8 @@ static bool build_vector_mul_mat_postops_dispatch(const DispatchMatchContext & c common_to_config_value(match.output_size)); dispatch.kernel.compile_parameters.emplace("ggml.matmul.vector_postops.apply_bias", common_to_config_value(match.has_bias ? int64_t{ 1 } : int64_t{ 0 })); - dispatch.kernel.compile_parameters.emplace("ggml.matmul.vector_postops.apply_residual", + dispatch.kernel.compile_parameters.emplace( + "ggml.matmul.vector_postops.apply_residual", common_to_config_value(match.has_residual ? int64_t{ 1 } : int64_t{ 0 })); VectorPublishPlan publish = make_vector_publish_plan(context, *match.residual_output, match.token_count, match.output_size, dispatch_match); @@ -1635,18 +1600,17 @@ static bool build_vector_mul_mat_postops_dispatch(const DispatchMatchContext & c dispatch.bindings.push_back(activation); if (weight_layout != nullptr) { - dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count, weight_layout, - match.weight->type, match.input_size, match.output_size, - match.weight->byte_count }); + dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count, weight_layout, 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 }); } - if (match.has_bias) { - dispatch.bindings.push_back({ match.bias->id, 0, match.bias->byte_count }); - } - if (match.has_residual) { - dispatch.bindings.push_back({ match.residual_input->id, 0, match.residual_input->byte_count }); - } + const Value * placeholder = match.has_bias ? match.bias : match.residual_input; + dispatch.bindings.push_back(match.has_bias ? DispatchBinding{ match.bias->id, 0, match.bias->byte_count } : + DispatchBinding{ placeholder->id, 0, placeholder->byte_count }); + dispatch.bindings.push_back(match.has_residual ? + DispatchBinding{ match.residual_input->id, 0, match.residual_input->byte_count } : + DispatchBinding{ placeholder->id, 0, placeholder->byte_count }); dispatch.bindings.push_back({ match.residual_output->id, 0, match.residual_output->byte_count }); dispatch.bindings.push_back(vector_next_output_binding(*match.residual_output, publish)); @@ -1706,9 +1670,8 @@ static bool match_q6_vector_final_projection_q8_dispatch(const DispatchMatchCont match = common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatF32F32WmmaKernel, false); } if (!match.matched() || !context.graph.has_index() || match.weight->type != GGML_TYPE_Q6_K || - match.token_count != 1 || match.weight->alias_source.value >= 0 || - match.input_size % 256 != 0 || match.output_size % 64 != 0 || - match.output_size > kMulMatVectorPackedMaxOutputSize || + match.token_count != 1 || match.weight->alias_source.value >= 0 || match.input_size % 256 != 0 || + match.output_size % 64 != 0 || match.output_size > kMulMatVectorPackedMaxOutputSize || !context.graph.index().consumers(match.output->id).empty()) { return false; } @@ -1743,8 +1706,7 @@ static bool build_mul_mat_postops_dispatch(const DispatchMatchContext & context, common_to_config_value(match.input_size)); dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_postops.output_size", common_to_config_value(match.output_size)); - const bool pack_rows = match.token_count <= 5 && - common_mul_mat_has_skinny_reduction_shape(match.input_size) && + const bool pack_rows = match.token_count <= 5 && common_mul_mat_has_skinny_reduction_shape(match.input_size) && match.output_size % 64 == 0 && match.weight->alias_source.value < 0; const bool can_pack_q4 = pack_rows && match.weight_format == CommonMulMatWeightFormat::Q4K; const bool can_pack_q6 = pack_rows && match.weight_format == CommonMulMatWeightFormat::Q6K; @@ -1755,9 +1717,8 @@ static bool build_mul_mat_postops_dispatch(const DispatchMatchContext & context, return false; } if (can_pack_q4 || can_pack_q6) { - if (common_prepare_q8_1_x4_input(context, *match.input, match.input_size, match.token_count, - dispatch_match, activation, - CommonQ8ActivationPolicy::ExistingAlternateOnly)) { + if (common_prepare_q8_1_x4_input(context, *match.input, match.input_size, match.token_count, dispatch_match, + activation, CommonQ8ActivationPolicy::ExistingAlternateOnly)) { pack_q4 = can_pack_q4; pack_q6 = can_pack_q6; dispatch.kernel.compile_parameters.emplace("ggml.mul_mat.activation_format", @@ -1768,25 +1729,27 @@ static bool build_mul_mat_postops_dispatch(const DispatchMatchContext & context, return false; } const CommonMulMatWeightFormat format = pack_q4 ? CommonMulMatWeightFormat::Q4KRow64 : - pack_q6 ? CommonMulMatWeightFormat::Q6KRow64 : match.weight_format; - dispatch.kernel.compile_parameters.emplace( - "ggml.mul_mat_postops.weight_format", + pack_q6 ? CommonMulMatWeightFormat::Q6KRow64 : + match.weight_format; + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_postops.weight_format", common_to_config_value(common_mul_mat_format_config_value(format))); + dispatch.kernel.compile_parameters.emplace( + "ggml.mul_mat_postops.epilogue", common_to_config_value((match.has_bias ? int64_t{ 1 } : int64_t{ 0 }) | + (match.has_residual ? int64_t{ 2 } : int64_t{ 0 }))); dispatch.bindings.push_back(activation); if (pack_q4 || pack_q6) { const char * layout = pack_q4 ? kQ4KPackedK256Row64Layout : kQ6KPackedK256Row64ScaleRowLayout; - dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count, layout, - match.weight->type, match.input_size, match.output_size, - match.weight->byte_count }); + dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count, layout, 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 }); } - if (match.has_bias) { - dispatch.bindings.push_back({ match.bias->id, 0, match.bias->byte_count }); - } - if (match.has_residual) { - dispatch.bindings.push_back({ match.residual_input->id, 0, match.residual_input->byte_count }); - } + const Value * placeholder = match.has_bias ? match.bias : match.residual_input; + dispatch.bindings.push_back(match.has_bias ? DispatchBinding{ match.bias->id, 0, match.bias->byte_count } : + DispatchBinding{ placeholder->id, 0, placeholder->byte_count }); + dispatch.bindings.push_back(match.has_residual ? + DispatchBinding{ match.residual_input->id, 0, match.residual_input->byte_count } : + DispatchBinding{ placeholder->id, 0, placeholder->byte_count }); dispatch.bindings.push_back({ match.residual_output->id, 0, match.residual_output->byte_count }); if (!append_covered_node_index_once(context.graph, context.covered_nodes, context.root_node, @@ -1811,22 +1774,19 @@ static bool build_mul_mat_postops_dispatch(const DispatchMatchContext & context, static bool match_skinny_mul_mat_postops_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { - const MulMatPostOpsMatch match = match_mul_mat_postops(context, kMulMatSkinnyF32F32Kernel, - kMulMatSkinnyBiasF32F32Kernel, - kMulMatSkinnyAddF32F32Kernel, - kMulMatSkinnyBiasAddF32F32Kernel); + const MulMatPostOpsMatch match = + match_mul_mat_postops(context, kMulMatSkinnyF32F32Kernel, kMulMatSkinnyBiasAddF32F32Kernel, + kMulMatSkinnyBiasAddF32F32Kernel, kMulMatSkinnyBiasAddF32F32Kernel); if (!common_mul_mat_postops_is_skinny_route(match)) { return false; } return build_mul_mat_postops_dispatch(context, match, dispatch_match); } -static bool match_tiled_mul_mat_postops_dispatch(const DispatchMatchContext & context, - DispatchMatch & dispatch_match) { - const MulMatPostOpsMatch match = match_mul_mat_postops(context, kMulMatTiledF32F32Kernel, - kMulMatTiledBiasF32F32Kernel, - kMulMatTiledAddF32F32Kernel, - kMulMatTiledBiasAddF32F32Kernel); +static bool match_tiled_mul_mat_postops_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + const MulMatPostOpsMatch match = + match_mul_mat_postops(context, kMulMatTiledF32F32Kernel, kMulMatTiledBiasAddF32F32Kernel, + kMulMatTiledBiasAddF32F32Kernel, kMulMatTiledBiasAddF32F32Kernel); if (!common_mul_mat_postops_is_tiled_route(match)) { return false; } @@ -1954,8 +1914,7 @@ static bool match_tiled_mul_mat_dispatch(const DispatchMatchContext & context, D return build_mul_mat_dispatch(context, match, dispatch_match); } -static bool match_skinny_mul_mat_unary_dispatch(const DispatchMatchContext & context, - DispatchMatch & dispatch_match) { +static bool match_skinny_mul_mat_unary_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { CommonMulMatMatch match = common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatSkinnyF32F32Kernel, false); if (!match.matched() || !try_match_fused_unary(context, match) || !common_mul_mat_is_skinny_route(match)) { @@ -1964,12 +1923,10 @@ static bool match_skinny_mul_mat_unary_dispatch(const DispatchMatchContext & con return build_mul_mat_dispatch(context, match, dispatch_match, CommonQ8ActivationPolicy::AllowStandaloneQuantize); } -static bool match_tiled_mul_mat_unary_dispatch(const DispatchMatchContext & context, - DispatchMatch & dispatch_match) { +static bool match_tiled_mul_mat_unary_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { CommonMulMatMatch match = common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatTiledF32F32Kernel, false); - if (!match.matched() || !try_match_fused_unary(context, match) || - !common_mul_mat_is_tiled_route(match)) { + if (!match.matched() || !try_match_fused_unary(context, match) || !common_mul_mat_is_tiled_route(match)) { return false; } return build_mul_mat_dispatch(context, match, dispatch_match); @@ -1988,10 +1945,12 @@ static bool match_decode_mul_mat_dispatch(const DispatchMatchContext & context, return false; } if ((match.weight_format == CommonMulMatWeightFormat::Q4K || - match.weight_format == CommonMulMatWeightFormat::Q6K) && match.output_size % 64 == 0 && - common_is_supported_dense_output_size(match.output_size) && match.weight->alias_source.value < 0) { + match.weight_format == CommonMulMatWeightFormat::Q6K) && + match.output_size % 64 == 0 && common_is_supported_dense_output_size(match.output_size) && + match.weight->alias_source.value < 0) { match.kernel = kMulMatF32F32WmmaKernel; - return build_mul_mat_dispatch(context, match, dispatch_match, CommonQ8ActivationPolicy::AllowStandaloneQuantize); + return build_mul_mat_dispatch(context, match, dispatch_match, + CommonQ8ActivationPolicy::AllowStandaloneQuantize); } else { build_decode_mul_mat_dispatch(match, dispatch_match, context.root_index); } @@ -2012,7 +1971,7 @@ static bool match_decode_mul_mat_add_dispatch(const DispatchMatchContext & conte postops.projection_output = match.root.output; postops.residual_input = match.residual_input; postops.residual_output = match.residual_output; - postops.kernel = kMulMatVectorAddF32F32Kernel; + postops.kernel = kMulMatVectorBiasAddF32F32Kernel; postops.weight_format = match.root.weight_format; postops.input_size = match.root.input_size; postops.output_size = match.root.output_size; diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rmsnorm.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rmsnorm.cpp index 2c805e3bf33c..2a078613af33 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rmsnorm.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rmsnorm.cpp @@ -1,5 +1,6 @@ #include "dispatch-rmsnorm.h" +#include "dispatch-binary-common.h" #include "dispatch-layout-utils.h" #include "dispatch-mul-mat-common.h" #include "ggml.h" @@ -18,6 +19,10 @@ namespace ggml::hrx { namespace { static constexpr KernelCatalogRef kRmsNormBinaryF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_rmsnorm_binary_f32"); +static constexpr KernelCatalogRef kRmsNormBinaryStridedF32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_rmsnorm_binary_strided_f32"); +static constexpr KernelCatalogRef kRmsNormMulAddF32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_rmsnorm_mul_add_f32"); static constexpr KernelCatalogRef kRmsNormBinaryF32F16Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_rmsnorm_binary_f32_f16"); static constexpr KernelCatalogRef kRmsNormBinaryF32K16Kernel = @@ -82,8 +87,7 @@ static bool supported_rmsnorm_input_layout(const Value & input, return true; } - if (input.ne[1] != token_count || input.ne[2] != 1 || input.ne[3] != 1 || - input.nb[1] % sizeof(float) != 0) { + if (input.ne[1] != token_count || input.ne[2] != 1 || input.ne[3] != 1 || input.nb[1] % sizeof(float) != 0) { return false; } input_stride = static_cast(input.nb[1] / sizeof(float)); @@ -94,14 +98,6 @@ static bool supported_rmsnorm_input_layout(const Value & input, return strided_f32_storage_span_bytes(input, input_span_bytes); } -static bool is_binary_op(ggml_op op) { - return op == GGML_OP_ADD || op == GGML_OP_SUB || op == GGML_OP_MUL || op == GGML_OP_DIV; -} - -static bool binary_kind_requires_order(BinaryKind kind) { - return kind == BinaryKind::Sub || kind == BinaryKind::Div; -} - static bool is_packed_q8_consumer(const Graph & graph, const GraphNode * consumer, const Value & input) { if (consumer == nullptr || consumer->op != GGML_OP_MUL_MAT || consumer->inputs.size() != 2 || consumer->inputs[1] != input.id) { @@ -109,20 +105,19 @@ static bool is_packed_q8_consumer(const Graph & graph, const GraphNode * consume } const Value * weight = graph_value(graph, consumer->inputs[0]); const Value * output = graph_value(graph, consumer->output); - if (weight == nullptr || output == nullptr || input.type != GGML_TYPE_F32 || - output->type != GGML_TYPE_F32 || !input.contiguous || !weight->contiguous || !output->contiguous) { + if (weight == nullptr || output == nullptr || input.type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32 || + !input.contiguous || !weight->contiguous || !output->contiguous) { return false; } - const bool prefill = (weight->type == GGML_TYPE_Q5_K || weight->type == GGML_TYPE_IQ4_XS) && - input.ne[1] >= 256 && input.ne[1] <= 2048 && input.ne[1] % 256 == 0; + const bool prefill = (weight->type == GGML_TYPE_Q5_K || weight->type == GGML_TYPE_IQ4_XS) && input.ne[1] >= 256 && + input.ne[1] <= 2048 && input.ne[1] % 256 == 0; const bool decode = input.ne[1] >= 1 && input.ne[1] <= 5 && weight->alias_source.value < 0 && (weight->type == GGML_TYPE_Q4_K || (weight->type == GGML_TYPE_Q6_K && !graph.index().consumers(output->id).empty())); return (prefill || decode) && input.ne[0] >= 256 && input.ne[0] <= 32768 && input.ne[0] % 256 == 0 && - input.ne[2] == 1 && input.ne[3] == 1 && - weight->ne[0] == input.ne[0] && weight->ne[1] >= 64 && weight->ne[1] <= 262144 && weight->ne[1] % 64 == 0 && - weight->ne[2] == 1 && weight->ne[3] == 1 && output->ne[0] == weight->ne[1] && output->ne[1] == input.ne[1] && - output->ne[2] == 1 && output->ne[3] == 1; + input.ne[2] == 1 && input.ne[3] == 1 && weight->ne[0] == input.ne[0] && weight->ne[1] >= 64 && + weight->ne[1] <= 262144 && weight->ne[1] % 64 == 0 && weight->ne[2] == 1 && weight->ne[3] == 1 && + output->ne[0] == weight->ne[1] && output->ne[1] == input.ne[1] && output->ne[2] == 1 && output->ne[3] == 1; } static bool has_packed_q8_consumer(const Graph & graph, const Value & value) { @@ -166,6 +161,8 @@ struct RmsNormBinaryMatch { size_t binary_node_index = 0; int64_t hidden_size = 0; int64_t token_count = 0; + int64_t input_stride = 0; + size_t input_span = 0; BinaryKind op = BinaryKind::Add; float epsilon = 0.0f; @@ -174,6 +171,20 @@ struct RmsNormBinaryMatch { } }; +struct RmsNormMulAddMatch { + RmsNormBinaryMatch rms_binary; + const GraphNode * add_node = nullptr; + const Value * residual = nullptr; + const Value * output = nullptr; + size_t add_node_index = 0; + BinaryKind output_op = BinaryKind::Add; + bool binary_lhs = true; + + bool matched() const { + return rms_binary.matched() && add_node != nullptr && residual != nullptr && output != nullptr; + } +}; + struct RmsNormMatch { const GraphNode * rms_node = nullptr; const Value * input = nullptr; @@ -245,6 +256,34 @@ template static bool pairwise_distinct_storage_roots(const std::array return true; } +static bool rope_has_set_rows_consumer(const Graph & graph, const GraphNode & rope) { + std::vector pending = { rope.output }; + for (size_t i = 0; i < pending.size(); ++i) { + for (const GraphNode * consumer : graph.index().consumers(pending[i])) { + if (consumer == nullptr) { + continue; + } + if (consumer->op == GGML_OP_SET_ROWS) { + return true; + } + if (is_layout_alias_node(graph, *consumer)) { + pending.push_back(consumer->output); + } + } + } + return false; +} + +static bool fused_input_is_available(const Graph & graph, ValueId value, const std::vector & covered_nodes) { + const GraphNode * producer = graph.index().producer(value); + if (producer == nullptr) { + return true; + } + size_t producer_index = 0; + return graph.index().node_index(producer, producer_index) && producer_index < covered_nodes.size() && + covered_nodes[producer_index]; +} + static RmsNormMulRopeMatch match_rmsnorm_mul_rope_f32(const Graph & graph, const GraphNode * rms) { RmsNormMulRopeMatch match; if (rms == nullptr || rms->op != GGML_OP_RMS_NORM || rms->inputs.size() != 1 || !graph.has_index()) { @@ -290,8 +329,7 @@ static RmsNormMulRopeMatch match_rmsnorm_mul_rope_f32(const Graph & graph, const !positions->contiguous || weight->ne[0] != input->ne[0] || weight->ne[1] != 1 || weight->ne[2] != 1 || weight->ne[3] != 1 || input->ne[0] < 2 || input->ne[0] > 512 || input->ne[0] % 2 != 0 || input->ne[1] < 1 || input->ne[1] > 1048576 || input->ne[2] < 1 || input->ne[2] > 1048576 || input->ne[3] < 1 || - input->ne[3] > 1024 || input->nb[0] != sizeof(float) || output->nb[0] != sizeof(float) || - input->ne[2] > std::numeric_limits::max() / 4 || positions->element_count < 4 * input->ne[2]) { + input->ne[3] > 1024 || input->nb[0] != sizeof(float) || output->nb[0] != sizeof(float)) { return {}; } for (int i = 1; i < GGML_MAX_DIMS; ++i) { @@ -301,15 +339,21 @@ static RmsNormMulRopeMatch match_rmsnorm_mul_rope_f32(const Graph & graph, const } const bool sectioned_mode = rope_params->mode == GGML_ROPE_TYPE_MROPE || rope_params->mode == GGML_ROPE_TYPE_IMROPE; + const bool lfm_neox_mode = rope_params->mode == GGML_ROPE_TYPE_NEOX && input->ne[0] == 64 && + rope_params->n_dims == 64 && !rope_has_set_rows_consumer(graph, *rope); int64_t section_count = 0; - for (int section : rope_params->sections) { - if (section < 0) { - return {}; + if (sectioned_mode) { + for (int section : rope_params->sections) { + if (section < 0) { + return {}; + } + section_count += section; } - section_count += section; } - if (!sectioned_mode || rope_params->n_dims < 2 || rope_params->n_dims > input->ne[0] || - rope_params->n_dims % 2 != 0 || section_count != rope_params->n_dims / 2 || !std::isfinite(rms_params->eps) || + const int64_t required_positions = sectioned_mode ? 4 * input->ne[2] : input->ne[2]; + if ((!sectioned_mode && !lfm_neox_mode) || positions->element_count < required_positions || + rope_params->n_dims < 2 || rope_params->n_dims > input->ne[0] || rope_params->n_dims % 2 != 0 || + (sectioned_mode && section_count != rope_params->n_dims / 2) || !std::isfinite(rms_params->eps) || rms_params->eps <= 0.0f || !std::isfinite(rope_params->freq_base) || rope_params->freq_base <= 0.0f || !std::isfinite(rope_params->freq_scale) || rope_params->freq_scale <= 0.0f || !std::isfinite(rope_params->attn_factor) || rope_params->attn_factor <= 0.0f || @@ -383,7 +427,10 @@ static RmsNormMatch match_rmsnorm_f32(const Graph & graph, const GraphNode * nod return match; } -static RmsNormBinaryMatch match_rmsnorm_binary_f32(const Graph & graph, const GraphNode * node, size_t node_index) { +static RmsNormBinaryMatch match_rmsnorm_binary_f32(const Graph & graph, + const GraphNode * node, + size_t node_index, + bool allow_strided_input = false) { RmsNormBinaryMatch match; if (node == nullptr || node->op != GGML_OP_RMS_NORM || node->inputs.size() != 1 || !graph.has_index()) { return match; @@ -399,28 +446,18 @@ static RmsNormBinaryMatch match_rmsnorm_binary_f32(const Graph & graph, const Gr } const GraphNode * binary_node = consumers.front(); size_t binary_node_index; - if (binary_node == nullptr || !is_binary_op(binary_node->op) || binary_node->inputs.size() != 2 || - !graph.index().node_index(binary_node, binary_node_index)) { + const CommonBinaryMatch binary_match = + common_match_binary_consumer(binary_node, node->output, kCommonArithmeticBinaryKinds); + if (!binary_match.matched() || !graph.index().node_index(binary_node, binary_node_index)) { return {}; } - const BinaryParams * binary_params = op_params_as(binary_node->params); - if (binary_params == nullptr || !binary_kind_supported(binary_params->op)) { + if (!binary_match.producer_is_lhs && common_binary_kind_requires_order(binary_match.kind)) { return {}; } - const bool rms_is_lhs = binary_node->inputs[0] == node->output; - const bool rms_is_rhs = binary_node->inputs[1] == node->output; - if (!rms_is_lhs && !rms_is_rhs) { - return {}; - } - if (rms_is_rhs && binary_kind_requires_order(binary_params->op)) { - return {}; - } - - const ValueId rhs_id = rms_is_lhs ? binary_node->inputs[1] : binary_node->inputs[0]; const Value * input = graph_value(graph, node->inputs[0]); const Value * rms = graph_value(graph, node->output); - const Value * rhs = graph_value(graph, rhs_id); + const Value * rhs = graph_value(graph, binary_match.operand); const Value * output = graph_value(graph, binary_node->output); if (input == nullptr || rms == nullptr || rhs == nullptr || output == nullptr) { return {}; @@ -429,7 +466,7 @@ static RmsNormBinaryMatch match_rmsnorm_binary_f32(const Graph & graph, const Gr output->type != GGML_TYPE_F32) { return {}; } - if (!input->contiguous || !rms->contiguous || !rhs->contiguous || !output->contiguous) { + if (!rms->contiguous || !rhs->contiguous || !output->contiguous) { return {}; } if (!same_shape(*input, *rms) || !same_shape(*input, *output)) { @@ -447,6 +484,13 @@ static RmsNormBinaryMatch match_rmsnorm_binary_f32(const Graph & graph, const Gr if (!is_supported_token_count(token_count)) { return {}; } + int64_t input_stride = 0; + size_t input_span = 0; + if (!supported_rmsnorm_input_layout(*input, hidden_size, token_count, input_stride, input_span) || + (!allow_strided_input && input_stride != hidden_size) || + !pairwise_distinct_storage_roots(std::array{ input, rhs, output })) { + return {}; + } match.rms_node = node; match.binary_node = binary_node; @@ -457,11 +501,69 @@ static RmsNormBinaryMatch match_rmsnorm_binary_f32(const Graph & graph, const Gr match.binary_node_index = binary_node_index; match.hidden_size = hidden_size; match.token_count = token_count; - match.op = binary_params->op; + match.input_stride = input_stride; + match.input_span = input_span; + match.op = binary_match.kind; match.epsilon = rms_params->eps; return match; } +static RmsNormMulAddMatch match_rmsnorm_mul_add_f32(const DispatchMatchContext & context) { + RmsNormMulAddMatch match; + match.rms_binary = match_rmsnorm_binary_f32(context.graph, context.root_node, context.root_index); + if (!match.rms_binary.matched()) { + return {}; + } + + const std::vector & consumers = context.graph.index().consumers(match.rms_binary.output->id); + if (consumers.size() != 1) { + return {}; + } + const GraphNode * add = consumers.front(); + size_t add_index = 0; + const CommonBinaryMatch output_binary = + common_match_binary_consumer(add, match.rms_binary.output->id, kCommonArithmeticBinaryKinds); + if (!output_binary.matched() || !context.graph.index().node_index(add, add_index)) { + return {}; + } + const Value * residual = graph_value(context.graph, output_binary.operand); + const Value * output = graph_value(context.graph, add->output); + const Value * rms_output = graph_value(context.graph, match.rms_binary.rms_node->output); + if (residual == nullptr || output == nullptr || rms_output == nullptr || residual->type != GGML_TYPE_F32 || + output->type != GGML_TYPE_F32 || !residual->contiguous || !output->contiguous || + !same_shape(*residual, *match.rms_binary.input) || !same_shape(*output, *match.rms_binary.input) || + !pairwise_distinct_storage_roots(std::array{ + match.rms_binary.input, match.rms_binary.rhs, rms_output, match.rms_binary.output, residual, output })) { + return {}; + } + const size_t nodes_size = context.covered_nodes.size(); + if (match.rms_binary.rms_node_index >= nodes_size || match.rms_binary.binary_node_index >= nodes_size || + add_index >= nodes_size || context.covered_nodes[match.rms_binary.rms_node_index] || + context.covered_nodes[match.rms_binary.binary_node_index] || context.covered_nodes[add_index]) { + return {}; + } + const auto available = [&](const Value & value) { + const GraphNode * producer = context.graph.index().producer(value.id); + if (producer == nullptr) { + return true; + } + size_t producer_index = 0; + return context.graph.index().node_index(producer, producer_index) && producer_index < nodes_size && + context.covered_nodes[producer_index]; + }; + if (!available(*match.rms_binary.input) || !available(*match.rms_binary.rhs) || !available(*residual)) { + return {}; + } + + match.add_node = add; + match.residual = residual; + match.output = output; + match.add_node_index = add_index; + match.output_op = output_binary.kind; + match.binary_lhs = output_binary.producer_is_lhs; + return match; +} + static AddRmsNormBinarySymmetricI4Match match_add_rmsnorm_binary_symmetric_i4(const Graph & graph, const GraphNode * node, size_t node_index) { @@ -512,9 +614,7 @@ static AddRmsNormBinarySymmetricI4Match match_add_rmsnorm_binary_symmetric_i4(co return match; } -static RmsNormGateMatch match_rmsnorm_gate(const Graph & graph, - const GraphNode * node, - size_t node_index) { +static RmsNormGateMatch match_rmsnorm_gate(const Graph & graph, const GraphNode * node, size_t node_index) { RmsNormGateMatch match; const RmsNormBinaryMatch rms_binary = match_rmsnorm_binary_f32(graph, node, node_index); if (!rms_binary.matched() || rms_binary.op != BinaryKind::Mul || rms_binary.hidden_size % 64 != 0) { @@ -588,6 +688,12 @@ static bool match_rmsnorm_mul_rope_f32_dispatch(const DispatchMatchContext & con if (!fused.matched()) { return false; } + if (fused.rope_params->mode == GGML_ROPE_TYPE_NEOX && + (!fused_input_is_available(context.graph, fused.input->id, context.covered_nodes) || + !fused_input_is_available(context.graph, fused.weight->id, context.covered_nodes) || + !fused_input_is_available(context.graph, fused.positions->id, context.covered_nodes))) { + return false; + } Dispatch dispatch; dispatch.kernel = make_kernel_specialization(kRmsNormMulRopeF32Kernel); @@ -609,8 +715,10 @@ static bool match_rmsnorm_mul_rope_f32_dispatch(const DispatchMatchContext & con config.emplace("ggml.rmsnorm_mul_rope.output_stride3", to_config_value(static_cast(fused.output->nb[3] / sizeof(float)))); config.emplace("ggml.rmsnorm_mul_rope.n_dims", to_config_value(static_cast(fused.rope_params->n_dims))); + const bool neox_mode = fused.rope_params->mode == GGML_ROPE_TYPE_NEOX; config.emplace("ggml.rmsnorm_mul_rope.section0", - to_config_value(static_cast(fused.rope_params->sections[0]))); + to_config_value(neox_mode ? static_cast(fused.rope_params->n_dims / 2) : + static_cast(fused.rope_params->sections[0]))); config.emplace("ggml.rmsnorm_mul_rope.section1", to_config_value(static_cast(fused.rope_params->sections[1]))); config.emplace("ggml.rmsnorm_mul_rope.section2", @@ -697,9 +805,8 @@ static bool match_rmsnorm_binary_symmetric_i4_dispatch(const DispatchMatchContex } const RmsNormBinaryMatch fused = match_rmsnorm_binary_f32(context.graph, &nodes[context.root_index], context.root_index); - if (!fused.matched() || fused.op != BinaryKind::Mul || - !is_supported_symmetric_i4_hidden_size(fused.hidden_size) || fused.token_count > 16 || - fused.rms_node_index >= context.covered_nodes.size() || + if (!fused.matched() || fused.op != BinaryKind::Mul || !is_supported_symmetric_i4_hidden_size(fused.hidden_size) || + fused.token_count > 16 || fused.rms_node_index >= context.covered_nodes.size() || fused.binary_node_index >= context.covered_nodes.size() || context.covered_nodes[fused.rms_node_index] || context.covered_nodes[fused.binary_node_index] || !common_has_symmetric_i4_lowrow_consumer(context.graph, *fused.output) || @@ -751,16 +858,18 @@ static bool match_rmsnorm_gate_dispatch(const DispatchMatchContext & context, Di return false; } - const bool use_i4 = fused.gate_op == UnaryKind::Silu && - common_has_symmetric_i4_lowrow_consumer(context.graph, *fused.output); + const bool use_i4 = + fused.gate_op == UnaryKind::Silu && common_has_symmetric_i4_lowrow_consumer(context.graph, *fused.output); if (!use_i4) { if (fused.hidden_size > 1024) { return false; } const Value * packed_input = fused.output; - const GraphNode * consumer = common_find_only_consumer_with_op(context.graph, packed_input->id, GGML_OP_MUL_MAT); + const GraphNode * consumer = + common_find_only_consumer_with_op(context.graph, packed_input->id, GGML_OP_MUL_MAT); if (consumer == nullptr) { - const GraphNode * reshape = common_find_only_consumer_with_op(context.graph, packed_input->id, GGML_OP_RESHAPE); + const GraphNode * reshape = + common_find_only_consumer_with_op(context.graph, packed_input->id, GGML_OP_RESHAPE); const Value * reshaped = reshape != nullptr ? graph_value(context.graph, reshape->output) : nullptr; if (reshaped != nullptr && is_layout_alias_node(context.graph, *reshape) && reshaped->contiguous && same_full_value_range(*packed_input, *reshaped)) { @@ -774,20 +883,21 @@ static bool match_rmsnorm_gate_dispatch(const DispatchMatchContext & context, Di projection.weight->alias_source.value < 0 && common_mul_mat_uses_k16_major_f16(projection.weight_format, projection.input_size, projection.output_size, projection.token_count); - const bool q8_output = !packed_output && packed_input->ne[1] <= 5 && - is_packed_q8_consumer(context.graph, consumer, *packed_input); - const size_t activation_bytes = q8_output ? q8_1_x4_byte_count(fused.token_count, fused.hidden_size) : - fused.output->byte_count / 2; + const bool q8_output = + !packed_output && packed_input->ne[1] <= 5 && is_packed_q8_consumer(context.graph, consumer, *packed_input); + const size_t activation_bytes = + q8_output ? q8_1_x4_byte_count(fused.token_count, fused.hidden_size) : fused.output->byte_count / 2; const ValueId activation = context.next_plan_value; const char * activation_name = q8_output ? "common.rmsnorm_gate.q8_1_x4" : - packed_output ? "common.rmsnorm_gate.k16_major_f16" : "common.rmsnorm_gate.f16"; + packed_output ? "common.rmsnorm_gate.k16_major_f16" : + "common.rmsnorm_gate.f16"; Dispatch dispatch; - dispatch.kernel = make_kernel_specialization(q8_output ? kRmsNormGateF32Q8_1X4Kernel : kRmsNormGateF32F16Kernel); + dispatch.kernel = + make_kernel_specialization(q8_output ? kRmsNormGateF32Q8_1X4Kernel : kRmsNormGateF32F16Kernel); dispatch.kernel.integer_parameters.emplace("token_count", fused.token_count); dispatch.kernel.compile_parameters.emplace("ggml.rmsnorm_gate_f32.hidden_size", to_config_value(fused.hidden_size)); - dispatch.kernel.compile_parameters.emplace("ggml.rmsnorm_gate_f32.rms_epsilon", - to_config_value(fused.epsilon)); + dispatch.kernel.compile_parameters.emplace("ggml.rmsnorm_gate_f32.rms_epsilon", to_config_value(fused.epsilon)); dispatch.kernel.compile_parameters.emplace("ggml.rmsnorm_gate_f32.gate_op", std::to_string(unary_kind_config_value(fused.gate_op))); dispatch.kernel.compile_parameters.emplace("ggml.rmsnorm_gate_f32.f16_output_row_width", @@ -798,12 +908,15 @@ static bool match_rmsnorm_gate_dispatch(const DispatchMatchContext & context, Di dispatch.bindings.push_back({ fused.output->id, 0, fused.output->byte_count }); dispatch.bindings.push_back({ activation, 0, activation_bytes }); Status status; - const bool recorded = packed_output ? + const bool recorded = + packed_output ? match.metadata.append_generated_resource( - { packed_input->id, GeneratedResourceRole::F16K16Major, activation, activation_bytes, {} }, status) : + { packed_input->id, GeneratedResourceRole::F16K16Major, activation, activation_bytes, {} }, + status) : match.metadata.append_alternate_value( - { fused.output->id, activation, q8_output ? GGML_TYPE_Q8_1 : GGML_TYPE_F16, - activation_bytes, activation_name }, status); + { fused.output->id, activation, q8_output ? GGML_TYPE_Q8_1 : GGML_TYPE_F16, activation_bytes, + activation_name }, + status); if (!recorded) { match.status.append(status); return false; @@ -874,8 +987,8 @@ static bool has_qualified_rmsnorm_k16_consumer(const Graph & graph, const RmsNor common_match_mul_mat_any_format(graph, consumer, kRmsNormBinaryF32K16Kernel, false); if (projection.matched() && projection.input->id == rms.output->id && projection.weight->alias_source.value < 0 && - common_mul_mat_uses_k16_major_f16(projection.weight_format, projection.input_size, - projection.output_size, projection.token_count)) { + common_mul_mat_uses_k16_major_f16(projection.weight_format, projection.input_size, projection.output_size, + projection.token_count)) { return true; } } @@ -883,8 +996,8 @@ static bool has_qualified_rmsnorm_k16_consumer(const Graph & graph, const RmsNor } static bool has_qualified_rmsnorm_f16_consumer(const Graph & graph, const RmsNormBinaryMatch & rms) { - if (!graph.has_index() || rms.token_count < 128 || rms.output->ne[1] != rms.token_count || - rms.output->ne[2] != 1 || rms.output->ne[3] != 1) { + if (!graph.has_index() || rms.token_count < 128 || rms.output->ne[1] != rms.token_count || rms.output->ne[2] != 1 || + rms.output->ne[3] != 1) { return false; } for (const GraphNode * consumer : graph.index().consumers(rms.output->id)) { @@ -913,17 +1026,20 @@ static bool match_rmsnorm_binary_f32_dispatch(const DispatchMatchContext & conte return false; } const RmsNormBinaryMatch rms_match = - match_rmsnorm_binary_f32(context.graph, &nodes[context.root_index], context.root_index); + match_rmsnorm_binary_f32(context.graph, &nodes[context.root_index], context.root_index, true); if (!rms_match.matched() || rms_match.rms_node_index >= context.covered_nodes.size() || rms_match.binary_node_index >= context.covered_nodes.size() || context.covered_nodes[rms_match.rms_node_index] || context.covered_nodes[rms_match.binary_node_index]) { return false; } - const bool packed_output = has_qualified_rmsnorm_k16_consumer(context.graph, rms_match); - const bool f16_output = !packed_output && has_qualified_rmsnorm_f16_consumer(context.graph, rms_match); + const bool strided_input = rms_match.input_stride != rms_match.hidden_size; + const bool packed_output = !strided_input && has_qualified_rmsnorm_k16_consumer(context.graph, rms_match); + const bool f16_output = + !strided_input && !packed_output && has_qualified_rmsnorm_f16_consumer(context.graph, rms_match); Dispatch dispatch; - dispatch.kernel = make_kernel_specialization(packed_output ? kRmsNormBinaryF32K16Kernel : + dispatch.kernel = make_kernel_specialization(strided_input ? kRmsNormBinaryStridedF32Kernel : + packed_output ? kRmsNormBinaryF32K16Kernel : f16_output ? kRmsNormBinaryF32F16Kernel : kRmsNormBinaryF32Kernel); dispatch.kernel.integer_parameters.emplace("token_count", rms_match.token_count); @@ -933,7 +1049,14 @@ static bool match_rmsnorm_binary_f32_dispatch(const DispatchMatchContext & conte to_config_value(rms_match.epsilon)); dispatch.kernel.compile_parameters.emplace("ggml.rmsnorm_binary_f32.op", std::to_string(binary_kind_config_value(rms_match.op))); - dispatch.bindings.push_back({ rms_match.input->id, 0, rms_match.input->byte_count }); + if (strided_input) { + dispatch.kernel.compile_parameters.emplace("ggml.rmsnorm_binary_f32.input_stride", + to_config_value(rms_match.input_stride)); + } + dispatch.bindings.push_back( + strided_input ? + DispatchBinding{ rms_match.input->storage_root, rms_match.input->storage_offset, rms_match.input_span } : + DispatchBinding{ rms_match.input->id, 0, rms_match.input->byte_count }); dispatch.bindings.push_back({ rms_match.rhs->id, 0, rms_match.rhs->byte_count }); dispatch.bindings.push_back({ rms_match.output->id, 0, rms_match.output->byte_count }); @@ -967,6 +1090,34 @@ static bool match_rmsnorm_binary_f32_dispatch(const DispatchMatchContext & conte return true; } +static bool match_rmsnorm_mul_add_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { + const RmsNormMulAddMatch fused = match_rmsnorm_mul_add_f32(context); + if (!fused.matched()) { + return false; + } + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kRmsNormMulAddF32Kernel); + dispatch.kernel.integer_parameters.emplace("token_count", fused.rms_binary.token_count); + dispatch.kernel.compile_parameters.emplace("ggml.rmsnorm_mul_add_f32.hidden_size", + to_config_value(fused.rms_binary.hidden_size)); + dispatch.kernel.compile_parameters.emplace("ggml.rmsnorm_mul_add_f32.rms_epsilon", + to_config_value(fused.rms_binary.epsilon)); + common_set_binary_compile_parameters(dispatch.kernel, "ggml.rmsnorm_mul_add_f32.normalized_op", nullptr, + fused.rms_binary.op, true); + common_set_binary_compile_parameters(dispatch.kernel, "ggml.rmsnorm_mul_add_f32.output_op", + "ggml.rmsnorm_mul_add_f32.binary_lhs", fused.output_op, fused.binary_lhs); + dispatch.bindings.push_back({ fused.rms_binary.input->id, 0, fused.rms_binary.input->byte_count }); + dispatch.bindings.push_back({ fused.rms_binary.rhs->id, 0, fused.rms_binary.rhs->byte_count }); + dispatch.bindings.push_back({ fused.residual->id, 0, fused.residual->byte_count }); + dispatch.bindings.push_back({ fused.output->id, 0, fused.output->byte_count }); + match.covered_nodes.push_back(fused.rms_binary.rms_node_index); + match.covered_nodes.push_back(fused.rms_binary.binary_node_index); + match.covered_nodes.push_back(fused.add_node_index); + match.dispatches.push_back(std::move(dispatch)); + return true; +} + static bool match_rmsnorm_binary_q8_1_x4_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { const std::vector & nodes = context.graph.nodes(); if (context.root_index >= nodes.size()) { @@ -1088,6 +1239,14 @@ void register_rmsnorm_dispatches(DispatchRegistryBuilder & registry) { DispatchSource::Common, match_rmsnorm_mul_rope_f32_dispatch, }); + registry.add({ + "common.rmsnorm_mul_add_f32", + GGML_OP_RMS_NORM, + DispatchMatchKind::Fused, + 250, + DispatchSource::Common, + match_rmsnorm_mul_add_f32_dispatch, + }); registry.add({ "common.rmsnorm_binary_q8_1_x4", GGML_OP_RMS_NORM, diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rope-set-rows.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rope-set-rows.cpp index 0399fe094524..0d7542c7a493 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rope-set-rows.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rope-set-rows.cpp @@ -6,6 +6,7 @@ #include "graph/graph-matcher.h" #include "kernel-corpus/kernel-corpus-catalog-verify.h" +#include #include #include #include @@ -17,7 +18,11 @@ namespace ggml::hrx { namespace { -static constexpr KernelCatalogRef kRopeF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_rope_f32"); +static constexpr KernelCatalogRef kRopeF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_rope_f32"); +static constexpr KernelCatalogRef kRopeConcatF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_rope_concat_f32"); +static constexpr KernelCatalogRef kRopeScaleF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_rope_scale_f32"); +static constexpr KernelCatalogRef kRopeTokenScaleF32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_rope_token_scale_f32"); static constexpr KernelCatalogRef kSetRowsKernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_set_rows"); static constexpr KernelCatalogRef kRopeSetRowsF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_rope_set_rows_f32"); @@ -273,6 +278,28 @@ struct RopeSetRowsMatch { bool matched() const { return rope.matched() && set_rows.matched() && published_rows != nullptr; } }; +struct RopeScaleMatch { + RopeMatch rope; + const GraphNode * scale_node = nullptr; + const Value * scales = nullptr; + const Value * output = nullptr; + float scalar = 1.0f; + bool per_token = false; + + bool matched() const { return rope.matched() && scale_node != nullptr && output != nullptr; } +}; + +struct RopeConcatMatch { + RopeMatch rope; + const GraphNode * concat = nullptr; + const Value * prefix = nullptr; + const Value * output = nullptr; + size_t concat_idx = 0; + size_t input_span_bytes = 0; + + bool matched() const { return rope.matched() && concat != nullptr && prefix != nullptr && output != nullptr; } +}; + static RopeMatch match_rope_f32(const Graph & graph, const GraphNode * node, Status * status = nullptr) { RopeMatch match; if (node == nullptr || node->op != GGML_OP_ROPE || node->inputs.size() < 2 || node->inputs.size() > 3) { @@ -490,6 +517,31 @@ static Dispatch make_rope_dispatch(const DispatchMatchContext & context, return dispatch; } +static Dispatch make_rope_scale_dispatch(const DispatchMatchContext & context, + const RopeScaleMatch & match, + DispatchMatch & dispatch_match) { + const auto [theta, freq_factors] = + add_rope_frequency_bindings(context, match.rope, dispatch_match, "common.rope_scale"); + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(match.per_token ? kRopeTokenScaleF32Kernel : kRopeScaleF32Kernel); + dispatch.kernel.integer_parameters.emplace("token_count", match.rope.token_count); + add_rope_compile_parameters(dispatch, match.rope, "ggml.rope_f32"); + if (!match.per_token) { + dispatch.kernel.compile_parameters.emplace("ggml.rope_scale_f32.scale", rope_mscale_config_value(match.scalar)); + } + dispatch.bindings.push_back({ match.rope.positions->id, 0, match.rope.positions->byte_count }); + dispatch.bindings.push_back( + { match.rope.input->storage_root, match.rope.input->storage_offset, match.rope.input_span_bytes }); + dispatch.bindings.push_back({ theta, 0, match.rope.theta_bytes }); + dispatch.bindings.push_back({ freq_factors, 0, match.rope.freq_factors_bytes }); + if (match.per_token) { + dispatch.bindings.push_back( + { match.scales->storage_root, match.scales->storage_offset, match.scales->byte_count }); + } + dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + return dispatch; +} + static Dispatch make_set_rows_dispatch(const SetRowsMatch & match) { Dispatch dispatch; dispatch.kernel = make_kernel_specialization(kSetRowsKernel); @@ -525,6 +577,99 @@ static Dispatch make_rope_set_rows_dispatch(const DispatchMatchContext & context return dispatch; } +static Dispatch make_rope_concat_dispatch(const DispatchMatchContext & context, + const RopeConcatMatch & match, + DispatchMatch & dispatch_match) { + const auto [theta, freq_factors] = + add_rope_frequency_bindings(context, match.rope, dispatch_match, "common.rope_concat"); + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kRopeConcatF32Kernel); + dispatch.kernel.integer_parameters.emplace("token_count", match.rope.token_count); + const char * prefix = "ggml.rope_concat_f32."; + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + "prefix_size", + to_config_value(match.prefix->ne[0])); + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + "rope_size", + to_config_value(match.rope.head_size)); + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + "n_dims", to_config_value(match.rope.n_dims)); + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + "head_count", + to_config_value(match.rope.head_count)); + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + "token_capacity", + to_config_value(match.rope.token_count)); + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + "input_stride1", + to_config_value(match.rope.input_stride1)); + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + "input_stride2", + to_config_value(match.rope.input_stride2)); + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + "input_span", + to_config_value(match.input_span_bytes / sizeof(float))); + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + "mscale", + rope_mscale_config_value(match.rope.rope_mscale)); + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + "mode", to_config_value(match.rope.mode)); + dispatch.bindings.push_back({ match.rope.positions->id, 0, match.rope.positions->byte_count }); + dispatch.bindings.push_back({ match.prefix->storage_root, match.prefix->storage_offset, match.input_span_bytes }); + dispatch.bindings.push_back({ theta, 0, match.rope.theta_bytes }); + dispatch.bindings.push_back({ freq_factors, 0, match.rope.freq_factors_bytes }); + dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + return dispatch; +} + +static RopeConcatMatch match_rope_concat_f32(const Graph & graph, const GraphNode * node) { + RopeConcatMatch match; + match.rope = match_rope_f32(graph, node); + if (!match.rope.matched() || match.rope.mode != GGML_ROPE_TYPE_NEOX || !graph.has_index() || + !graph.index().has_single_consumer(node->output)) { + return {}; + } + + const GraphNode * concat = graph.index().consumers(node->output).front(); + if (concat == nullptr || concat->op != GGML_OP_CONCAT || concat->inputs.size() != 2 || + concat->inputs[1] != node->output || !graph.index().node_index(concat, match.concat_idx)) { + return {}; + } + + const Value * prefix = graph_value(graph, concat->inputs[0]); + const Value * output = graph_value(graph, concat->output); + if (prefix == nullptr || output == nullptr || prefix->type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32 || + prefix->ne[0] < 4 || prefix->ne[0] > 1024 || prefix->ne[0] % 4 != 0 || + prefix->ne[0] + match.rope.head_size > 1024 || !is_packed_f32_rope_layout(*output) || + output->ne[0] != prefix->ne[0] + match.rope.head_size || prefix->nb[0] != sizeof(float) || + prefix->storage_root != match.rope.input->storage_root || prefix->storage_root == output->storage_root || + match.rope.input->storage_offset != + prefix->storage_offset + static_cast(prefix->ne[0]) * sizeof(float)) { + return {}; + } + for (int dim = 1; dim < GGML_MAX_DIMS; ++dim) { + if (prefix->ne[dim] != match.rope.input->ne[dim] || prefix->ne[dim] != output->ne[dim] || + prefix->nb[dim] != match.rope.input->nb[dim]) { + return {}; + } + } + + size_t prefix_span_bytes = 0; + if (!strided_f32_storage_span_bytes(*prefix, prefix_span_bytes)) { + return {}; + } + const size_t rope_offset_bytes = match.rope.input->storage_offset - prefix->storage_offset; + match.input_span_bytes = std::max(prefix_span_bytes, rope_offset_bytes + match.rope.input_span_bytes); + match.concat = concat; + match.prefix = prefix; + match.output = output; + return match; +} + +static bool match_rope_concat_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + const RopeConcatMatch match = match_rope_concat_f32(context.graph, context.root_node); + if (!match.matched() || + !append_covered_node_index_once(context.graph, context.covered_nodes, context.root_node, + dispatch_match.covered_nodes) || + !append_covered_node_index_once(context.graph, context.covered_nodes, match.concat, + dispatch_match.covered_nodes)) { + return false; + } + + dispatch_match.dispatches.push_back(make_rope_concat_dispatch(context, match, dispatch_match)); + return true; +} + static bool find_rope_set_rows_consumer(const Graph & graph, const GraphNode & rope, RopeSetRowsMatch & match) { if (!graph.has_index() || !graph.index().has_single_consumer(rope.output)) { return false; @@ -571,6 +716,84 @@ static RopeSetRowsMatch match_rope_set_rows_f32(const Graph & graph, const Graph return match; } +static RopeScaleMatch match_rope_scale_f32(const Graph & graph, const GraphNode * node) { + RopeScaleMatch match; + match.rope = match_rope_f32(graph, node); + if (!match.rope.matched() || !graph.has_index() || !graph.index().has_single_consumer(node->output)) { + return {}; + } + + const GraphNode * scale_node = graph.index().consumers(node->output).front(); + if (scale_node == nullptr) { + return {}; + } + const Value * output = graph_value(graph, scale_node->output); + if (output == nullptr || output->type != GGML_TYPE_F32 || !output->contiguous || output->alias_source.value >= 0 || + !same_shape(*match.rope.output, *output)) { + return {}; + } + + if (scale_node->op == GGML_OP_SCALE && scale_node->inputs.size() == 1 && + scale_node->inputs.front() == node->output) { + const ScaleParams * params = op_params_as(scale_node->params); + if (params == nullptr || params->bias != 0.0f || !std::isfinite(params->scale)) { + return {}; + } + match.scalar = params->scale; + } else if (scale_node->op == GGML_OP_MUL && scale_node->inputs.size() == 2) { + const BinaryParams * params = op_params_as(scale_node->params); + if (params == nullptr || params->op != BinaryKind::Mul) { + return {}; + } + ValueId scale_id; + if (scale_node->inputs[0] == node->output) { + scale_id = scale_node->inputs[1]; + } else if (scale_node->inputs[1] == node->output) { + scale_id = scale_node->inputs[0]; + } else { + return {}; + } + const Value * scales = graph_value(graph, scale_id); + if (scales == nullptr || scales->type != GGML_TYPE_F32 || !scales->contiguous || scales->ne[0] != 1 || + scales->ne[1] != 1 || scales->ne[2] != match.rope.token_count || scales->ne[3] != 1) { + return {}; + } + match.scales = scales; + match.per_token = true; + } else { + return {}; + } + + match.scale_node = scale_node; + match.output = output; + return match; +} + +static bool value_is_available(const Graph & graph, ValueId value, const std::vector & covered_nodes) { + const GraphNode * producer = graph.index().producer(value); + if (producer == nullptr) { + return true; + } + size_t producer_index = 0; + return graph.index().node_index(producer, producer_index) && producer_index < covered_nodes.size() && + covered_nodes[producer_index]; +} + +static bool match_rope_scale_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + const RopeScaleMatch match = match_rope_scale_f32(context.graph, context.root_node); + if (!match.matched() || + (match.per_token && !value_is_available(context.graph, match.scales->id, context.covered_nodes)) || + !append_covered_node_index_once(context.graph, context.covered_nodes, context.root_node, + dispatch_match.covered_nodes) || + !append_covered_node_index_once(context.graph, context.covered_nodes, match.scale_node, + dispatch_match.covered_nodes)) { + return false; + } + + dispatch_match.dispatches.push_back(make_rope_scale_dispatch(context, match, dispatch_match)); + return true; +} + static bool match_rope_set_rows_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { const RopeSetRowsMatch match = match_rope_set_rows_f32(context.graph, context.root_node); if (!match.matched()) { @@ -615,6 +838,22 @@ static bool match_set_rows_dispatch(const DispatchMatchContext & context, Dispat } // namespace void register_rope_set_rows_dispatches(DispatchRegistryBuilder & registry) { + registry.add({ + "common.rope_scale.f32", + GGML_OP_ROPE, + DispatchMatchKind::Fused, + 300, + DispatchSource::Common, + match_rope_scale_f32_dispatch, + }); + registry.add({ + "common.rope_concat.f32", + GGML_OP_ROPE, + DispatchMatchKind::Fused, + 300, + DispatchSource::Common, + match_rope_concat_f32_dispatch, + }); registry.add({ "common.rope_set_rows.f32", GGML_OP_ROPE, diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-scale.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-scale.cpp index 5c23046162a1..a5dff78f70aa 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-scale.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-scale.cpp @@ -1,5 +1,7 @@ #include "dispatch-scale.h" +#include "dispatch-binary-common.h" +#include "dispatch-layout-utils.h" #include "ggml.h" #include "kernel-corpus/kernel-corpus-catalog-verify.h" @@ -12,7 +14,8 @@ namespace ggml::hrx { namespace { -static constexpr KernelCatalogRef kScaleF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_scale_f32"); +static constexpr KernelCatalogRef kScaleF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_scale_f32"); +static constexpr KernelCatalogRef kScaleAddF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_scale_add_f32"); static bool same_shape(const Value & lhs, const Value & rhs) { for (int i = 0; i < GGML_MAX_DIMS; ++i) { @@ -51,6 +54,42 @@ static bool distinct_storage(const Value & input, const Value & output) { return input.storage != output.storage; } +static bool storage_ranges_disjoint(const Value & lhs, + size_t lhs_byte_count, + const Value & rhs, + size_t rhs_byte_count) { + if (lhs.storage != rhs.storage) { + return true; + } + if (lhs.storage_offset > std::numeric_limits::max() - lhs_byte_count || + rhs.storage_offset > std::numeric_limits::max() - rhs_byte_count) { + return false; + } + return lhs.storage_offset + lhs_byte_count <= rhs.storage_offset || + rhs.storage_offset + rhs_byte_count <= lhs.storage_offset; +} + +static bool value_is_available(const Graph & graph, ValueId value, const std::vector & covered_nodes) { + const GraphNode * producer = graph.index().producer(value); + if (producer == nullptr) { + return true; + } + size_t producer_index = 0; + return graph.index().node_index(producer, producer_index) && producer_index < covered_nodes.size() && + covered_nodes[producer_index]; +} + +static bool value_is_available_or_layout_alias(const Graph & graph, + ValueId value, + const std::vector & covered_nodes) { + if (value_is_available(graph, value, covered_nodes)) { + return true; + } + const GraphNode * producer = graph.index().producer(value); + return producer != nullptr && is_layout_alias_node(graph, *producer) && producer->inputs.size() == 1 && + value_is_available(graph, producer->inputs[0], covered_nodes); +} + static bool supported_source_layout(const Graph & graph, const Value & value) { if (value.alias_source.value < 0) { return true; @@ -69,6 +108,91 @@ static std::string format_float_config(float value) { return out.str(); } +static void add_strided_source_config(Dispatch & dispatch, + const char * prefix, + const Value & source, + size_t source_byte_count) { + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + "stride1", + std::to_string(source.nb[1] / sizeof(float))); + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + "stride2", + std::to_string(source.nb[2] / sizeof(float))); + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + "stride3", + std::to_string(source.nb[3] / sizeof(float))); + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + "span", + std::to_string(source_byte_count / sizeof(float))); +} + +static bool match_scale_add_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { + const GraphNode * scale_node = context.root_node; + if (scale_node == nullptr || scale_node->op != GGML_OP_SCALE || scale_node->inputs.size() != 1 || + !context.graph.has_index() || context.root_index >= context.covered_nodes.size() || + context.covered_nodes[context.root_index]) { + return false; + } + + const ScaleParams * scale_params = op_params_as(scale_node->params); + if (scale_params == nullptr || !context.graph.index().has_single_consumer(scale_node->output)) { + return false; + } + const GraphNode * binary_node = context.graph.index().consumers(scale_node->output).front(); + const CommonBinaryMatch binary = + common_match_binary_consumer(binary_node, scale_node->output, kCommonArithmeticBinaryKinds); + if (!binary.matched()) { + return false; + } + + size_t binary_index = 0; + if (!context.graph.index().node_index(binary_node, binary_index) || binary_index >= context.covered_nodes.size() || + context.covered_nodes[binary_index]) { + return false; + } + + const Value * input = graph_value(context.graph, scale_node->inputs[0]); + const Value * scaled = graph_value(context.graph, scale_node->output); + const Value * residual = graph_value(context.graph, binary.operand); + const Value * output = graph_value(context.graph, binary_node->output); + if (input == nullptr || scaled == nullptr || residual == nullptr || output == nullptr || + input->type != GGML_TYPE_F32 || scaled->type != GGML_TYPE_F32 || residual->type != GGML_TYPE_F32 || + output->type != GGML_TYPE_F32 || !same_shape(*input, *scaled) || !same_shape(*input, *residual) || + !same_shape(*input, *output) || !positive_shape(*output) || !output->contiguous || + !packed_f32_layout(*output) || scaled->alias_source.value >= 0 || output->alias_source.value >= 0 || + static_cast(output->element_count) > std::numeric_limits::max()) { + return false; + } + + size_t input_byte_count = 0; + size_t residual_byte_count = 0; + if (!strided_f32_storage_span_bytes(*input, input_byte_count) || + !strided_f32_storage_span_bytes(*residual, residual_byte_count) || + !storage_ranges_disjoint(*input, input_byte_count, *output, output->byte_count) || + !storage_ranges_disjoint(*residual, residual_byte_count, *output, output->byte_count) || + !value_is_available_or_layout_alias(context.graph, input->id, context.covered_nodes) || + !value_is_available_or_layout_alias(context.graph, residual->id, context.covered_nodes)) { + return false; + } + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kScaleAddF32Kernel); + dispatch.kernel.integer_parameters.emplace("element_count", output->element_count); + dispatch.kernel.compile_parameters.emplace("ggml.scale_add_f32.ne0", std::to_string(output->ne[0])); + dispatch.kernel.compile_parameters.emplace("ggml.scale_add_f32.ne1", std::to_string(output->ne[1])); + dispatch.kernel.compile_parameters.emplace("ggml.scale_add_f32.ne2", std::to_string(output->ne[2])); + dispatch.kernel.compile_parameters.emplace("ggml.scale_add_f32.scale", format_float_config(scale_params->scale)); + dispatch.kernel.compile_parameters.emplace("ggml.scale_add_f32.bias", format_float_config(scale_params->bias)); + common_set_binary_compile_parameters(dispatch.kernel, "ggml.scale_add_f32.binary_op", + "ggml.scale_add_f32.scaled_lhs", binary.kind, binary.producer_is_lhs); + add_strided_source_config(dispatch, "ggml.scale_add_f32.input_", *input, input_byte_count); + add_strided_source_config(dispatch, "ggml.scale_add_f32.residual_", *residual, residual_byte_count); + dispatch.bindings.push_back({ input->storage_root, input->storage_offset, input_byte_count }); + dispatch.bindings.push_back({ residual->storage_root, residual->storage_offset, residual_byte_count }); + dispatch.bindings.push_back({ output->id, 0, output->byte_count }); + + match.covered_nodes.push_back(context.root_index); + match.covered_nodes.push_back(binary_index); + match.dispatches.push_back(std::move(dispatch)); + return true; +} + static bool match_scale_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { const GraphNode * node = context.root_node; if (node == nullptr || node->op != GGML_OP_SCALE || node->inputs.size() != 1) { @@ -110,6 +234,14 @@ static bool match_scale_f32_dispatch(const DispatchMatchContext & context, Dispa } // namespace void register_scale_dispatch(DispatchRegistryBuilder & registry) { + registry.add({ + "common.scale_add_f32", + GGML_OP_SCALE, + DispatchMatchKind::Fused, + 1000, + DispatchSource::Common, + match_scale_add_f32_dispatch, + }); registry.add({ "common.scale_f32", GGML_OP_SCALE, diff --git a/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-attention-qkv.cpp b/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-attention-qkv.cpp index 5cbc4e0e53f9..a974bc1bee0b 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-attention-qkv.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-attention-qkv.cpp @@ -6,6 +6,7 @@ #include "kernel-corpus/kernel-corpus-catalog-verify.h" #include +#include #include #include #include @@ -259,10 +260,44 @@ struct AttentionQkvMatch { std::vector layout_nodes; KernelCatalogRef kernel = {}; AttentionQkvProjectionKind kind = AttentionQkvProjectionKind::Query; + bool use_f16_activation = false; bool matched() const { return root.matched() && kernel.id != kUncatalogedKernelId; } }; +static bool prefill_v_cache_fusion_disabled() { + const char * value = std::getenv("GGML_HRX_DISABLE_PREFILL_V_CACHE_FUSION"); + if (value == nullptr || value[0] == '\0') { + return false; + } + return std::strcmp(value, "0") != 0 && std::strcmp(value, "false") != 0 && std::strcmp(value, "FALSE") != 0 && + std::strcmp(value, "off") != 0 && std::strcmp(value, "OFF") != 0; +} + +static bool distinct_storage(const Graph & graph, const Value & lhs, const Value & rhs) { + return !graph.values().same_storage(lhs.id, rhs.id); +} + +static bool can_fuse_q6_f16_prefill_v_cache_cell(const DispatchMatchContext & context, + const AttentionMatMulMatch & root, + const AttentionSetRowsMatch & set_rows) { + if (prefill_v_cache_fusion_disabled() || root.weight->type != GGML_TYPE_Q6_K || root.input_size != 1536 || + root.output_size != 256 || root.token_count != 64 || set_rows.output_format != 16 || + set_rows.cache_row_count != 512 || root.output->kind != ValueKind::Transient || + root.weight->alias_source.value >= 0 || !distinct_storage(context.graph, *root.input, *root.weight) || + !distinct_storage(context.graph, *root.input, *set_rows.indices) || + !distinct_storage(context.graph, *root.input, *set_rows.output) || + !distinct_storage(context.graph, *root.weight, *set_rows.indices) || + !distinct_storage(context.graph, *root.weight, *set_rows.output) || + !distinct_storage(context.graph, *set_rows.indices, *set_rows.output)) { + return false; + } + + const size_t activation_bytes = static_cast(root.input_size * root.token_count) * sizeof(ggml_fp16_t); + return find_alternate_value(context.graph, context.plan, root.input->id, GGML_TYPE_F16, activation_bytes) != + nullptr; +} + static bool can_use_packed_lowtoken_attention_matmul(const AttentionMatMulMatch & root) { if (root.token_count < 1 || root.token_count > 5 || root.output_size % 64 != 0 || root.weight->alias_source.value >= 0) { @@ -483,15 +518,16 @@ static AttentionQkvMatch match_attention_qkv_projection(const DispatchMatchConte } if (consumer->op == GGML_OP_SET_ROWS) { - if (!is_supported_decode_token_count(root.token_count) && - !common_mul_mat_dense_float_format(root.weight_format)) { - return {}; - } match.set_rows = match_attention_set_rows(context.graph, consumer, after_projection); if (!match.set_rows.matched() || match.set_rows.token_count != root.token_count || match.set_rows.output_size != root.output_size) { return {}; } + const bool quantized_prefill = !is_supported_decode_token_count(root.token_count) && + !common_mul_mat_dense_float_format(root.weight_format); + if (quantized_prefill && !can_fuse_q6_f16_prefill_v_cache_cell(context, root, match.set_rows)) { + return {}; + } match.root = root; match.kernel = is_supported_decode_token_count(root.token_count) ? kAttentionVMatMulSetRowsVectorF32F32Kernel : kAttentionVMatMulSetRowsTiledF32F32Kernel; @@ -500,6 +536,7 @@ static AttentionQkvMatch match_attention_qkv_projection(const DispatchMatchConte match.kernel = kAttentionVMatMulSetRowsLegacyF32F32Kernel; } match.kind = AttentionQkvProjectionKind::Value; + match.use_f16_activation = quantized_prefill; return match; } @@ -595,6 +632,16 @@ static bool match_attention_qkv_dispatch(const DispatchMatchContext & context, D bool pack_q6 = match.root.weight_format == CommonMulMatWeightFormat::Q6KRow64; bool use_q8_activation = false; DispatchBinding activation = { match.root.input->id, 0, match.root.input->byte_count }; + if (match.use_f16_activation) { + const size_t activation_bytes = + static_cast(match.root.input_size * match.root.token_count) * sizeof(ggml_fp16_t); + const CommandPlanAlternateValue * alternate = + find_alternate_value(context.graph, context.plan, match.root.input->id, GGML_TYPE_F16, activation_bytes); + if (alternate == nullptr) { + return false; + } + activation = { alternate->alternate_value, 0, activation_bytes }; + } if (pack_q4 || pack_q6) { if (common_prepare_q8_1_x4_input(context, *match.root.input, match.root.input_size, match.root.token_count, dispatch_match, activation, @@ -632,6 +679,8 @@ static bool match_attention_qkv_dispatch(const DispatchMatchContext & context, D add_attention_qkv_compile_parameters(dispatch, match); if (use_q8_activation) { dispatch.kernel.compile_parameters.emplace("ggml.mul_mat.activation_format", std::to_string(GGML_TYPE_Q8_1)); + } else if (match.use_f16_activation) { + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat.activation_format", std::to_string(GGML_TYPE_F16)); } dispatch.bindings.push_back(activation); if (pack_q4 || pack_q6) { diff --git a/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-ssm-conv.cpp b/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-ssm-conv.cpp index 15ea9323451f..f731ba5da23e 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-ssm-conv.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-ssm-conv.cpp @@ -1,13 +1,14 @@ #include "dispatch-ssm-conv.h" +#include "../common/dispatch-binary-common.h" #include "../common/dispatch-mul-mat-common.h" - #include "ggml.h" #include "graph/graph-matcher.h" #include "kernel-corpus/kernel-corpus-catalog-verify.h" #include #include +#include #include #include #include @@ -32,6 +33,8 @@ static constexpr KernelCatalogRef kSsmConvRollbackKernel = static constexpr KernelCatalogRef kSsmConvGenericKernel = GGML_HRX_KERNEL_REF("loom_libs", "llm_ssm_conv_f32"); static constexpr KernelCatalogRef kSsmConvGenericBinaryKernel = GGML_HRX_KERNEL_REF("loom_libs", "llm_ssm_conv_binary_f32"); +static constexpr KernelCatalogRef kLfmDconv3Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "llm_ssm_conv_lfm_dconv3_state_f32"); static const Value * graph_value(const Graph & graph, ValueId id) { return graph.values().find(id); @@ -153,6 +156,151 @@ struct SsmConvCoreMatch { bool has_binary_fusion() const { return binary != nullptr; } }; +struct LfmDconv3Match { + const GraphNode * concat = nullptr; + const GraphNode * transpose = nullptr; + const GraphNode * ssm = nullptr; + const GraphNode * tail_view = nullptr; + const GraphNode * cache_view = nullptr; + const GraphNode * cache_copy = nullptr; + const Value * state = nullptr; + const Value * x = nullptr; + const Value * filter = nullptr; + const Value * output = nullptr; + const Value * cache = nullptr; + int64_t token_count = 0; + + bool matched() const { + return concat != nullptr && transpose != nullptr && ssm != nullptr && tail_view != nullptr && + cache_view != nullptr && cache_copy != nullptr && state != nullptr && x != nullptr && + filter != nullptr && output != nullptr && cache != nullptr; + } +}; + +static bool lfm_dconv3_disabled() { + return std::getenv("GGML_HRX_DISABLE_LFM_DCONV3_FUSION") != nullptr; +} + +static LfmDconv3Match match_lfm_dconv3(const Graph & graph, const GraphNode * node) { + LfmDconv3Match match; + if (lfm_dconv3_disabled() || node == nullptr || node->op != GGML_OP_CONCAT || node->inputs.size() != 2 || + !graph.has_index()) { + return match; + } + + const Value * state = graph_value(graph, node->inputs[0]); + const Value * x_transposed = graph_value(graph, node->inputs[1]); + const Value * window = graph_value(graph, node->output); + if (!is_f32(state) || !is_f32(x_transposed) || !is_f32(window)) { + return {}; + } + const GraphNode * transpose = graph.index().producer(x_transposed->id); + if (transpose == nullptr || transpose->op != GGML_OP_TRANSPOSE || transpose->inputs.size() != 1) { + return {}; + } + const Value * x_layout = graph_value(graph, transpose->inputs[0]); + const Value * x = root_alias_value(graph, x_layout); + if (!is_f32(x_layout) || !x_layout->contiguous || !is_f32(x) || !x->contiguous) { + return {}; + } + + const int64_t token_count = x_layout->ne[1]; + if ((token_count != 1 && token_count != 64) || !is_shape(*state, 2, 2048, 1, 1) || + !is_shape(*x_layout, 2048, token_count, 1, 1) || !is_shape(*x_transposed, token_count, 2048, 1, 1) || + !is_shape(*window, token_count + 2, 2048, 1, 1) || state->nb[0] != sizeof(float) || + state->nb[1] != 2 * sizeof(float) || x_layout->nb[0] != sizeof(float) || + x_layout->nb[1] != 2048 * sizeof(float) || x_transposed->nb[0] != 2048 * sizeof(float) || + x_transposed->nb[1] != sizeof(float) || window->nb[0] != sizeof(float) || + window->nb[1] != static_cast(token_count + 2) * sizeof(float)) { + return {}; + } + + const GraphNode * ssm = nullptr; + const GraphNode * tail_view = nullptr; + const auto & consumers = graph.index().consumers(window->id); + if (consumers.size() != 2) { + return {}; + } + for (const GraphNode * consumer : consumers) { + if (consumer != nullptr && consumer->op == GGML_OP_SSM_CONV && ssm == nullptr) { + ssm = consumer; + } else if (consumer != nullptr && consumer->op == GGML_OP_VIEW && tail_view == nullptr) { + tail_view = consumer; + } else { + return {}; + } + } + if (ssm == nullptr || tail_view == nullptr || ssm->inputs.size() != 2 || ssm->inputs[0] != window->id || + tail_view->inputs.size() != 1) { + return {}; + } + const Value * filter = graph_value(graph, ssm->inputs[1]); + const Value * output = graph_value(graph, ssm->output); + const Value * tail = graph_value(graph, tail_view->output); + if (!is_f32(filter) || !is_f32(output) || !is_f32(tail) || !filter->contiguous || !output->contiguous || + !is_shape(*filter, 3, 2048, 1, 1) || !is_shape(*output, 2048, token_count, 1, 1) || + !is_shape(*tail, 2, 2048, 1, 1) || tail->alias_source != window->id || + tail->storage_offset < window->storage_offset || + tail->storage_offset - window->storage_offset != static_cast(token_count) * sizeof(float)) { + return {}; + } + const auto & tail_consumers = graph.index().consumers(tail->id); + if (tail_consumers.size() != 1 || tail_consumers.front() == nullptr || tail_consumers.front()->op != GGML_OP_CPY || + tail_consumers.front()->inputs.size() != 2 || tail_consumers.front()->inputs[0] != tail->id) { + return {}; + } + const GraphNode * cache_copy = tail_consumers.front(); + const Value * cache_target = graph_value(graph, cache_copy->inputs[1]); + const Value * cache = graph_value(graph, cache_copy->output); + const GraphNode * cache_view = cache_target != nullptr ? graph.index().producer(cache_target->id) : nullptr; + if (!is_f32(cache_target) || !is_f32(cache) || cache_view == nullptr || cache_view->op != GGML_OP_VIEW || + cache_view->inputs.size() != 1 || !same_full_value_range(*cache_target, *cache) || + cache->byte_count != static_cast(4096) * sizeof(float)) { + return {}; + } + if (!distinct_storage(*state, *x) || !distinct_storage(*state, *filter) || !distinct_storage(*state, *output) || + !distinct_storage(*state, *cache) || !distinct_storage(*x, *filter) || !distinct_storage(*x, *output) || + !distinct_storage(*x, *cache) || !distinct_storage(*filter, *output) || !distinct_storage(*filter, *cache) || + !distinct_storage(*output, *cache)) { + return {}; + } + + match.concat = node; + match.transpose = transpose; + match.ssm = ssm; + match.tail_view = tail_view; + match.cache_view = cache_view; + match.cache_copy = cache_copy; + match.state = state; + match.x = x; + match.filter = filter; + match.output = output; + match.cache = cache; + match.token_count = token_count; + return match; +} + +static bool match_lfm_dconv3_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + const LfmDconv3Match match = match_lfm_dconv3(context.graph, context.root_node); + if (!match.matched() || !append_covered_node(context, match.concat, dispatch_match) || + !append_covered_node(context, match.ssm, dispatch_match) || + !append_covered_node(context, match.tail_view, dispatch_match) || + !append_covered_node(context, match.cache_view, dispatch_match) || + !append_covered_node(context, match.cache_copy, dispatch_match)) { + return false; + } + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kLfmDconv3Kernel); + set_compile_parameter(dispatch.kernel, "llm.ssm_conv.lfm_dconv3.n_t", match.token_count); + dispatch.bindings.push_back({match.state->id, 0, match.state->byte_count}); + dispatch.bindings.push_back({match.x->id, 0, match.x->byte_count}); + dispatch.bindings.push_back({match.filter->id, 0, match.filter->byte_count}); + dispatch.bindings.push_back({match.output->id, 0, match.output->byte_count}); + dispatch.bindings.push_back({match.cache->id, 0, match.cache->byte_count}); + dispatch_match.dispatches.push_back(std::move(dispatch)); + return true; +} + static SsmConvPrefillMatch match_ssm_conv_prefill(const Graph & graph, const GraphNode * node) { SsmConvPrefillMatch match; if (node == nullptr || node->op != GGML_OP_CONCAT || node->inputs.size() != 2 || !graph.has_index()) { @@ -391,18 +539,13 @@ static bool try_match_binary_fusion(const Graph & graph, SsmConvCoreMatch & matc } const GraphNode * binary = consumers.front(); - const BinaryParams * params = op_params_as(binary->params); - if (params == nullptr || params->op != BinaryKind::Mul) { + const CommonBinaryMatch binary_match = + common_match_binary_consumer(binary, match.conv_output->id, kCommonArithmeticBinaryKinds); + if (!binary_match.matched()) { return false; } - const bool conv_is_lhs = binary->inputs[0] == match.conv_output->id; - const bool conv_is_rhs = binary->inputs[1] == match.conv_output->id; - if (conv_is_lhs == conv_is_rhs) { - return false; - } - - const Value * operand = graph_value(graph, conv_is_lhs ? binary->inputs[1] : binary->inputs[0]); + const Value * operand = graph_value(graph, binary_match.operand); const Value * output = graph_value(graph, binary->output); if (!is_f32(operand) || output == nullptr || output->type != GGML_TYPE_F32 || !packed_f32_layout(*operand) || !packed_f32_layout(*output) || !same_shape(*operand, *match.conv_output) || @@ -414,8 +557,8 @@ static bool try_match_binary_fusion(const Graph & graph, SsmConvCoreMatch & matc match.binary = binary; match.binary_operand = operand; match.output = output; - match.binary_op = params->op; - match.binary_lhs = conv_is_lhs; + match.binary_op = binary_match.kind; + match.binary_lhs = binary_match.producer_is_lhs; return true; } @@ -472,9 +615,9 @@ static bool match_ssm_conv_prefill_dispatch(const DispatchMatchContext & context return true; } - const bool optimized_prefill = - match.token_count == 512 && match.sequence_count == 1 && match.cache_updates.size() == 1 && - match.hidden_size >= 8192 && match.hidden_size <= 10240; + const bool optimized_prefill = match.token_count == 512 && match.sequence_count == 1 && + match.cache_updates.size() == 1 && match.hidden_size >= 8192 && + match.hidden_size <= 10240; if (!optimized_prefill) { Dispatch ssm_dispatch; const bool decode = match.token_count == 1 && match.sequence_count == 1 && match.cache_updates.size() == 1; @@ -573,9 +716,8 @@ static bool match_ssm_conv_generic_dispatch(const DispatchMatchContext & context dispatch.bindings.push_back({ match.window->id, 0, match.window->byte_count }); dispatch.bindings.push_back({ match.filter->id, 0, match.filter->byte_count }); if (match.has_binary_fusion()) { - dispatch.kernel.compile_parameters.emplace("llm.ssm_conv.generic.binary_op", - std::to_string(binary_kind_config_value(match.binary_op))); - dispatch.kernel.compile_parameters.emplace("llm.ssm_conv.generic.binary_lhs", match.binary_lhs ? "1" : "0"); + common_set_binary_compile_parameters(dispatch.kernel, "llm.ssm_conv.generic.binary_op", + "llm.ssm_conv.generic.binary_lhs", match.binary_op, match.binary_lhs); dispatch.bindings.push_back({ match.binary_operand->id, 0, match.binary_operand->byte_count }); } dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); @@ -584,12 +726,12 @@ static bool match_ssm_conv_generic_dispatch(const DispatchMatchContext & context } static bool match_mul_mat_conv4_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { - const CommonMulMatMatch mat = common_match_mul_mat_any_format( - context.graph, context.root_node, kMulMatConv4Kernel, false); + const CommonMulMatMatch mat = + common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatConv4Kernel, false); if (!mat.matched() || !context.graph.has_index() || mat.token_count != 512 || (mat.weight->type != GGML_TYPE_Q4_K && mat.weight->type != GGML_TYPE_Q6_K) || - mat.weight->alias_source.value >= 0 || mat.input_size % 256 != 0 || - mat.output_size % 64 != 0 || mat.output_size / 64 < 32 || + mat.weight->alias_source.value >= 0 || mat.input_size % 256 != 0 || mat.output_size % 64 != 0 || + mat.output_size / 64 < 32 || !common_mul_mat_uses_k16_major_f16(mat.weight_format, mat.input_size, mat.output_size, mat.token_count)) { return false; } @@ -614,8 +756,8 @@ static bool match_mul_mat_conv4_dispatch(const DispatchMatchContext & context, D value = graph_value(context.graph, consumer->output); } const SsmConvPrefillMatch conv = match_ssm_conv_prefill(context.graph, concat); - if (!conv.matched() || conv.x->id != mat.output->id || conv.sequence_count != 1 || - conv.token_count != 512 || conv.cache_updates.size() != 1) { + if (!conv.matched() || conv.x->id != mat.output->id || conv.sequence_count != 1 || conv.token_count != 512 || + conv.cache_updates.size() != 1) { return false; } const Value * window = graph_value(context.graph, conv.concat->output); @@ -627,7 +769,8 @@ static bool match_mul_mat_conv4_dispatch(const DispatchMatchContext & context, D for (const Value * input : {conv.state, conv.filter}) { const GraphNode * producer = context.graph.index().producer(input->id); size_t index = 0; - if (producer != nullptr && (!context.graph.index().node_index(producer, index) || !context.covered_nodes[index])) { + if (producer != nullptr && + (!context.graph.index().node_index(producer, index) || !context.covered_nodes[index])) { if (input == conv.filter) { return false; } @@ -654,8 +797,8 @@ static bool match_mul_mat_conv4_dispatch(const DispatchMatchContext & context, D } DispatchBinding activation; - if (!common_prepare_k16_major_f16_input(context, *mat.input, mat.input_size, mat.token_count, - dispatch_match, activation)) { + if (!common_prepare_k16_major_f16_input(context, *mat.input, mat.input_size, mat.token_count, dispatch_match, + activation)) { return false; } Dispatch dispatch; @@ -713,6 +856,14 @@ void register_llm_ssm_conv_dispatch(DispatchRegistryBuilder & registry) { DispatchSource::Llm, match_ssm_conv_prefill_dispatch, }); + registry.add({ + "llm.ssm_conv.lfm_dconv3", + GGML_OP_CONCAT, + DispatchMatchKind::Fused, + 201, + DispatchSource::Llm, + match_lfm_dconv3_dispatch, + }); registry.add({ "llm.ssm_conv.generic_f32", GGML_OP_SSM_CONV, diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/hrx/gather_add_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/hrx/gather_add_f32.loom index f645dedd9d6d..c8e46116e5d6 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/hrx/gather_add_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/hrx/gather_add_f32.loom @@ -9,6 +9,8 @@ // GGML GET_ROWS operations followed by ADD. amdgpu.target @ggml_gather_add_gfx11_wave64 {subgroup_size = 64} +config.decl @ggml.gather_add_f32.binary_op : %value: index where [range(%value, 0, 3)] + kernel.def target(@ggml_gather_add_gfx11_wave64) export("ggml_gather_add_f32") @ggml_gather_add_f32(%source_token_count: index, %output_token_count: index, %hidden_size: index) { %one = index.constant 1 : index %twofiftysix = index.constant 256 : index @@ -20,6 +22,7 @@ kernel.def target(@ggml_gather_add_gfx11_wave64) export("ggml_gather_add_f32") @ %source_count = index.assume %source_token_count [range(%source_token_count, 1, 2048)] : index %output_count = index.assume %output_token_count [range(%output_token_count, 1, 2048)] : index %width = index.assume %hidden_size [range(%hidden_size, 1, 65536)] : index + %binary_op = config.get @ggml.gather_add_f32.binary_op : index %column_workgroup = kernel.workgroup.id : index %output_row0 = kernel.workgroup.id : index %workitem = kernel.workitem.id : index @@ -43,7 +46,19 @@ kernel.def target(@ggml_gather_add_gfx11_wave64) export("ggml_gather_add_f32") @ %attention_value = view.load %attention_view[%source_row, %safe_column] : view<[%source_count]x[%width]xf32> -> f32 %residual_value = view.load %residual_view[%source_row, %safe_column] : view<[%source_count]x[%width]xf32> -> f32 %sum = scalar.addf %attention_value, %residual_value : f32 - view.store %sum, %output_view[%output_row, %safe_column] : f32, view<[%launch_output_count]x[%width]xf32> + %difference = scalar.subf %attention_value, %residual_value : f32 + %product = scalar.mulf %attention_value, %residual_value : f32 + %quotient = scalar.divf %attention_value, %residual_value : f32 + %op0 = index.constant 0 : index + %op1 = index.constant 1 : index + %op2 = index.constant 2 : index + %is_add = index.cmp eq, %binary_op, %op0 : index + %is_sub = index.cmp eq, %binary_op, %op1 : index + %is_mul = index.cmp eq, %binary_op, %op2 : index + %add_or_div = scf.select %is_add, %sum, %quotient : f32 + %sub_or_other = scf.select %is_sub, %difference, %add_or_div : f32 + %result = scf.select %is_mul, %product, %sub_or_other : f32 + view.store %result, %output_view[%output_row, %safe_column] : f32, view<[%launch_output_count]x[%width]xf32> } kernel.return } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/manifest.json b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/manifest.json index 9bfdc8a97301..1f6348c02a23 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/manifest.json +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/manifest.json @@ -14,6 +14,9 @@ { "path": "ops/elementwise/scale_f32.loom" }, + { + "path": "ops/elementwise/scale_add_f32.loom" + }, { "path": "ops/elementwise/scale_bias_f32.loom" }, @@ -47,6 +50,12 @@ { "path": "ops/normalization/rmsnorm_binary_f32.loom" }, + { + "path": "ops/normalization/rmsnorm_binary_strided_f32.loom" + }, + { + "path": "ops/normalization/rmsnorm_mul_add_f32.loom" + }, { "path": "ops/normalization/rmsnorm_binary_f32_f16.loom" }, @@ -165,11 +174,41 @@ "path": "ops/matmul_id/vector_weighted.loom" }, { - "path": "../qwen_moe/qwen3_moe/attention_prepare_quantized.loom" + "path": "motifs/normalization/rmsnorm_q8_1_x4.loom" + }, + { + "path": "motifs/matmul/q4_q8_1_x4_decode.loom" + }, + { + "path": "motifs/matmul/q4_q8_1_x4_dense.loom" + }, + { + "path": "motifs/matmul/q4_q6_q8_1_x4_qkv.loom" }, { "path": "motifs/matmul/q6_f32_decode.loom" }, + { + "path": "motifs/matmul/q6_q8_1_x4_decode.loom" + }, + { + "path": "motifs/routing/router_projection_storage.loom" + }, + { + "path": "motifs/routing/router_topk_row.loom" + }, + { + "path": "motifs/routing/router_projection_fused.loom" + }, + { + "path": "motifs/routing/weighted_reduce_f16_f32.loom" + }, + { + "path": "motifs/routing/weighted_reduce_next_rmsnorm_f32.loom" + }, + { + "path": "motifs/attention/rmsnorm_neox_cache_f16.loom" + }, { "path": "motifs/matmul_id/pair_input_f32_dequant_accumulate_f32.loom" }, @@ -242,6 +281,12 @@ { "path": "ops/rows/get_rows_f32.loom" }, + { + "path": "ops/rows/get_rows_scale_f32.loom" + }, + { + "path": "ops/rows/gather_add_rmsnorm_binary_f32.loom" + }, { "path": "ops/rows/get_rows_rmsnorm_binary_q8_1_x4_f16.loom" }, @@ -332,6 +377,9 @@ { "path": "ops/ssm_conv/generic_f32.loom" }, + { + "path": "ops/ssm_conv/lfm_dconv3_state_f32.loom" + }, { "path": "ops/ssm_conv/generic_binary_f32.loom" }, @@ -353,6 +401,12 @@ { "path": "ops/rope/rope_f32.loom" }, + { + "path": "ops/rope/rope_concat_f32.loom" + }, + { + "path": "ops/rope/rope_scale_f32.loom" + }, { "path": "ops/rope/rope_set_rows_f32.loom" }, @@ -449,6 +503,12 @@ { "path": "motifs/quantize_q8_1_x4.loom" }, + { + "path": "motifs/normalization/rmsnorm_q8_1_x4.loom" + }, + { + "path": "motifs/quantize_q8_1_x4_group.loom" + }, { "path": "motifs/publish_f32.loom" }, @@ -951,6 +1011,49 @@ "motifs/elementwise/scale_apply.loom" ] }, + { + "name": "ggml_scale_add_f32", + "family": "loom_libs", + "symbol": "ggml_scale_add_f32", + "source": "ops/elementwise/scale_add_f32.loom", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "bindings": [ + "input", + "residual", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/elementwise/scale_add_f32.loom" + ], + "library_sources": [ + "motifs/elementwise/scale_apply.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + "compile_dependencies": [ + "motifs/elementwise/scale_apply.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, { "name": "ggml_binary_f32", "family": "loom_libs", @@ -1239,6 +1342,96 @@ "motifs/elementwise/binary_apply.loom" ] }, + { + "name": "ggml_rmsnorm_binary_strided_f32", + "family": "loom_libs", + "symbol": "ggml_rmsnorm_binary_strided_f32", + "source": "ops/normalization/rmsnorm_binary_strided_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "rhs", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/normalization/rmsnorm_binary_strided_f32.loom" + ], + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + "compile_dependencies": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + { + "name": "ggml_rmsnorm_mul_add_f32", + "family": "loom_libs", + "symbol": "ggml_rmsnorm_mul_add_f32", + "source": "ops/normalization/rmsnorm_mul_add_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "residual", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/normalization/rmsnorm_mul_add_f32.loom" + ], + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + "compile_dependencies": [ + "motifs/normalization/rmsnorm.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, { "name": "ggml_rmsnorm_binary_f32_k16", "family": "loom_libs", @@ -2030,9 +2223,9 @@ ] }, { - "name": "ggml_mul_mat_vector_bias_f32_f32", + "name": "ggml_mul_mat_vector_bias_residual_f32_f32", "family": "loom_libs", - "symbol": "ggml_mul_mat_vector_bias_f32_f32", + "symbol": "ggml_mul_mat_vector_bias_residual_f32_f32", "source": "ops/matmul/vector_postops_f32_f32.loom", "workload_parameters": [ { @@ -2050,13 +2243,15 @@ "input", "weight", "bias", - "output", + "residual_input", + "residual_output", "next_output" ], "binding_access": [ "read", "read", "read", + "read", "read_write", "write" ], @@ -2097,10 +2292,10 @@ ] }, { - "name": "ggml_mul_mat_vector_residual_f32_f32", + "name": "llm_attention_q_matmul_rope_tiled_f32_f32", "family": "loom_libs", - "symbol": "ggml_mul_mat_vector_residual_f32_f32", - "source": "ops/matmul/vector_postops_f32_f32.loom", + "symbol": "llm_attention_q_matmul_rope_tiled_f32_f32", + "source": "ops/llm/llm_attention_q_matmul_rope_tiled_f32_f32.loom", "workload_parameters": [ { "name": "token_count", @@ -2116,28 +2311,29 @@ "bindings": [ "input", "weight", - "residual_input", - "residual_output", - "next_output" + "positions", + "theta", + "freq_factors", + "output" ], "binding_access": [ "read", "read", "read", - "read_write", - "write" + "read", + "read", + "read_write" ], "target_selector": "", "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/matmul/vector_postops_f32_f32.loom" + "ops/llm/llm_attention_q_matmul_rope_tiled_f32_f32.loom" ], "library_sources": [ - "motifs/matmul/vector_activation.loom", - "motifs/matmul/vector_decode_pair.loom", - "motifs/matmul/vector_publish.loom", - "motifs/matmul/vector_postops.loom", + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/llm_attention_qkv_matmul_postops.loom", + "motifs/rope/rope_f32.loom", "motifs/dequant.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", @@ -2149,10 +2345,9 @@ ] }, "compile_dependencies": [ - "motifs/matmul/vector_activation.loom", - "motifs/matmul/vector_decode_pair.loom", - "motifs/matmul/vector_publish.loom", - "motifs/matmul/vector_postops.loom", + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/llm_attention_qkv_matmul_postops.loom", + "motifs/rope/rope_f32.loom", "motifs/dequant.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", @@ -2164,10 +2359,10 @@ ] }, { - "name": "ggml_mul_mat_vector_bias_residual_f32_f32", + "name": "llm_attention_k_matmul_rope_set_rows_tiled_f32_f32", "family": "loom_libs", - "symbol": "ggml_mul_mat_vector_bias_residual_f32_f32", - "source": "ops/matmul/vector_postops_f32_f32.loom", + "symbol": "llm_attention_k_matmul_rope_set_rows_tiled_f32_f32", + "source": "ops/llm/llm_attention_k_matmul_rope_set_rows_tiled_f32_f32.loom", "workload_parameters": [ { "name": "token_count", @@ -2183,30 +2378,31 @@ "bindings": [ "input", "weight", - "bias", - "residual_input", - "residual_output", - "next_output" + "positions", + "indices", + "theta", + "freq_factors", + "cache" ], "binding_access": [ "read", "read", "read", "read", - "read_write", - "write" + "read", + "read", + "read_write" ], "target_selector": "", "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/matmul/vector_postops_f32_f32.loom" + "ops/llm/llm_attention_k_matmul_rope_set_rows_tiled_f32_f32.loom" ], "library_sources": [ - "motifs/matmul/vector_activation.loom", - "motifs/matmul/vector_decode_pair.loom", - "motifs/matmul/vector_publish.loom", - "motifs/matmul/vector_postops.loom", + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/llm_attention_qkv_matmul_postops.loom", + "motifs/rope/rope_f32.loom", "motifs/dequant.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", @@ -2218,10 +2414,9 @@ ] }, "compile_dependencies": [ - "motifs/matmul/vector_activation.loom", - "motifs/matmul/vector_decode_pair.loom", - "motifs/matmul/vector_publish.loom", - "motifs/matmul/vector_postops.loom", + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/llm_attention_qkv_matmul_postops.loom", + "motifs/rope/rope_f32.loom", "motifs/dequant.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", @@ -2233,10 +2428,10 @@ ] }, { - "name": "llm_attention_q_matmul_rope_tiled_f32_f32", + "name": "llm_attention_v_matmul_set_rows_tiled_f32_f32", "family": "loom_libs", - "symbol": "llm_attention_q_matmul_rope_tiled_f32_f32", - "source": "ops/llm/llm_attention_q_matmul_rope_tiled_f32_f32.loom", + "symbol": "llm_attention_v_matmul_set_rows_tiled_f32_f32", + "source": "ops/llm/llm_attention_v_matmul_set_rows_tiled_f32_f32.loom", "workload_parameters": [ { "name": "token_count", @@ -2252,14 +2447,10 @@ "bindings": [ "input", "weight", - "positions", - "theta", - "freq_factors", - "output" + "indices", + "cache" ], "binding_access": [ - "read", - "read", "read", "read", "read", @@ -2269,12 +2460,11 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/llm/llm_attention_q_matmul_rope_tiled_f32_f32.loom" + "ops/llm/llm_attention_v_matmul_set_rows_tiled_f32_f32.loom" ], "library_sources": [ "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", "motifs/llm_attention_qkv_matmul_postops.loom", - "motifs/rope/rope_f32.loom", "motifs/dequant.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", @@ -2288,7 +2478,6 @@ "compile_dependencies": [ "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", "motifs/llm_attention_qkv_matmul_postops.loom", - "motifs/rope/rope_f32.loom", "motifs/dequant.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", @@ -2300,140 +2489,10 @@ ] }, { - "name": "llm_attention_k_matmul_rope_set_rows_tiled_f32_f32", + "name": "llm_attention_q_matmul_rope_vector_f32_f32", "family": "loom_libs", - "symbol": "llm_attention_k_matmul_rope_set_rows_tiled_f32_f32", - "source": "ops/llm/llm_attention_k_matmul_rope_set_rows_tiled_f32_f32.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "input", - "weight", - "positions", - "indices", - "theta", - "freq_factors", - "cache" - ], - "binding_access": [ - "read", - "read", - "read", - "read", - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "ops/llm/llm_attention_k_matmul_rope_set_rows_tiled_f32_f32.loom" - ], - "library_sources": [ - "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", - "motifs/llm_attention_qkv_matmul_postops.loom", - "motifs/rope/rope_f32.loom", - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/bf16_f16.loom", - "motifs/f32_f16.loom" - ] - }, - "compile_dependencies": [ - "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", - "motifs/llm_attention_qkv_matmul_postops.loom", - "motifs/rope/rope_f32.loom", - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/bf16_f16.loom", - "motifs/f32_f16.loom" - ] - }, - { - "name": "llm_attention_v_matmul_set_rows_tiled_f32_f32", - "family": "loom_libs", - "symbol": "llm_attention_v_matmul_set_rows_tiled_f32_f32", - "source": "ops/llm/llm_attention_v_matmul_set_rows_tiled_f32_f32.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "input", - "weight", - "indices", - "cache" - ], - "binding_access": [ - "read", - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "ops/llm/llm_attention_v_matmul_set_rows_tiled_f32_f32.loom" - ], - "library_sources": [ - "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", - "motifs/llm_attention_qkv_matmul_postops.loom", - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/bf16_f16.loom", - "motifs/f32_f16.loom" - ] - }, - "compile_dependencies": [ - "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", - "motifs/llm_attention_qkv_matmul_postops.loom", - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/bf16_f16.loom", - "motifs/f32_f16.loom" - ] - }, - { - "name": "llm_attention_q_matmul_rope_vector_f32_f32", - "family": "loom_libs", - "symbol": "llm_attention_q_matmul_rope_vector_f32_f32", - "source": "ops/llm/llm_attention_q_matmul_rope_vector_f32_f32.loom", + "symbol": "llm_attention_q_matmul_rope_vector_f32_f32", + "source": "ops/llm/llm_attention_q_matmul_rope_vector_f32_f32.loom", "workload_parameters": [ { "name": "token_count", @@ -3748,14 +3807,14 @@ "motifs/matmul_id/vector_weighted.loom", "motifs/matmul_id/vector_q4k_q8.loom", "motifs/matmul/q6_f32_decode.loom", - "../qwen_moe/qwen3_moe/attention_prepare_quantized.loom" + "motifs/normalization/rmsnorm_q8_1_x4.loom" ] }, "compile_dependencies": [ "motifs/matmul_id/vector_weighted.loom", "motifs/matmul_id/vector_q4k_q8.loom", "motifs/matmul/q6_f32_decode.loom", - "../qwen_moe/qwen3_moe/attention_prepare_quantized.loom" + "motifs/normalization/rmsnorm_q8_1_x4.loom" ] }, { @@ -3845,14 +3904,14 @@ "motifs/matmul_id/vector_weighted.loom", "motifs/matmul_id/vector_q4k_q8.loom", "motifs/matmul/q6_f32_decode.loom", - "../qwen_moe/qwen3_moe/attention_prepare_quantized.loom" + "motifs/normalization/rmsnorm_q8_1_x4.loom" ] }, "compile_dependencies": [ "motifs/matmul_id/vector_weighted.loom", "motifs/matmul_id/vector_q4k_q8.loom", "motifs/matmul/q6_f32_decode.loom", - "../qwen_moe/qwen3_moe/attention_prepare_quantized.loom" + "motifs/normalization/rmsnorm_q8_1_x4.loom" ] }, { @@ -6143,6 +6202,140 @@ "motifs/f32_f16.loom" ] }, + { + "name": "ggml_get_rows_scale_f32", + "family": "loom_libs", + "symbol": "ggml_get_rows_scale_f32", + "source": "ops/rows/get_rows_scale_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ], + "bindings": [ + "token_ids", + "weight", + "raw_output", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/rows/get_rows_scale_f32.loom" + ], + "library_sources": [ + "motifs/rows/get_rows_decode.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/rows/get_rows_decode.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_gather_add_rmsnorm_binary_f32", + "family": "loom_libs", + "symbol": "ggml_gather_add_rmsnorm_binary_f32", + "source": "ops/rows/gather_add_rmsnorm_binary_f32.loom", + "workload_parameters": [ + { + "name": "source_token_count", + "type": "index" + }, + { + "name": "output_token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "source_token_count", + "type": "index" + }, + { + "name": "output_token_count", + "type": "index" + } + ], + "bindings": [ + "attention", + "residual", + "output_ids", + "raw_output", + "weight", + "normalized_output" + ], + "binding_access": [ + "read", + "read", + "read", + "write", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/rows/gather_add_rmsnorm_binary_f32.loom" + ], + "library_sources": [ + "motifs/elementwise/binary_apply.loom" + ] + }, + "compile_dependencies": [ + "motifs/elementwise/binary_apply.loom" + ] + }, { "name": "ggml_get_rows_rmsnorm_binary_q8_1_x4_f16", "family": "loom_libs", @@ -6353,17 +6546,154 @@ ] }, { - "name": "ggml_set_rows", + "name": "ggml_rope_scale_f32", "family": "loom_libs", - "symbol": "ggml_set_rows", - "source": "ops/rows/set_rows.loom", + "symbol": "ggml_rope_scale_f32", + "source": "ops/rope/rope_scale_f32.loom", "workload_parameters": [ { "name": "token_count", "type": "index" - }, + } + ], + "launch_parameters": [ { - "name": "cache_row_count", + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "positions", + "input", + "theta", + "freq_factors", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/rope/rope_scale_f32.loom" + ], + "library_sources": [ + "motifs/rope/rope_f32.loom" + ] + }, + "compile_dependencies": [ + "motifs/rope/rope_f32.loom" + ] + }, + { + "name": "ggml_rope_token_scale_f32", + "family": "loom_libs", + "symbol": "ggml_rope_token_scale_f32", + "source": "ops/rope/rope_scale_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "positions", + "input", + "theta", + "freq_factors", + "scales", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/rope/rope_scale_f32.loom" + ], + "library_sources": [ + "motifs/rope/rope_f32.loom" + ] + }, + "compile_dependencies": [ + "motifs/rope/rope_f32.loom" + ] + }, + { + "name": "ggml_rope_concat_f32", + "family": "loom_libs", + "symbol": "ggml_rope_concat_f32", + "source": "ops/rope/rope_concat_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "positions", + "input", + "theta", + "freq_factors", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/rope/rope_concat_f32.loom" + ], + "library_sources": [ + "motifs/rope/rope_f32.loom" + ] + }, + "compile_dependencies": [ + "motifs/rope/rope_f32.loom" + ] + }, + { + "name": "ggml_set_rows", + "family": "loom_libs", + "symbol": "ggml_set_rows", + "source": "ops/rows/set_rows.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "cache_row_count", "type": "index" }, { @@ -6501,6 +6831,37 @@ "motifs/elementwise/binary_apply.loom" ] }, + { + "name": "llm_ssm_conv_lfm_dconv3_state_f32", + "family": "loom_libs", + "symbol": "llm_ssm_conv_lfm_dconv3_state_f32", + "source": "ops/ssm_conv/lfm_dconv3_state_f32.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "state", + "x", + "filter", + "output", + "cache" + ], + "binding_access": [ + "read", + "read", + "read", + "write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/ssm_conv/lfm_dconv3_state_f32.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, { "name": "llm_ssm_conv_binary_f32", "family": "loom_libs", @@ -7429,9 +7790,9 @@ ] }, { - "name": "ggml_mul_mat_skinny_input_f32_bias_publish_f32", + "name": "ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", "family": "loom_libs", - "symbol": "ggml_mul_mat_skinny_input_f32_bias_publish_f32", + "symbol": "ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", "source": "ops/matmul/skinny_input_f32_publish_f32.loom", "workload_parameters": [ { @@ -7449,12 +7810,14 @@ "input", "weight", "bias", - "output" + "residual_input", + "residual_output" ], "binding_access": [ "read", "read", "read", + "read", "read_write" ], "target_selector": "", @@ -7494,10 +7857,10 @@ ] }, { - "name": "ggml_mul_mat_skinny_input_f32_residual_publish_f32", + "name": "ggml_mul_mat_tiled_input_f32_publish_f32", "family": "loom_libs", - "symbol": "ggml_mul_mat_skinny_input_f32_residual_publish_f32", - "source": "ops/matmul/skinny_input_f32_publish_f32.loom", + "symbol": "ggml_mul_mat_tiled_input_f32_publish_f32", + "source": "ops/matmul/tiled_input_f32_publish_f32.loom", "workload_parameters": [ { "name": "token_count", @@ -7513,11 +7876,9 @@ "bindings": [ "input", "weight", - "residual_input", - "residual_output" + "output" ], "binding_access": [ - "read", "read", "read", "read_write" @@ -7526,12 +7887,12 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/matmul/skinny_input_f32_publish_f32.loom" + "ops/matmul/tiled_input_f32_publish_f32.loom" ], "library_sources": [ - "motifs/matmul/skinny_input_f32_dequant_accumulate_f32.loom", - "motifs/matmul/skinny_postprocess_f32.loom", - "motifs/matmul/skinny_publish_f32.loom", + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", "motifs/dequant.loom", "motifs/elementwise/unary_apply.loom", "motifs/q4_k_f16.loom", @@ -7544,9 +7905,9 @@ ] }, "compile_dependencies": [ - "motifs/matmul/skinny_input_f32_dequant_accumulate_f32.loom", - "motifs/matmul/skinny_postprocess_f32.loom", - "motifs/matmul/skinny_publish_f32.loom", + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", "motifs/dequant.loom", "motifs/elementwise/unary_apply.loom", "motifs/q4_k_f16.loom", @@ -7559,10 +7920,10 @@ ] }, { - "name": "ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", + "name": "ggml_mul_mat_tiled_input_f32_publish_f32_f16_alternate", "family": "loom_libs", - "symbol": "ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", - "source": "ops/matmul/skinny_input_f32_publish_f32.loom", + "symbol": "ggml_mul_mat_tiled_input_f32_publish_f32_f16_alternate", + "source": "ops/matmul/tiled_input_f32_publish_f32.loom", "workload_parameters": [ { "name": "token_count", @@ -7578,27 +7939,25 @@ "bindings": [ "input", "weight", - "bias", - "residual_input", - "residual_output" + "output", + "f16_output" ], "binding_access": [ "read", "read", - "read", - "read", - "read_write" + "read_write", + "write" ], "target_selector": "", "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/matmul/skinny_input_f32_publish_f32.loom" + "ops/matmul/tiled_input_f32_publish_f32.loom" ], "library_sources": [ - "motifs/matmul/skinny_input_f32_dequant_accumulate_f32.loom", - "motifs/matmul/skinny_postprocess_f32.loom", - "motifs/matmul/skinny_publish_f32.loom", + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", "motifs/dequant.loom", "motifs/elementwise/unary_apply.loom", "motifs/q4_k_f16.loom", @@ -7611,9 +7970,9 @@ ] }, "compile_dependencies": [ - "motifs/matmul/skinny_input_f32_dequant_accumulate_f32.loom", - "motifs/matmul/skinny_postprocess_f32.loom", - "motifs/matmul/skinny_publish_f32.loom", + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", "motifs/dequant.loom", "motifs/elementwise/unary_apply.loom", "motifs/q4_k_f16.loom", @@ -7626,9 +7985,9 @@ ] }, { - "name": "ggml_mul_mat_tiled_input_f32_publish_f32", + "name": "ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", "family": "loom_libs", - "symbol": "ggml_mul_mat_tiled_input_f32_publish_f32", + "symbol": "ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", "source": "ops/matmul/tiled_input_f32_publish_f32.loom", "workload_parameters": [ { @@ -7645,9 +8004,13 @@ "bindings": [ "input", "weight", - "output" + "bias", + "residual_input", + "residual_output" ], "binding_access": [ + "read", + "read", "read", "read", "read_write" @@ -7689,10 +8052,10 @@ ] }, { - "name": "ggml_mul_mat_tiled_input_f32_publish_f32_f16_alternate", + "name": "ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32_k16", "family": "loom_libs", - "symbol": "ggml_mul_mat_tiled_input_f32_publish_f32_f16_alternate", - "source": "ops/matmul/tiled_input_f32_publish_f32.loom", + "symbol": "ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32_k16", + "source": "ops/matmul/tiled_pair_input_f32_binary_publish_f32_k16.loom", "workload_parameters": [ { "name": "token_count", @@ -7707,11 +8070,13 @@ ], "bindings": [ "input", - "weight", + "gate_weight", + "up_weight", "output", "f16_output" ], "binding_access": [ + "read", "read", "read", "read_write", @@ -7721,12 +8086,13 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/matmul/tiled_input_f32_publish_f32.loom" + "ops/matmul/tiled_pair_input_f32_binary_publish_f32_k16.loom" ], "library_sources": [ - "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", - "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_pair_postprocess_f32.loom", "motifs/matmul/tiled_publish.loom", + "motifs/normalization/publish_f16.loom", "motifs/dequant.loom", "motifs/elementwise/unary_apply.loom", "motifs/q4_k_f16.loom", @@ -7739,9 +8105,10 @@ ] }, "compile_dependencies": [ - "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", - "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_pair_postprocess_f32.loom", "motifs/matmul/tiled_publish.loom", + "motifs/normalization/publish_f16.loom", "motifs/dequant.loom", "motifs/elementwise/unary_apply.loom", "motifs/q4_k_f16.loom", @@ -7754,10 +8121,10 @@ ] }, { - "name": "ggml_mul_mat_tiled_input_f32_bias_publish_f32", + "name": "ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32", "family": "loom_libs", - "symbol": "ggml_mul_mat_tiled_input_f32_bias_publish_f32", - "source": "ops/matmul/tiled_input_f32_publish_f32.loom", + "symbol": "ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32", + "source": "ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom", "workload_parameters": [ { "name": "token_count", @@ -7772,8 +8139,8 @@ ], "bindings": [ "input", - "weight", - "bias", + "gate_weight", + "up_weight", "output" ], "binding_access": [ @@ -7786,415 +8153,11 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/matmul/tiled_input_f32_publish_f32.loom" + "ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom" ], "library_sources": [ - "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", - "motifs/matmul/tiled_postprocess_f32.loom", - "motifs/matmul/tiled_publish.loom", - "motifs/dequant.loom", - "motifs/elementwise/unary_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/bf16_f16.loom", - "motifs/f32_f16.loom" - ] - }, - "compile_dependencies": [ - "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", - "motifs/matmul/tiled_postprocess_f32.loom", - "motifs/matmul/tiled_publish.loom", - "motifs/dequant.loom", - "motifs/elementwise/unary_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/bf16_f16.loom", - "motifs/f32_f16.loom" - ] - }, - { - "name": "ggml_mul_mat_tiled_input_f32_residual_publish_f32", - "family": "loom_libs", - "symbol": "ggml_mul_mat_tiled_input_f32_residual_publish_f32", - "source": "ops/matmul/tiled_input_f32_publish_f32.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "input", - "weight", - "residual_input", - "residual_output" - ], - "binding_access": [ - "read", - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "ops/matmul/tiled_input_f32_publish_f32.loom" - ], - "library_sources": [ - "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", - "motifs/matmul/tiled_postprocess_f32.loom", - "motifs/matmul/tiled_publish.loom", - "motifs/dequant.loom", - "motifs/elementwise/unary_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/bf16_f16.loom", - "motifs/f32_f16.loom" - ] - }, - "compile_dependencies": [ - "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", - "motifs/matmul/tiled_postprocess_f32.loom", - "motifs/matmul/tiled_publish.loom", - "motifs/dequant.loom", - "motifs/elementwise/unary_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/bf16_f16.loom", - "motifs/f32_f16.loom" - ] - }, - { - "name": "ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", - "family": "loom_libs", - "symbol": "ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", - "source": "ops/matmul/tiled_input_f32_publish_f32.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "input", - "weight", - "bias", - "residual_input", - "residual_output" - ], - "binding_access": [ - "read", - "read", - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "ops/matmul/tiled_input_f32_publish_f32.loom" - ], - "library_sources": [ - "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", - "motifs/matmul/tiled_postprocess_f32.loom", - "motifs/matmul/tiled_publish.loom", - "motifs/dequant.loom", - "motifs/elementwise/unary_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/bf16_f16.loom", - "motifs/f32_f16.loom" - ] - }, - "compile_dependencies": [ - "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", - "motifs/matmul/tiled_postprocess_f32.loom", - "motifs/matmul/tiled_publish.loom", - "motifs/dequant.loom", - "motifs/elementwise/unary_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/bf16_f16.loom", - "motifs/f32_f16.loom" - ] - }, - { - "name": "ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32_k16", - "family": "loom_libs", - "symbol": "ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32_k16", - "source": "ops/matmul/tiled_pair_input_f32_binary_publish_f32_k16.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "input", - "gate_weight", - "up_weight", - "output", - "f16_output" - ], - "binding_access": [ - "read", - "read", - "read", - "read_write", - "write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "ops/matmul/tiled_pair_input_f32_binary_publish_f32_k16.loom" - ], - "library_sources": [ - "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", - "motifs/matmul/tiled_pair_postprocess_f32.loom", - "motifs/matmul/tiled_publish.loom", - "motifs/normalization/publish_f16.loom", - "motifs/dequant.loom", - "motifs/elementwise/unary_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/bf16_f16.loom", - "motifs/f32_f16.loom" - ] - }, - "compile_dependencies": [ - "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", - "motifs/matmul/tiled_pair_postprocess_f32.loom", - "motifs/matmul/tiled_publish.loom", - "motifs/normalization/publish_f16.loom", - "motifs/dequant.loom", - "motifs/elementwise/unary_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/bf16_f16.loom", - "motifs/f32_f16.loom" - ] - }, - { - "name": "ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32", - "family": "loom_libs", - "symbol": "ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32", - "source": "ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "input", - "gate_weight", - "up_weight", - "output" - ], - "binding_access": [ - "read", - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom" - ], - "library_sources": [ - "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", - "motifs/matmul/tiled_pair_postprocess_f32.loom", - "motifs/matmul/tiled_postprocess_f32.loom", - "motifs/matmul/tiled_publish.loom", - "motifs/dequant.loom", - "motifs/elementwise/unary_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/bf16_f16.loom", - "motifs/f32_f16.loom" - ] - }, - "compile_dependencies": [ - "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", - "motifs/matmul/tiled_pair_postprocess_f32.loom", - "motifs/matmul/tiled_postprocess_f32.loom", - "motifs/matmul/tiled_publish.loom", - "motifs/dequant.loom", - "motifs/elementwise/unary_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/bf16_f16.loom", - "motifs/f32_f16.loom" - ] - }, - { - "name": "ggml_mul_mat_tiled_pair_input_f32_binary_bias_publish_f32", - "family": "loom_libs", - "symbol": "ggml_mul_mat_tiled_pair_input_f32_binary_bias_publish_f32", - "source": "ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "input", - "gate_weight", - "up_weight", - "bias", - "output" - ], - "binding_access": [ - "read", - "read", - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom" - ], - "library_sources": [ - "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", - "motifs/matmul/tiled_pair_postprocess_f32.loom", - "motifs/matmul/tiled_postprocess_f32.loom", - "motifs/matmul/tiled_publish.loom", - "motifs/dequant.loom", - "motifs/elementwise/unary_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/bf16_f16.loom", - "motifs/f32_f16.loom" - ] - }, - "compile_dependencies": [ - "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", - "motifs/matmul/tiled_pair_postprocess_f32.loom", - "motifs/matmul/tiled_postprocess_f32.loom", - "motifs/matmul/tiled_publish.loom", - "motifs/dequant.loom", - "motifs/elementwise/unary_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/bf16_f16.loom", - "motifs/f32_f16.loom" - ] - }, - { - "name": "ggml_mul_mat_tiled_pair_input_f32_binary_residual_publish_f32", - "family": "loom_libs", - "symbol": "ggml_mul_mat_tiled_pair_input_f32_binary_residual_publish_f32", - "source": "ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "input", - "gate_weight", - "up_weight", - "residual_input", - "residual_output" - ], - "binding_access": [ - "read", - "read", - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom" - ], - "library_sources": [ - "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", - "motifs/matmul/tiled_pair_postprocess_f32.loom", + "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_pair_postprocess_f32.loom", "motifs/matmul/tiled_postprocess_f32.loom", "motifs/matmul/tiled_publish.loom", "motifs/dequant.loom", @@ -8339,6 +8302,21 @@ "--output-format=jsonl" ], "source": "ops/rows/get_rows_f32.loom" + }, + { + "name": "get_rows_scale_f32_plan_test", + "args": [ + "$(location ops/rows/get_rows_scale_f32.loom)", + "--benchmark=@ggml_get_rows_scale_f32_f32", + "--config=ggml.get_rows_f32.token_capacity=4", + "--config=ggml.get_rows_f32.hidden_capacity=4", + "--config=ggml.get_rows_f32.weight_format=32", + "--config=ggml.get_rows_scale_f32.scale=2.0", + "--config=ggml.get_rows_scale_f32.bias=0.25", + "--dry-run", + "--output-format=jsonl" + ], + "source": "ops/rows/get_rows_scale_f32.loom" } ], "sources": [], diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/attention/rmsnorm_neox_cache_f16.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/attention/rmsnorm_neox_cache_f16.loom new file mode 100644 index 000000000000..876b2be71732 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/attention/rmsnorm_neox_cache_f16.loom @@ -0,0 +1,168 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @qwen3_moe.attention.postprocess.head_body(%publish_output: i1, %token_count: index, %cache_row_count: index, %head_domain0: index, %token0: index, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer, %query_input: buffer, %key_input: buffer, %value_input: buffer, %query_norm_weight: buffer, %key_norm_weight: buffer, %inverse_frequencies: buffer, %query_output: buffer, %key_cache: buffer, %value_cache: buffer) + +config.decl @qwen3_moe.model.rms_epsilon : f32 + +config.decl @qwen3_moe.attention.head_size : %value: index where [range(%value, 4, 1024), mul(%value, 4)] + +config.decl @qwen3_moe.attention.query_size : %value: index where [range(%value, 1, 262144)] + +config.decl @qwen3_moe.attention.key_value_size : %value: index where [range(%value, 1, 262144)] + +config.decl @qwen3_moe.attention.rope_mscale : f32 + +func.def inline @qwen3_moe_rmsnorm_neox_packet(%row_sum: f32, %head_size: f32, %epsilon: f32, %position: f32, %rope_mscale: f32, %inverse_frequencies: vector<2xf32>, %low_values: vector<2xf32>, %high_values: vector<2xf32>, %low_weights: vector<2xf32>, %high_weights: vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) { + %mean = scalar.divf %row_sum, %head_size : f32 + %biased_mean = scalar.addf %mean, %epsilon : f32 + %scale = scalar.rsqrtf %biased_mean : f32 + %scale_vector = vector.splat %scale : vector<2xf32> + %position_vector = vector.splat %position : vector<2xf32> + %inverse_two_pi = scalar.constant 0.15915494309189535 : f32 + %inverse_two_pi_vector = vector.splat %inverse_two_pi : vector<2xf32> + %normalized_low = vector.mulf %low_values, %scale_vector : vector<2xf32> + %normalized_high = vector.mulf %high_values, %scale_vector : vector<2xf32> + %scaled_low = vector.mulf %normalized_low, %low_weights : vector<2xf32> + %scaled_high = vector.mulf %normalized_high, %high_weights : vector<2xf32> + %angles = vector.mulf %position_vector, %inverse_frequencies : vector<2xf32> + %turns = vector.mulf %angles, %inverse_two_pi_vector : vector<2xf32> + %cosines = vector.costurnsf %turns : vector<2xf32> + %sines = vector.sinturnsf %turns : vector<2xf32> + %low_cosines = vector.mulf %scaled_low, %cosines : vector<2xf32> + %high_sines = vector.mulf %scaled_high, %sines : vector<2xf32> + %low_sines = vector.mulf %scaled_low, %sines : vector<2xf32> + %high_cosines = vector.mulf %scaled_high, %cosines : vector<2xf32> + %rotated_low0 = vector.subf %low_cosines, %high_sines : vector<2xf32> + %rotated_high0 = vector.addf %low_sines, %high_cosines : vector<2xf32> + %rope_mscale_vector = vector.splat %rope_mscale : vector<2xf32> + %rotated_low = vector.mulf %rotated_low0, %rope_mscale_vector : vector<2xf32> + %rotated_high = vector.mulf %rotated_high0, %rope_mscale_vector : vector<2xf32> + func.return %rotated_low, %rotated_high : vector<2xf32>, vector<2xf32> +} +// Processes one head domain after its raw projection row is visible. Callers +// may use a larger workgroup than the packet count; inactive workitems +// contribute zero to normalization and perform no loads or stores. +template.def<@qwen3_moe.attention.postprocess.head_body> device @qwen3_moe_attention_postprocess_head_body(%publish_output: i1, %token_count: index, %cache_row_count: index, %head_domain0: index, %token0: index, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer, %query_input: buffer, %key_input: buffer, %value_input: buffer, %query_norm_weight: buffer, %key_norm_weight: buffer, %inverse_frequencies: buffer, %query_output: buffer, %key_cache: buffer, %value_cache: buffer) { + %query_size0 = config.get @qwen3_moe.attention.query_size : index + %key_value_size0 = config.get @qwen3_moe.attention.key_value_size : index + %head_size0 = config.get @qwen3_moe.attention.head_size : index + %rope_mscale = config.get @qwen3_moe.attention.rope_mscale : f32 + %epsilon = config.get @qwen3_moe.model.rms_epsilon : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_cache_row_count = index.assume %cache_row_count [range(%cache_row_count, 1, 1048576)] : index + %query_size, %key_value_size, %head_size = index.assume %query_size0, %key_value_size0, %head_size0 [range(%query_size0, 1, 262144), range(%key_value_size0, 1, 262144), range(%head_size0, 4, 1024), mul(%head_size0, 4), mul(%query_size0, %head_size0), mul(%key_value_size0, %head_size0)] : index, index, index + %channel0 = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c0_f32 = scalar.constant 0.0 : f32 + %c0_offset = index.constant 0 : offset + %token, %launch_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %half_head_size = index.div %head_size, %c2 : index + %pair_packet_count = index.div %head_size, %c4 : index + %query_head_count = index.div %query_size, %head_size : index + %key_value_head_count = index.div %key_value_size, %head_size : index + %key_value_domain_count = index.mul %key_value_head_count, %c2 : index + %head_domain_count = index.add %query_head_count, %key_value_domain_count : index + %head_domain = index.assume %head_domain0 [lt(%head_domain0, %head_domain_count)] : index + %key_domain_end = index.add %query_head_count, %key_value_head_count : index + %is_query = index.cmp ult, %head_domain, %query_head_count : index + %is_query_or_key = index.cmp ult, %head_domain, %key_domain_end : index + %key_value_head = index.rem %head_domain, %key_value_head_count : index + %active_channel = index.cmp ult, %channel0, %pair_packet_count : index + %head_size_i32 = index.cast %head_size : index to i32 + %head_size_f32 = scalar.sitofp %head_size_i32 : i32 to f32 + %positions_noalias, %key_cache_indices_noalias, %value_cache_indices_noalias, %query_input_noalias, %key_input_noalias, %value_input_noalias, %query_norm_weight_noalias, %key_norm_weight_noalias, %inverse_frequencies_noalias, %query_output_noalias, %key_cache_noalias, %value_cache_noalias = buffer.assume.noalias %positions, %key_cache_indices, %value_cache_indices, %query_input, %key_input, %value_input, %query_norm_weight, %key_norm_weight, %inverse_frequencies, %query_output, %key_cache, %value_cache : buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer + %positions_view = buffer.view %positions_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi32> + %key_cache_indices_view = buffer.view %key_cache_indices_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi64> + %value_cache_indices_view = buffer.view %value_cache_indices_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi64> + %query_input_view = buffer.view %query_input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> + %key_input_view = buffer.view %key_input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> + %value_input_view = buffer.view %value_input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> + %query_norm_weight_view = buffer.view %query_norm_weight_noalias[%c0_offset] : buffer -> view<[%head_size]xf32> + %key_norm_weight_view = buffer.view %key_norm_weight_noalias[%c0_offset] : buffer -> view<[%head_size]xf32> + %inverse_frequencies_view = buffer.view %inverse_frequencies_noalias[%c0_offset] : buffer -> view<[%half_head_size]xf32> + %query_output_view = buffer.view %query_output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> + %key_cache_words_view = buffer.view %key_cache_noalias[%c0_offset] : buffer -> view<[%bounded_cache_row_count]x[%key_value_head_count]x[%half_head_size]xi32> + %value_cache_words_view = buffer.view %value_cache_noalias[%c0_offset] : buffer -> view<[%bounded_cache_row_count]x[%key_value_head_count]x[%half_head_size]xi32> + %row_sum = scf.if %is_query_or_key -> (f32) { + %partial_sum = scf.if %publish_output -> (f32) { + %channel_sum = scf.if %active_channel -> (f32) { + %channel = index.assume %channel0 [lt(%channel0, %pair_packet_count)] : index + %reduction_channel = index.mul %channel, %c4 : index + %reduction_values = scf.if %is_query -> (vector<4xf32>) { + %query_values = vector.load %query_input_view[%token, %head_domain, %reduction_channel] : view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> -> vector<4xf32> + scf.yield %query_values : vector<4xf32> + } else { + %key_values = vector.load %key_input_view[%token, %key_value_head, %reduction_channel] : view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> -> vector<4xf32> + scf.yield %key_values : vector<4xf32> + } + %squares = vector.mulf %reduction_values, %reduction_values : vector<4xf32> + %sum = vector.reduce %squares, %c0_f32 : vector<4xf32>, f32 + scf.yield %sum : f32 + } else { + scf.yield %c0_f32 : f32 + } + scf.yield %channel_sum : f32 + } else { + scf.yield %c0_f32 : f32 + } + %sum = kernel.workgroup.reduce %partial_sum : f32 + scf.yield %sum : f32 + } else { + scf.yield %c0_f32 : f32 + } + scf.if %publish_output { + scf.if %active_channel { + %channel = index.assume %channel0 [lt(%channel0, %pair_packet_count)] : index + %pair_channel = index.mul %channel, %c2 : index + %paired_channel = index.add %pair_channel, %half_head_size : index + %paired_word = index.add %channel, %pair_packet_count : index + %reduction_channel = index.mul %channel, %c4 : index + scf.if %is_query { + %low_values = vector.load %query_input_view[%token, %head_domain, %pair_channel] : view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> -> vector<2xf32> + %high_values = vector.load %query_input_view[%token, %head_domain, %paired_channel] : view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> -> vector<2xf32> + %low_weights = vector.load %query_norm_weight_view[%pair_channel] : view<[%head_size]xf32> -> vector<2xf32> + %high_weights = vector.load %query_norm_weight_view[%paired_channel] : view<[%head_size]xf32> -> vector<2xf32> + %inverse_frequencies_packet = vector.load %inverse_frequencies_view[%pair_channel] : view<[%half_head_size]xf32> -> vector<2xf32> + %position_i32 = view.load %positions_view[%token] : view<[%launch_token_count]xi32> -> i32 + %position = scalar.sitofp %position_i32 : i32 to f32 + %rotated_low, %rotated_high = func.call @qwen3_moe_rmsnorm_neox_packet(%row_sum, %head_size_f32, %epsilon, %position, %rope_mscale, %inverse_frequencies_packet, %low_values, %high_values, %low_weights, %high_weights) : (f32, f32, f32, f32, f32, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) + vector.store %rotated_low, %query_output_view[%token, %head_domain, %pair_channel] : vector<2xf32>, view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> + vector.store %rotated_high, %query_output_view[%token, %head_domain, %paired_channel] : vector<2xf32>, view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> + } else { + scf.if %is_query_or_key { + %low_values = vector.load %key_input_view[%token, %key_value_head, %pair_channel] : view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> -> vector<2xf32> + %high_values = vector.load %key_input_view[%token, %key_value_head, %paired_channel] : view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> -> vector<2xf32> + %low_weights = vector.load %key_norm_weight_view[%pair_channel] : view<[%head_size]xf32> -> vector<2xf32> + %high_weights = vector.load %key_norm_weight_view[%paired_channel] : view<[%head_size]xf32> -> vector<2xf32> + %inverse_frequencies_packet = vector.load %inverse_frequencies_view[%pair_channel] : view<[%half_head_size]xf32> -> vector<2xf32> + %position_i32 = view.load %positions_view[%token] : view<[%launch_token_count]xi32> -> i32 + %position = scalar.sitofp %position_i32 : i32 to f32 + %cache_index_raw = view.load %key_cache_indices_view[%token] : view<[%launch_token_count]xi64> -> i64 + %cache_index_i64 = scalar.assume %cache_index_raw [range(%cache_index_raw, 0, 1048575)] : i64 + %cache_index0 = index.cast %cache_index_i64 : i64 to index + %cache_index = index.assume %cache_index0 [lt(%cache_index0, %bounded_cache_row_count)] : index + %rotated_low, %rotated_high = func.call @qwen3_moe_rmsnorm_neox_packet(%row_sum, %head_size_f32, %epsilon, %position, %rope_mscale, %inverse_frequencies_packet, %low_values, %high_values, %low_weights, %high_weights) : (f32, f32, f32, f32, f32, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) + %half_low = vector.fptrunc %rotated_low : vector<2xf32> to vector<2xf16> + %half_high = vector.fptrunc %rotated_high : vector<2xf32> to vector<2xf16> + %packed_low = vector.bitcast %half_low : vector<2xf16> to vector<1xi32> + %packed_high = vector.bitcast %half_high : vector<2xf16> to vector<1xi32> + vector.store %packed_low, %key_cache_words_view[%cache_index, %key_value_head, %channel] : vector<1xi32>, view<[%bounded_cache_row_count]x[%key_value_head_count]x[%half_head_size]xi32> + vector.store %packed_high, %key_cache_words_view[%cache_index, %key_value_head, %paired_word] : vector<1xi32>, view<[%bounded_cache_row_count]x[%key_value_head_count]x[%half_head_size]xi32> + } else { + %cache_index_raw = view.load %value_cache_indices_view[%token] : view<[%launch_token_count]xi64> -> i64 + %cache_index_i64 = scalar.assume %cache_index_raw [range(%cache_index_raw, 0, 1048575)] : i64 + %cache_index0 = index.cast %cache_index_i64 : i64 to index + %cache_index = index.assume %cache_index0 [lt(%cache_index0, %bounded_cache_row_count)] : index + %values = vector.load %value_input_view[%token, %key_value_head, %reduction_channel] : view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> -> vector<4xf32> + %half_values = vector.fptrunc %values : vector<4xf32> to vector<4xf16> + %packed_values = vector.bitcast %half_values : vector<4xf16> to vector<2xi32> + vector.store %packed_values, %value_cache_words_view[%cache_index, %key_value_head, %pair_channel] : vector<2xi32>, view<[%bounded_cache_row_count]x[%key_value_head_count]x[%half_head_size]xi32> + } + } + } + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/q4_q6_q8_1_x4_qkv.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/q4_q6_q8_1_x4_qkv.loom new file mode 100644 index 000000000000..63f8c97a55a9 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/q4_q6_q8_1_x4_qkv.loom @@ -0,0 +1,113 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @qwen3_moe.attention.qkv_quantized.body(%value_uses_q6_index: index, %publish_output: i1, %token_count: index, %token0: index, %q8_input: buffer, %query_weight: buffer, %key_weight: buffer, %value_weight: buffer, %query_output: buffer, %key_output: buffer, %value_output: buffer) + +config.decl @qwen3_moe.model.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] + +config.decl @qwen3_moe.attention.query_size : %value: index where [range(%value, 1, 262144)] + +config.decl @qwen3_moe.attention.key_value_size : %value: index where [range(%value, 1, 262144)] + +func.decl @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32) + +func.decl @ggml_q6k_q8_1_x4_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32) + +// Device body shared by the ordinary projection and completion-fused +// postprocess exports. Each caller owns its boundary after the raw row stores. +template.def<@qwen3_moe.attention.qkv_quantized.body> device @qwen3_moe_attention_qkv_quantized_body(%value_uses_q6_index: index, %publish_output: i1, %token_count: index, %token0: index, %q8_input: buffer, %query_weight: buffer, %key_weight: buffer, %value_weight: buffer, %query_output: buffer, %key_output: buffer, %value_output: buffer) { + %hidden_size0 = config.get @qwen3_moe.model.hidden_size : index + %query_size0 = config.get @qwen3_moe.attention.query_size : index + %key_value_size0 = config.get @qwen3_moe.attention.key_value_size : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %hidden_size = index.assume %hidden_size0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128)] : index + %query_size, %key_value_size = index.assume %query_size0, %key_value_size0 [range(%query_size0, 1, 262144), range(%key_value_size0, 1, 262144), mul(%query_size0, %key_value_size0)] : index, index + %channel_tile = kernel.workgroup.id : index + %subgroup0 = kernel.subgroup.id : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c8 = index.constant 8 : index + %c128 = index.constant 128 : index + %c144_bytes = index.constant 144 : offset + %c210_bytes = index.constant 210 : offset + %c256 = index.constant 256 : index + %c0_i32 = scalar.constant 0 : i32 + %c0_f32 = scalar.constant 0.0 : f32 + %c0_offset = index.constant 0 : offset + %token, %launch_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %channel_base = index.mul %channel_tile, %c8 : index + %global_channel = index.add %channel_base, %subgroup : index + %key_value_end = index.add %query_size, %key_value_size : index + %key_value_output_size = index.mul %key_value_size, %c2 : index + %total_output_size = index.add %query_size, %key_value_output_size : index + %valid_channel = index.cmp ult, %global_channel, %total_output_size : index + %is_query = index.cmp ult, %global_channel, %query_size : index + %is_key = index.cmp ult, %global_channel, %key_value_end : index + %key_value_channel = index.rem %global_channel, %key_value_size : index + %value_uses_q6 = index.cmp eq, %value_uses_q6_index, %c1 : index + %lane_i32 = index.cast %lane : index to i32 + %is_lane_zero = scalar.cmpi eq, %lane_i32, %c0_i32 : i32 + %quant_block_count = index.div %hidden_size, %c256 : index + %q4_row_bytes = index.scale %quant_block_count, %c144_bytes : index, offset -> offset + %q6_row_bytes = index.scale %quant_block_count, %c210_bytes : index, offset -> offset + %q8_group_count = index.div %hidden_size, %c128 : index + %q8_row_bytes = index.scale %q8_group_count, %c144_bytes : index, offset -> offset + %q8_row_byte_base = index.scale %token, %q8_row_bytes : index, offset -> offset + %q8_noalias, %query_weight_noalias, %key_weight_noalias, %value_weight_noalias, %query_output_noalias, %key_output_noalias, %value_output_noalias = buffer.assume.noalias %q8_input, %query_weight, %key_weight, %value_weight, %query_output, %key_output, %value_output : buffer, buffer, buffer, buffer, buffer, buffer, buffer + %lane_sum = scf.if %publish_output -> (f32) { + %channel_sum = scf.if %valid_channel -> (f32) { + %projection_sum = scf.if %is_query -> (f32) { + %row_byte_base = index.scale %global_channel, %q4_row_bytes : index, offset -> offset + %sum = func.call @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%hidden_size, %query_weight_noalias, %row_byte_base, %q8_noalias, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) + scf.yield %sum : f32 + } else { + %key_or_value_sum = scf.if %is_key -> (f32) { + %row_byte_base = index.scale %key_value_channel, %q4_row_bytes : index, offset -> offset + %sum = func.call @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%hidden_size, %key_weight_noalias, %row_byte_base, %q8_noalias, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) + scf.yield %sum : f32 + } else { + %value_sum = scf.if %value_uses_q6 -> (f32) { + %row_byte_base = index.scale %key_value_channel, %q6_row_bytes : index, offset -> offset + %sum = func.call @ggml_q6k_q8_1_x4_row_lane(%hidden_size, %value_weight_noalias, %row_byte_base, %q8_noalias, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) + scf.yield %sum : f32 + } else { + %row_byte_base = index.scale %key_value_channel, %q4_row_bytes : index, offset -> offset + %sum = func.call @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%hidden_size, %value_weight_noalias, %row_byte_base, %q8_noalias, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) + scf.yield %sum : f32 + } + scf.yield %value_sum : f32 + } + scf.yield %key_or_value_sum : f32 + } + scf.yield %projection_sum : f32 + } else { + scf.yield %c0_f32 : f32 + } + scf.yield %channel_sum : f32 + } else { + scf.yield %c0_f32 : f32 + } + %dot = kernel.subgroup.reduce %lane_sum : f32 + %query_output_view = buffer.view %query_output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%query_size]xf32> + %key_output_view = buffer.view %key_output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%key_value_size]xf32> + %value_output_view = buffer.view %value_output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%key_value_size]xf32> + scf.if %publish_output { + scf.if %valid_channel { + scf.if %is_lane_zero { + scf.if %is_query { + view.store %dot, %query_output_view[%token, %global_channel] : f32, view<[%launch_token_count]x[%query_size]xf32> + } else { + scf.if %is_key { + view.store %dot, %key_output_view[%token, %key_value_channel] : f32, view<[%launch_token_count]x[%key_value_size]xf32> + } else { + view.store %dot, %value_output_view[%token, %key_value_channel] : f32, view<[%launch_token_count]x[%key_value_size]xf32> + } + } + } + } + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/q4_q8_1_x4_decode.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/q4_q8_1_x4_decode.loom new file mode 100644 index 000000000000..7534dcbe7b3c --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/q4_q8_1_x4_decode.loom @@ -0,0 +1,200 @@ +// Packed Q4_K x Q8_1 x4 row contraction shared by dense wrappers. +// The symbols and bodies are preserved from the retired routed gate-up packet. + +func.def inline @qwen3_moe_q4k_scale_from_header(%scale0: i32, %scale1: i32, %scale2: i32, %q4_group: index) -> (i32, i32) { + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c2_i32 = scalar.constant 2 : i32 + %c4_i32 = scalar.constant 4 : i32 + %c15_i32 = scalar.constant 15 : i32 + %c48_i32 = scalar.constant 48 : i32 + %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index + %is_low_group = index.cmp ult, %bounded_group, %c4 : index + %scale_lane = index.rem %bounded_group, %c4 : index + %scale_shift_index = index.mul %scale_lane, %c8 : index + %scale_shift = index.cast %scale_shift_index : index to i32 + %high_shift = scalar.addi %scale_shift, %c2_i32 : i32 + %minimum_shift = scalar.addi %scale_shift, %c4_i32 : i32 + %selected_scale_source = scf.select %is_low_group, %scale0, %scale2 : i32 + %selected_minimum_source = scf.select %is_low_group, %scale1, %scale2 : i32 + %selected_scale_high_shift = scf.select %is_low_group, %scale_shift, %high_shift : i32 + %selected_minimum_low_shift = scf.select %is_low_group, %scale_shift, %minimum_shift : i32 + %scale_low0 = scalar.shrui %selected_scale_source, %scale_shift : i32 + %scale_low = scalar.andi %scale_low0, %c15_i32 : i32 + %scale_high0 = scalar.shrui %scale0, %selected_scale_high_shift : i32 + %scale_high = scalar.andi %scale_high0, %c48_i32 : i32 + %scale = scalar.ori %scale_low, %scale_high : i32 + %minimum_low0 = scalar.shrui %selected_minimum_source, %selected_minimum_low_shift : i32 + %minimum_low = scalar.andi %minimum_low0, %c15_i32 : i32 + %minimum_high0 = scalar.shrui %scale1, %selected_scale_high_shift : i32 + %minimum_high = scalar.andi %minimum_high0, %c48_i32 : i32 + %minimum = scalar.ori %minimum_low, %minimum_high : i32 + func.return %scale, %minimum : i32, i32 +} + +func.def inline @qwen3_moe_q4k_chunk_pair_global(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group_pair: index, %q4_half: index, %header_words: vector<4xi32>) -> (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) { + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %block_bytes = index.constant 144 : offset + %code_byte_add = index.constant 16 : offset + %c4_i32 = scalar.constant 4 : i32 + %nibble_mask = vector.constant 252645135 : vector<4xi32> + %bounded_pair = index.assume %q4_group_pair [range(%q4_group_pair, 0, 3)] : index + %bounded_half = index.assume %q4_half [range(%q4_half, 0, 1)] : index + %block_byte_add = index.scale %q4_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %code_byte_base = index.add %block_byte_base, %code_byte_add : offset + %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<32xi32> + %header_halves = vector.bitcast %header_words : vector<4xi32> to vector<8xf16> + %d_f16 = vector.extract %header_halves[0] : vector<8xf16> -> f16 + %dmin_f16 = vector.extract %header_halves[1] : vector<8xf16> -> f16 + %scale0 = vector.extract %header_words[1] : vector<4xi32> -> i32 + %scale1 = vector.extract %header_words[2] : vector<4xi32> -> i32 + %scale2 = vector.extract %header_words[3] : vector<4xi32> -> i32 + %d = scalar.extf %d_f16 : f16 to f32 + %dmin = scalar.extf %dmin_f16 : f16 to f32 + %pair_code_base = index.mul %bounded_pair, %c8 : index + %half_code_add = index.mul %bounded_half, %c4 : index + %code_index0 = index.add %pair_code_base, %half_code_add : index + %code_index = index.assume %code_index0 [range(%code_index0, 0, 28)] : index + %packed_codes = vector.load %code_view[%code_index] : view<32xi32> -> vector<4xi32> + %low_codes = vector.andi %packed_codes, %nibble_mask : vector<4xi32> + %c4_i32v = vector.splat %c4_i32 : vector<4xi32> + %high_shifted = vector.shrui %packed_codes, %c4_i32v : vector<4xi32> + %high_codes = vector.andi %high_shifted, %nibble_mask : vector<4xi32> + %q4_low = vector.bitcast %low_codes : vector<4xi32> to vector<16xi8> + %q4_high = vector.bitcast %high_codes : vector<4xi32> to vector<16xi8> + %low_group = index.mul %bounded_pair, %c2 : index + %high_group = index.add %low_group, %c1 : index + %low_scale, %low_minimum = func.call @qwen3_moe_q4k_scale_from_header(%scale0, %scale1, %scale2, %low_group) : (i32, i32, i32, index) -> (i32, i32) + %high_scale, %high_minimum = func.call @qwen3_moe_q4k_scale_from_header(%scale0, %scale1, %scale2, %high_group) : (i32, i32, i32, index) -> (i32, i32) + %low_scale_f32 = scalar.uitofp %low_scale : i32 to f32 + %low_minimum_f32 = scalar.uitofp %low_minimum : i32 to f32 + %high_scale_f32 = scalar.uitofp %high_scale : i32 to f32 + %high_minimum_f32 = scalar.uitofp %high_minimum : i32 to f32 + %low_d_scale = scalar.mulf %d, %low_scale_f32 : f32 + %low_dmin_scale = scalar.mulf %dmin, %low_minimum_f32 : f32 + %high_d_scale = scalar.mulf %d, %high_scale_f32 : f32 + %high_dmin_scale = scalar.mulf %dmin, %high_minimum_f32 : f32 + func.return %q4_low, %low_d_scale, %low_dmin_scale, %q4_high, %high_d_scale, %high_dmin_scale : vector<16xi8>, f32, f32, vector<16xi8>, f32, f32 +} + +func.def inline @qwen3_moe_q4k_q8_1_dot(%q4_values: vector<16xi8>, %d_scale: f32, %dmin_scale: f32, %q8_values: vector<16xi8>, %q8_d: f32, %q8_s: f32) -> (f32) { + %c0_i32 = scalar.constant 0 : i32 + %c0_i32v = vector.constant 0 : vector<4xi32> + %half_f32 = scalar.constant 0.5 : f32 + %partial_dots = vector.dot4i %q4_values, %q8_values, %c0_i32v : vector<16xi8>, vector<16xi8>, vector<4xi32> + %q_sum = vector.reduce %partial_dots, %c0_i32 : vector<4xi32>, i32 + %q_sum_f32 = scalar.sitofp %q_sum : i32 to f32 + %scaled_dot0 = scalar.mulf %q8_d, %d_scale : f32 + %scaled_dot = scalar.mulf %scaled_dot0, %q_sum_f32 : f32 + %q8_half_sum = scalar.mulf %q8_s, %half_f32 : f32 + %minimum_correction = scalar.mulf %dmin_scale, %q8_half_sum : f32 + %contribution = scalar.subf %scaled_dot, %minimum_correction : f32 + func.return %contribution : f32 +} + +func.def inline @qwen3_moe_q4k_q8_1_x4_paired_block_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %q4_block: index, %block_lane: index) -> (f32) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_q4_block0 = index.assume %q4_block [range(%q4_block, 0, 127)] : index + %bounded_block_lane = index.assume %block_lane [range(%block_lane, 0, 7)] : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %q4_block_bytes = index.constant 144 : offset + %q8_group_bytes = index.constant 144 : offset + %q8_payload_byte_add = index.constant 16 : offset + %q4_block_count = index.div %bounded_input_size, %c256 : index + %q8_group_count = index.div %bounded_input_size, %c128 : index + %bounded_q4_block, %bounded_q4_block_count = index.assume %bounded_q4_block0, %q4_block_count [lt(%bounded_q4_block0, %q4_block_count)] : index, index + %q4_group_pair0 = index.div %bounded_block_lane, %c2 : index + %q4_group_pair = index.assume %q4_group_pair0 [range(%q4_group_pair0, 0, 3)] : index + %q4_half0 = index.rem %bounded_block_lane, %c2 : index + %q4_half = index.assume %q4_half0 [range(%q4_half0, 0, 1)] : index + %q8_group_in_block0 = index.div %q4_group_pair, %c2 : index + %q8_group_in_block = index.assume %q8_group_in_block0 [range(%q8_group_in_block0, 0, 1)] : index + %pair_in_q8_group0 = index.rem %q4_group_pair, %c2 : index + %pair_in_q8_group = index.assume %pair_in_q8_group0 [range(%pair_in_q8_group0, 0, 1)] : index + %q8_low_inner_block0 = index.mul %pair_in_q8_group, %c2 : index + %q8_low_inner_block = index.assume %q8_low_inner_block0 [range(%q8_low_inner_block0, 0, 2)] : index + %q8_high_inner_block0 = index.add %q8_low_inner_block, %c1 : index + %q8_high_inner_block = index.assume %q8_high_inner_block0 [range(%q8_high_inner_block0, 1, 3)] : index + %q8_half_word_add = index.mul %q4_half, %c4 : index + %q8_low_inner_word_base = index.mul %q8_low_inner_block, %c8 : index + %q8_low_word_index0 = index.add %q8_low_inner_word_base, %q8_half_word_add : index + %q8_low_word_index = index.assume %q8_low_word_index0 [range(%q8_low_word_index0, 0, 20)] : index + %q8_high_inner_word_base = index.mul %q8_high_inner_block, %c8 : index + %q8_high_word_index0 = index.add %q8_high_inner_word_base, %q8_half_word_add : index + %q8_high_word_index = index.assume %q8_high_word_index0 [range(%q8_high_word_index0, 8, 28)] : index + %q8_low_ds_index0 = index.mul %q8_low_inner_block, %c2 : index + %q8_low_ds_index = index.assume %q8_low_ds_index0 [range(%q8_low_ds_index0, 0, 4)] : index + %q8_block_group_base = index.mul %bounded_q4_block, %c2 : index + %q8_group0 = index.add %q8_block_group_base, %q8_group_in_block : index + %q8_group, %bounded_q8_group_count = index.assume %q8_group0, %q8_group_count [lt(%q8_group0, %q8_group_count)] : index, index + %q8_group_byte_add = index.scale %q8_group, %q8_group_bytes : index, offset -> offset + %q8_group_byte_base = index.add %q8_row_byte_base, %q8_group_byte_add : offset + %q8_payload_byte_base = index.add %q8_group_byte_base, %q8_payload_byte_add : offset + %q8_ds_view = buffer.view %q8_input[%q8_group_byte_base] : buffer -> view<8xf16> + %q8_words_view = buffer.view %q8_input[%q8_payload_byte_base] : buffer -> view<32xi32> + %q8_ds = vector.load %q8_ds_view[%q8_low_ds_index] : view<8xf16> -> vector<4xf16> + %q8_low_d_f16 = vector.extract %q8_ds[0] : vector<4xf16> -> f16 + %q8_low_s_f16 = vector.extract %q8_ds[1] : vector<4xf16> -> f16 + %q8_high_d_f16 = vector.extract %q8_ds[2] : vector<4xf16> -> f16 + %q8_high_s_f16 = vector.extract %q8_ds[3] : vector<4xf16> -> f16 + %q8_low_d = scalar.extf %q8_low_d_f16 : f16 to f32 + %q8_low_s = scalar.extf %q8_low_s_f16 : f16 to f32 + %q8_high_d = scalar.extf %q8_high_d_f16 : f16 to f32 + %q8_high_s = scalar.extf %q8_high_s_f16 : f16 to f32 + %q8_low_words = vector.load %q8_words_view[%q8_low_word_index] : view<32xi32> -> vector<4xi32> + %q8_high_words = vector.load %q8_words_view[%q8_high_word_index] : view<32xi32> -> vector<4xi32> + %q8_low_values = vector.bitcast %q8_low_words : vector<4xi32> to vector<16xi8> + %q8_high_values = vector.bitcast %q8_high_words : vector<4xi32> to vector<16xi8> + %q4_block_byte_add = index.scale %bounded_q4_block, %q4_block_bytes : index, offset -> offset + %q4_block_byte_base = index.add %weight_row_byte_base, %q4_block_byte_add : offset + %q4_header_view = buffer.view %weight[%q4_block_byte_base] : buffer -> view<4xi32> + %q4_header_words = vector.load %q4_header_view[0] : view<4xi32> -> vector<4xi32> + %q4_low, %low_d_scale, %low_dmin_scale, %q4_high, %high_d_scale, %high_dmin_scale = func.call @qwen3_moe_q4k_chunk_pair_global(%weight, %weight_row_byte_base, %bounded_q4_block, %q4_group_pair, %q4_half, %q4_header_words) : (buffer, offset, index, index, index, vector<4xi32>) -> (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) + %low = func.call @qwen3_moe_q4k_q8_1_dot(%q4_low, %low_d_scale, %low_dmin_scale, %q8_low_values, %q8_low_d, %q8_low_s) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) + %high = func.call @qwen3_moe_q4k_q8_1_dot(%q4_high, %high_d_scale, %high_dmin_scale, %q8_high_values, %q8_high_d, %q8_high_s) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) + %pair = scalar.addf %low, %high : f32 + func.return %pair : f32 +} + +func.def inline @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_lane = index.assume %lane [range(%lane, 0, 31)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c256 = index.constant 256 : index + %c1023 = index.constant 1023 : index + %c1024 = index.constant 1024 : index + %c0_f32 = scalar.constant 0.0 : f32 + %q4_block_count = index.div %bounded_input_size, %c256 : index + %padded_input_size = index.add %bounded_input_size, %c1023 : index + %iteration_count = index.div %padded_input_size, %c1024 : index + %lane_q4_block = index.div %bounded_lane, %c8 : index + %block_lane0 = index.rem %bounded_lane, %c8 : index + %block_lane = index.assume %block_lane0 [range(%block_lane0, 0, 7)] : index + %sum = scf.for %iteration = [%c0 to %iteration_count step %c1](%iteration_acc = %c0_f32 : f32) -> (f32) unroll { + %iteration_q4_block = index.mul %iteration, %c4 : index + %q4_block0 = index.add %iteration_q4_block, %lane_q4_block : index + %valid_q4_block = index.cmp ult, %q4_block0, %q4_block_count : index + %contribution = scf.if %valid_q4_block -> (f32) { + %q4_block, %bounded_q4_block_count = index.assume %q4_block0, %q4_block_count [lt(%q4_block0, %q4_block_count)] : index, index + %pair = func.call @qwen3_moe_q4k_q8_1_x4_paired_block_lane(%bounded_input_size, %weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %q4_block, %block_lane) : (index, buffer, offset, buffer, offset, index, index) -> (f32) + scf.yield %pair : f32 + } else { + scf.yield %c0_f32 : f32 + } + %next = scalar.addf %iteration_acc, %contribution : f32 + scf.yield %next : f32 + } + func.return %sum : f32 +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/q4_q8_1_x4_dense.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/q4_q8_1_x4_dense.loom new file mode 100644 index 000000000000..4acae68782f2 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/q4_q8_1_x4_dense.loom @@ -0,0 +1,75 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @qwen3_moe.dense.q4k_q8_1_x4.body(%publish_output: i1, %token_count: index, %token0: index, %q8_input: buffer, %weight: buffer, %output: buffer) + +config.decl @qwen3_moe.dense_quantized.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @qwen3_moe.dense_quantized.output_size : %value: index where [range(%value, 1, 262144)] + +config.decl @qwen3_moe.dense_quantized.output_accumulation : %value: index where [range(%value, 0, 1)] + +func.decl @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32) + +template.def<@qwen3_moe.dense.q4k_q8_1_x4.body> device @qwen3_moe_dense_linear_q4k_q8_1_x4_body(%publish_output: i1, %token_count: index, %token0: index, %q8_input: buffer, %weight: buffer, %output: buffer) { + %input_size = config.get @qwen3_moe.dense_quantized.input_size : index + %output_size = config.get @qwen3_moe.dense_quantized.output_size : index + %output_accumulation = config.get @qwen3_moe.dense_quantized.output_accumulation : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144)] : index + %channel_tile = kernel.workgroup.id : index + %subgroup0 = kernel.subgroup.id : index + %lane = kernel.subgroup.lane.id : index + %c1 = index.constant 1 : index + %c8 = index.constant 8 : index + %c128 = index.constant 128 : index + %c144_bytes = index.constant 144 : offset + %c256 = index.constant 256 : index + %c0_i32 = scalar.constant 0 : i32 + %c0_f32 = scalar.constant 0.0 : f32 + %c0_offset = index.constant 0 : offset + %token, %launch_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %channel_base = index.mul %channel_tile, %c8 : index + %channel = index.add %channel_base, %subgroup : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %lane_i32 = index.cast %lane : index to i32 + %is_lane_zero = scalar.cmpi eq, %lane_i32, %c0_i32 : i32 + %q4_block_count = index.div %bounded_input_size, %c256 : index + %weight_row_bytes = index.scale %q4_block_count, %c144_bytes : index, offset -> offset + %weight_row_byte_base = index.scale %channel, %weight_row_bytes : index, offset -> offset + %q8_group_count = index.div %bounded_input_size, %c128 : index + %q8_row_bytes = index.scale %q8_group_count, %c144_bytes : index, offset -> offset + %q8_row_byte_base = index.scale %token, %q8_row_bytes : index, offset -> offset + %q8_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %q8_input, %weight, %output : buffer, buffer, buffer + %lane_sum = scf.if %publish_output -> (f32) { + %channel_sum = scf.if %valid_channel -> (f32) { + %sum = func.call @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%bounded_input_size, %weight_noalias, %weight_row_byte_base, %q8_noalias, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) + scf.yield %sum : f32 + } else { + scf.yield %c0_f32 : f32 + } + scf.yield %channel_sum : f32 + } else { + scf.yield %c0_f32 : f32 + } + %dot = kernel.subgroup.reduce %lane_sum : f32 + %accumulates_output = index.cmp eq, %output_accumulation, %c1 : index + scf.if %publish_output { + scf.if %valid_channel { + scf.if %is_lane_zero { + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%bounded_output_size]xf32> + %result = scf.if %accumulates_output -> (f32) { + %residual = view.load %output_view[%token, %channel] : view<[%launch_token_count]x[%bounded_output_size]xf32> -> f32 + %sum = scalar.addf %residual, %dot : f32 + scf.yield %sum : f32 + } else { + scf.yield %dot : f32 + } + view.store %result, %output_view[%token, %channel] : f32, view<[%launch_token_count]x[%bounded_output_size]xf32> + } + } + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom new file mode 100644 index 000000000000..49df978e7e09 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom @@ -0,0 +1,204 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.linear_q6k_q8_1_x4.body(%publish_output: i1, %token_count: index, %token0: index, %input_size: index, %output_size: index, %q8_input: buffer, %weight: buffer) -> (f32, index, index, i1, index, index) + +func.decl @ggml_q8_1_x4_word(%q8_input: buffer, %row_byte_base: offset, %q8_block: index, %word_in_block: index) -> (vector<4xi8>, f32) + +func.decl @ggml_q6k_f16_vector4(%weight: buffer, %weight_row_byte_base: offset, %q6_block: index, %q6_group: index, %packet: index) -> (vector<4xf16>) + +func.def inline @ggml_q6k_sign_extend_dot4(%code: vector<1xi32>) -> (vector<4xi8>) { + %c1_i32v = vector.constant 1 : vector<1xi32> + %c2_i32v = vector.constant 2 : vector<1xi32> + %low5_mask = vector.constant 522133279 : vector<1xi32> + %bit5_mask = vector.constant 538976288 : vector<1xi32> + %sign_mask = vector.constant -522133280 : vector<1xi32> + %low5 = vector.andi %code, %low5_mask : vector<1xi32> + %bit5 = vector.andi %code, %bit5_mask : vector<1xi32> + %bit6 = vector.shli %bit5, %c1_i32v : vector<1xi32> + %bit7 = vector.shli %bit5, %c2_i32v : vector<1xi32> + %high01 = vector.ori %bit5, %bit6 : vector<1xi32> + %high = vector.ori %high01, %bit7 : vector<1xi32> + %sign = vector.xori %high, %sign_mask : vector<1xi32> + %signed_i32 = vector.ori %low5, %sign : vector<1xi32> + %signed = vector.bitcast %signed_i32 : vector<1xi32> to vector<4xi8> + func.return %signed : vector<4xi8> +} + +func.def inline @ggml_q6k_q8_1_x4_block_lane(%weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %q6_block: index, %lane: index) -> (f32) { + %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 + %block_bytes = index.constant 210 : offset + %qh_byte_add = index.constant 128 : offset + %scale_byte_add = index.constant 192 : offset + %d_byte_add = index.constant 208 : offset + %c4_i32v = vector.constant 4 : vector<1xi32> + %c0_i32v = vector.constant 0 : vector<1xi32> + %nibble_mask = vector.constant 252645135 : vector<1xi32> + %high_mask = vector.constant 808464432 : vector<1xi32> + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_lane = index.assume %lane [range(%lane, 0, 31)] : index + %block_byte_add = index.scale %q6_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %weight_row_byte_base, %block_byte_add : offset + %qh_byte_base = index.add %block_byte_base, %qh_byte_add : offset + %scale_byte_base = index.add %block_byte_base, %scale_byte_add : offset + %d_byte_base = index.add %block_byte_base, %d_byte_add : offset + %ql_view = buffer.view %weight[%block_byte_base] : buffer -> view<32xi32> + %qh_view = buffer.view %weight[%qh_byte_base] : buffer -> view<16xi32> + %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<16xi8> + %d_view = buffer.view %weight[%d_byte_base] : buffer -> view<1xf16> + %lane_mod8 = index.rem %bounded_lane, %c8 : index + %lane_mod16 = index.rem %bounded_lane, %c16 : index + %lane_div16 = index.div %bounded_lane, %c16 : index + %lane_div8_in_16 = index.div %lane_mod16, %c8 : index + %lane_div4_in_16 = index.div %lane_mod16, %c4 : index + %qh_high_base = index.mul %lane_div16, %c8 : index + %qh_index0 = index.add %qh_high_base, %lane_mod8 : index + %qh_index = index.assume %qh_index0 [range(%qh_index0, 0, 15)] : index + %ql_word = vector.load %ql_view[%bounded_lane] : view<32xi32> -> vector<1xi32> + %qh_word = vector.load %qh_view[%qh_index] : view<16xi32> -> vector<1xi32> + %qh_base_shift_index = index.mul %lane_div8_in_16, %c2 : index + %qh_base_shift_i32 = index.cast %qh_base_shift_index : index to i32 + %q8_block_base = index.mul %q6_block, %c8 : index + %q8_high_add = index.mul %lane_div16, %c4 : index + %q8_quadrant = index.add %q8_high_add, %lane_div8_in_16 : index + %scale_high_base = index.mul %lane_div16, %c8 : index + %scale_lane0 = index.add %scale_high_base, %lane_div4_in_16 : index + %d_f16 = view.load %d_view[0] : view<1xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %sum = scf.for %part = [%c0 to %c2 step %c1](%accumulator = %c0_f32 : f32) -> (f32) unroll { + %bounded_part = index.assume %part [range(%part, 0, 1)] : index + %part_shift_index = index.mul %bounded_part, %c4 : index + %part_shift_i32 = index.cast %part_shift_index : index to i32 + %part_shift = vector.splat %part_shift_i32 : vector<1xi32> + %ql_shifted = vector.shrui %ql_word, %part_shift : vector<1xi32> + %ql = vector.andi %ql_shifted, %nibble_mask : vector<1xi32> + %qh_shift_i32 = scalar.addi %qh_base_shift_i32, %part_shift_i32 : i32 + %qh_shift = vector.splat %qh_shift_i32 : vector<1xi32> + %qh_shifted = vector.shrui %qh_word, %qh_shift : vector<1xi32> + %qh_positioned = vector.shli %qh_shifted, %c4_i32v : vector<1xi32> + %qh = vector.andi %qh_positioned, %high_mask : vector<1xi32> + %code = vector.ori %ql, %qh : vector<1xi32> + %signed_weight = func.call @ggml_q6k_sign_extend_dot4(%code) : (vector<1xi32>) -> (vector<4xi8>) + %q8_part_add = index.mul %bounded_part, %c2 : index + %q8_block_part = index.add %q8_block_base, %q8_part_add : index + %q8_block = index.add %q8_block_part, %q8_quadrant : index + %q8_values, %q8_d = func.call @ggml_q8_1_x4_word(%q8_input, %q8_row_byte_base, %q8_block, %lane_mod8) : (buffer, offset, index, index) -> (vector<4xi8>, f32) + %scale_lane = index.add %scale_lane0, %part_shift_index : index + %scale_i8 = view.load %scale_view[%scale_lane] : view<16xi8> -> i8 + %scale = scalar.sitofp %scale_i8 : i8 to f32 + %dot = vector.dot4i %signed_weight, %q8_values, %c0_i32v : vector<4xi8>, vector<4xi8>, vector<1xi32> + %dot_i32 = vector.extract %dot[0] : vector<1xi32> -> i32 + %dot_f32 = scalar.sitofp %dot_i32 : i32 to f32 + %scaled0 = scalar.mulf %dot_f32, %scale : f32 + %scaled1 = scalar.mulf %scaled0, %d : f32 + %contribution = scalar.mulf %scaled1, %q8_d : f32 + %next = scalar.addf %accumulator, %contribution : f32 + scf.yield %next : f32 + } + func.return %sum : f32 +} + +// Computes one lane's partial for a complete Q6_K row. Callers choose how +// rows and activations are tiled, then reduce the returned value by subgroup. +func.def inline @ggml_q6k_q8_1_x4_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + %c0_f32 = scalar.constant 0.0 : f32 + %block_count = index.div %bounded_input_size, %c256 : index + %result = scf.for %block = [%c0 to %block_count step %c1](%block_acc = %c0_f32 : f32) -> (f32) { + %contribution = func.call @ggml_q6k_q8_1_x4_block_lane(%weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %block, %lane) : (buffer, offset, buffer, offset, index, index) -> (f32) + %next = scalar.addf %block_acc, %contribution : f32 + scf.yield %next : f32 + } + func.return %result : f32 +} + +// Maps four Q6_K blocks across the four 16-lane partitions of a wave64. Each +// physical lane contracts the pair of virtual wave32 lanes that own the same +// packed positions in the low and high halves of one block. The subgroup +// reduction therefore combines four complete blocks per loop iteration while +// preserving the canonical wave32 unpacking and dot-product primitive. +func.def inline @ggml_q6k_q8_1_x4_row_lane_wave64_block4(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_lane = index.assume %lane [range(%lane, 0, 63)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c16 = index.constant 16 : index + %c256 = index.constant 256 : index + %c0_f32 = scalar.constant 0.0 : f32 + %block_count = index.div %bounded_input_size, %c256 : index + %padded_block_count = index.add %block_count, %c3 : index + %iteration_count = index.div %padded_block_count, %c4 : index + %block_in_iteration = index.div %bounded_lane, %c16 : index + %virtual_lane0 = index.rem %bounded_lane, %c16 : index + %virtual_lane1 = index.add %virtual_lane0, %c16 : index + %result = scf.for %iteration = [%c0 to %iteration_count step %c1](%row_acc = %c0_f32 : f32) -> (f32) { + %block_base = index.mul %iteration, %c4 : index + %block0 = index.add %block_base, %block_in_iteration : index + %valid_block = index.cmp ult, %block0, %block_count : index + %block_acc = scf.if %valid_block -> (f32) { + %low = func.call @ggml_q6k_q8_1_x4_block_lane(%weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %block0, %virtual_lane0) : (buffer, offset, buffer, offset, index, index) -> (f32) + %high = func.call @ggml_q6k_q8_1_x4_block_lane(%weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %block0, %virtual_lane1) : (buffer, offset, buffer, offset, index, index) -> (f32) + %sum = scalar.addf %low, %high : f32 + scf.yield %sum : f32 + } else { + scf.yield %c0_f32 : f32 + } + %next = scalar.addf %row_acc, %block_acc : f32 + scf.yield %next : f32 + } + func.return %result : f32 +} + +// Contracts one output channel for the calling Q6_K projection kernel. The +// caller owns publication so the same canonical contraction can feed either a +// dense logits tensor or an endpoint reduction. +template.def<@ggml.linear_q6k_q8_1_x4.body> device @ggml_linear_q6k_q8_1_x4_body(%publish_output: i1, %token_count: index, %token0: index, %input_size: index, %output_size: index, %q8_input: buffer, %weight: buffer) -> (f32, index, index, i1, index, index) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144)] : index + %channel_tile = kernel.workgroup.id : index + %subgroup0 = kernel.subgroup.id : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c8 = index.constant 8 : index + %c128 = index.constant 128 : index + %c210_bytes = index.constant 210 : offset + %c256 = index.constant 256 : index + %c144_bytes = index.constant 144 : offset + %token, %launch_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %channel_base = index.mul %channel_tile, %c8 : index + %channel = index.add %channel_base, %subgroup : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %q6_block_count = index.div %bounded_input_size, %c256 : index + %weight_row_bytes = index.scale %q6_block_count, %c210_bytes : index, offset -> offset + %weight_row_byte_base = index.scale %channel, %weight_row_bytes : index, offset -> offset + %q8_group_count = index.div %bounded_input_size, %c128 : index + %q8_row_bytes = index.scale %q8_group_count, %c144_bytes : index, offset -> offset + %q8_row_byte_base = index.scale %token, %q8_row_bytes : index, offset -> offset + %lane_acc = scf.if %publish_output -> (f32) { + %channel_acc = scf.if %valid_channel -> (f32) { + %value = func.call @ggml_q6k_q8_1_x4_row_lane(%bounded_input_size, %weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) + scf.yield %value : f32 + } else { + %c0_f32 = scalar.constant 0.0 : f32 + scf.yield %c0_f32 : f32 + } + scf.yield %channel_acc : f32 + } else { + %c0_f32 = scalar.constant 0.0 : f32 + scf.yield %c0_f32 : f32 + } + %dot = kernel.subgroup.reduce %lane_acc : f32 + template.return %dot, %token, %channel, %valid_channel, %launch_token_count, %bounded_output_size : f32, index, index, i1, index, index +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm.loom index ae8cd9174e87..2c75f94808c7 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm.loom @@ -182,7 +182,7 @@ config.decl @ggml.rmsnorm_mul_rope.section0 : %value: index where [range(%value, config.decl @ggml.rmsnorm_mul_rope.section1 : %value: index where [range(%value, 0, 65536)] config.decl @ggml.rmsnorm_mul_rope.section2 : %value: index where [range(%value, 0, 65536)] config.decl @ggml.rmsnorm_mul_rope.section3 : %value: index where [range(%value, 0, 65536)] -config.decl @ggml.rmsnorm_mul_rope.mode : %value: index where [range(%value, 8, 40)] +config.decl @ggml.rmsnorm_mul_rope.mode : %value: index where [range(%value, 2, 40)] config.decl @ggml.rmsnorm_mul_rope.workgroup_size : %value: index where [range(%value, 32, 1024), mul(%value, 32)] config.decl @ggml.rmsnorm_mul_rope.epsilon : f32 config.decl @ggml.rmsnorm_mul_rope.freq_base : f32 @@ -195,6 +195,7 @@ template.def<@ggml.rmsnorm_mul_rope_f32.body> device @ggml_rmsnorm_mul_rope_f32_ %c1 = index.constant 1 : index %c2 = index.constant 2 : index %c3 = index.constant 3 : index + %cneox = index.constant 2 : index %c40 = index.constant 40 : index %c0_f32 = scalar.constant 0.0 : f32 %cneg2_f32 = scalar.constant -2.0 : f32 @@ -308,7 +309,9 @@ template.def<@ggml.rmsnorm_mul_rope_f32.body> device @ggml_rmsnorm_mul_rope_f32_ %plane = scf.select %is_im, %plane_im, %plane_mr : index %plane_off = index.mul %plane, %ne2_b : index - %pos_i0 = index.add %i2, %plane_off : index + %sectioned_pos_i = index.add %i2, %plane_off : index + %is_neox = index.cmp eq, %mode, %cneox : index + %pos_i0 = scf.select %is_neox, %i2, %sectioned_pos_i : index %pos_i = index.assume %pos_i0 [range(%pos_i0, 0, 1073741823)] : index %pos_i32 = view.load %pos_view[%pos_i] : view<1073741824xi32> -> i32 %pos_f = scalar.sitofp %pos_i32 : i32 to f32 diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm_binary.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm_binary.loom index c7837a6696f6..f3ee58daa77f 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm_binary.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm_binary.loom @@ -2,20 +2,31 @@ // SPDX-License-Identifier: Apache-2.0 template.decl @ggml.rmsnorm_f32.apply(%arg0: f32, %arg1: f32) -> (f32) + template.decl @ggml.rmsnorm_f32.row_scale(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: f32, %arg5: buffer) -> (f32, index, index, i1) + template.decl @ggml.binary_f32.apply(%arg0: index, %arg1: f32, %arg2: f32) -> (f32) + template.decl @ggml.normalization.publish_f16.k16_vector4(%values: vector<4xf32>, %channel: index, %token: index, %token_count: index, %hidden_size: index, %packet: index, %output: buffer) + template.decl @ggml.normalization.publish_q8.binary_vector4(%publish_word: i1, %token_output_byte_base: offset, %channel: index, %values: vector<4xf32>, %scratch_values: view<256xf32>, %scratch_d: view<32xf32>, %output: buffer) config.decl @ggml.rmsnorm_binary_f32.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] + config.decl @ggml.rmsnorm_binary_f32.rms_epsilon : f32 + config.decl @ggml.rmsnorm_binary_f32.op : %value: index where [range(%value, 0, 3)] + config.decl @ggml.rmsnorm_binary_q8_1_x4.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] + config.decl @ggml.rmsnorm_binary_q8_1_x4.rms_epsilon : f32 + config.decl @ggml.rmsnorm_binary_q8_1_x4.op : %value: index where [range(%value, 0, 3)] template.decl @ggml.rmsnorm_binary_f32.body(%publish_f16: i1, %token_count: index, %input: buffer, %rhs: buffer, %output: buffer, %f16_output: buffer) +template.decl @ggml.rmsnorm_binary_strided_f32.body(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer) + template.def<@ggml.rmsnorm_binary_f32.body> device @ggml_rmsnorm_binary_f32_body(%publish_f16: i1, %token_count: index, %input: buffer, %rhs: buffer, %output: buffer, %f16_output: buffer) { %hidden_size0 = config.get @ggml.rmsnorm_binary_f32.hidden_size : index %epsilon = config.get @ggml.rmsnorm_binary_f32.rms_epsilon : f32 @@ -48,6 +59,36 @@ template.def<@ggml.rmsnorm_binary_f32.body> device @ggml_rmsnorm_binary_f32_body template.return } +config.decl @ggml.rmsnorm_binary_f32.input_stride : %value: index where [range(%value, 128, 1048576)] + +template.def<@ggml.rmsnorm_binary_strided_f32.body> device @ggml_rmsnorm_binary_strided_f32_body(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer) { + %hidden_size0 = config.get @ggml.rmsnorm_binary_f32.hidden_size : index + %input_stride0 = config.get @ggml.rmsnorm_binary_f32.input_stride : index + %epsilon = config.get @ggml.rmsnorm_binary_f32.rms_epsilon : f32 + %op = config.get @ggml.rmsnorm_binary_f32.op : index + %hidden_size1 = index.assume %hidden_size0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128)] : index + %input_stride1 = index.assume %input_stride0 [range(%input_stride0, 128, 1048576)] : index + %hidden_size, %input_stride = index.assume %hidden_size1, %input_stride1 [le(%hidden_size1, %input_stride1)] : index, index + %token0 = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %c0_offset = index.constant 0 : offset + %c256 = index.constant 256 : index + %input_noalias, %rhs_noalias, %output_noalias = buffer.assume.noalias %input, %rhs, %output : buffer, buffer, buffer + %row_scale, %token, %launch_token_count, %valid_token = template.apply<@ggml.rmsnorm_f32.row_scale>(%token_count, %token0, %hidden_size, %input_stride, %epsilon, %input_noalias) : (index, index, index, index, f32, buffer) -> (f32, index, index, i1) + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%input_stride]xf32> + %rhs_view = buffer.view %rhs_noalias[%c0_offset] : buffer -> view<[%hidden_size]xf32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + scf.if %valid_token { + scf.for %channel = [%workitem to %hidden_size step %c256] { + %value = view.load %input_view[%token, %channel] : view<[%launch_token_count]x[%input_stride]xf32> -> f32 + %rhs_value = view.load %rhs_view[%channel] : view<[%hidden_size]xf32> -> f32 + %normalized = template.apply<@ggml.rmsnorm_f32.apply>(%value, %row_scale) : (f32, f32) -> (f32) + %result = template.apply<@ggml.binary_f32.apply>(%op, %normalized, %rhs_value) : (index, f32, f32) -> (f32) + view.store %result, %output_view[%token, %channel] : f32, view<[%launch_token_count]x[%hidden_size]xf32> + } + } + template.return +} template.decl @ggml.rmsnorm_binary_f32_k16.body(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer, %f16_output: buffer) @@ -241,7 +282,6 @@ template.def<@ggml.rmsnorm_binary_f32_k16.body> device @ggml_rmsnorm_binary_f32_ template.return } - template.decl @ggml.rmsnorm_binary_q8_1_x4.body(%publish_f16: i1, %token_count: index, %input: buffer, %rhs: buffer, %output: buffer, %q8_output: buffer, %f16_output: buffer) template.def<@ggml.rmsnorm_binary_q8_1_x4.body> device @ggml_rmsnorm_binary_q8_1_x4_body(%publish_f16: i1, %token_count: index, %input: buffer, %rhs: buffer, %output: buffer, %q8_output: buffer, %f16_output: buffer) { diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm_q8_1_x4.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm_q8_1_x4.loom new file mode 100644 index 000000000000..c3b9c47a181d --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm_q8_1_x4.loom @@ -0,0 +1,182 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @qwen3_moe.rmsnorm_quantize_q8_1_x4.body(%publish_normalized: i1, %reduction_subgroup_count0: index, %token_count: index, %token0: index, %input: buffer, %weight: buffer, %normalized_output: buffer, %q8_output: buffer) + +config.decl @qwen3_moe.model.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] + +config.decl @qwen3_moe.model.rms_epsilon : f32 + +// Shared RMSNorm and GGML Q8_1 x4 row producer. The caller owns launch geometry +// and supplies the logical token whose complete hidden row this workgroup owns. +// The attention export discards the ordinary F32 row, while feed-forward +// publishes it for the router without rereading and requantizing the normalized +// values. +template.def<@qwen3_moe.rmsnorm_quantize_q8_1_x4.body> device @qwen3_moe_rmsnorm_quantize_q8_1_x4_body(%publish_normalized: i1, %reduction_subgroup_count0: index, %token_count: index, %token0: index, %input: buffer, %weight: buffer, %normalized_output: buffer, %q8_output: buffer) { + %hidden_size0 = config.get @qwen3_moe.model.hidden_size : index + %epsilon = config.get @qwen3_moe.model.rms_epsilon : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %reduction_subgroup_count = index.assume %reduction_subgroup_count0 [range(%reduction_subgroup_count0, 1, 8)] : index + %hidden_size = index.assume %hidden_size0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128)] : index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %lane = kernel.subgroup.lane.id : index + %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 + %c32 = index.constant 32 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %c1024 = index.constant 1024 : index + %group_bytes = index.constant 144 : offset + %payload_byte_add = index.constant 16 : offset + %scratch_d_byte_add = index.constant 1024 : offset + %scratch_bytes = index.constant 1152 : offset + %c0_offset = index.constant 0 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %c1_f32 = scalar.constant 1.0 : f32 + %c127 = scalar.constant 127.0 : f32 + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %valid_token = index.cmp ult, %token0, %bounded_token_count : index + %safe_token0 = scf.select %valid_token, %token0, %c0 : index + %token, %launch_token_count = index.assume %safe_token0, %bounded_token_count [lt(%safe_token0, %bounded_token_count)] : index, index + %hidden_size_i32 = index.cast %hidden_size : index to i32 + %hidden_size_f32 = scalar.sitofp %hidden_size_i32 : i32 to f32 + %physical_group_count = index.div %hidden_size, %c128 : index + %row_bytes = index.scale %physical_group_count, %group_bytes : index, offset -> offset + %token_output_byte_base = index.scale %token, %row_bytes : index, offset -> offset + %input_noalias, %weight_noalias, %q8_output_noalias = buffer.assume.noalias %input, %weight, %q8_output : buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + %weight_view = buffer.view %weight_noalias[%c0_offset] : buffer -> view<[%hidden_size]xf32> + %normalized_output_view = buffer.view %normalized_output[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_values = buffer.view %scratch[%c0_offset] : buffer -> view<256xf32> + %scratch_d = buffer.view %scratch[%scratch_d_byte_add] : buffer -> view<32xf32> + // Reduce the complete row before any block-local quantization. + %thread_sum = scf.for %channel = [%workitem to %hidden_size step %c256](%running_sum = %c0_f32 : f32) -> (f32) { + %value = view.load %input_view[%token, %channel] : view<[%launch_token_count]x[%hidden_size]xf32> -> f32 + %square = scalar.mulf %value, %value : f32 + %next_sum = scalar.addf %running_sum, %square : f32 + scf.yield %next_sum : f32 + } + %subgroup_sum = kernel.subgroup.reduce %thread_sum : f32 + %is_subgroup_leader = index.cmp eq, %lane, %c0 : index + scf.if %is_subgroup_leader { + view.store %subgroup_sum, %scratch_values[%subgroup] : f32, view<256xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %is_reduction_subgroup = index.cmp eq, %subgroup, %c0 : index + %is_reduction_lane = index.cmp ult, %lane, %reduction_subgroup_count : index + %loads_subgroup_sum = scalar.andi %is_reduction_subgroup, %is_reduction_lane : i1 + %subgroup_partial = scf.if %loads_subgroup_sum -> (f32) { + %value = view.load %scratch_values[%lane] : view<256xf32> -> f32 + scf.yield %value : f32 + } else { + scf.yield %c0_f32 : f32 + } + %row_sum = kernel.subgroup.reduce %subgroup_partial : f32 + %writes_scale = scalar.andi %is_reduction_subgroup, %is_subgroup_leader : i1 + scf.if %writes_scale { + %mean = scalar.divf %row_sum, %hidden_size_f32 : f32 + %biased_mean = scalar.addf %mean, %epsilon : f32 + %scale = scalar.rsqrtf %biased_mean : f32 + view.store %scale, %scratch_values[%c0] : f32, view<256xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %row_scale = view.load %scratch_values[%c0] : view<256xf32> -> f32 + %row_scale_vector = vector.splat %row_scale : vector<4xf32> + %publishes_normalized = scalar.andi %publish_normalized, %valid_token : i1 + // Reuse one scratch frame for each 1024-element stripe. Hidden sizes need + // only be divisible by 128; inactive workitems in the final stripe carry + // zeros and never publish. + scf.for %stripe_base = [%c0 to %hidden_size step %c1024] { + %word_element_add = index.mul %workitem, %c4 : index + %channel = index.add %stripe_base, %word_element_add : index + %valid_word = index.cmp ult, %channel, %hidden_size : index + %mask = vector.mask.range [%channel to %hidden_size step %c1] : index -> vector<4xi1> + %input_values = vector.load.mask %input_view[%token, %channel], %mask, %c0_f32x4 : view<[%launch_token_count]x[%hidden_size]xf32>, vector<4xi1>, vector<4xf32> + %learned_weights = vector.load.mask %weight_view[%channel], %mask, %c0_f32x4 : view<[%hidden_size]xf32>, vector<4xi1>, vector<4xf32> + %normalized0 = vector.mulf %input_values, %row_scale_vector : vector<4xf32> + %normalized = vector.mulf %normalized0, %learned_weights : vector<4xf32> + scf.if %publishes_normalized { + vector.store.mask %normalized, %normalized_output_view[%token, %channel], %mask : vector<4xf32>, view<[%launch_token_count]x[%hidden_size]xf32>, vector<4xi1> + } + %absolute_values = vector.absf %normalized : vector<4xf32> + %thread_max = vector.reduce %absolute_values, %c0_f32 : vector<4xf32>, f32 + view.store %thread_max, %scratch_values[%workitem] : f32, view<256xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + %word_in_block = index.rem %workitem, %c8 : index + %block_in_stripe = index.div %workitem, %c8 : index + %is_block_leader = index.cmp eq, %word_in_block, %c0 : index + %writes_block_d = scalar.andi %valid_word, %is_block_leader : i1 + scf.if %writes_block_d { + %cohort_base = index.mul %block_in_stripe, %c8 : index + %cohort_maxima = vector.load %scratch_values[%cohort_base] : view<256xf32> -> vector<8xf32> + %amax = vector.reduce %cohort_maxima, %c0_f32 : vector<8xf32>, f32 + %d = scalar.divf %amax, %c127 : f32 + view.store %d, %scratch_d[%block_in_stripe] : f32, view<32xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %d = scf.if %valid_word -> (f32) { + %block_d = view.load %scratch_d[%block_in_stripe] : view<32xf32> -> f32 + scf.yield %block_d : f32 + } else { + scf.yield %c0_f32 : f32 + } + %d_nonzero = scalar.cmpf one, %d, %c0_f32 : f32 + %d_inverse = scf.if %d_nonzero -> (f32) { + %inverse = scalar.divf %c1_f32, %d : f32 + scf.yield %inverse : f32 + } else { + scf.yield %c0_f32 : f32 + } + %d_inverse_vector = vector.splat %d_inverse : vector<4xf32> + %scaled_values = vector.mulf %normalized, %d_inverse_vector : vector<4xf32> + %rounded_values = vector.roundf %scaled_values : vector<4xf32> + %quantized_values = vector.fptosi %rounded_values : vector<4xf32> to vector<4xi8> + %packed_word = vector.bitcast %quantized_values : vector<4xi8> to vector<1xi32> + %publishes_q8_word = scalar.andi %valid_word, %valid_token : i1 + scf.if %publishes_q8_word { + %q8_block = index.div %channel, %c32 : index + %physical_group = index.div %q8_block, %c4 : index + %block_in_group = index.rem %q8_block, %c4 : index + %group_byte_add = index.scale %physical_group, %group_bytes : index, offset -> offset + %group_byte_offset = index.add %token_output_byte_base, %group_byte_add : offset + %payload_byte_offset = index.add %group_byte_offset, %payload_byte_add : offset + %group_ds = buffer.view %q8_output_noalias[%group_byte_offset] : buffer -> view<8xf16> + %group_qs = buffer.view %q8_output_noalias[%payload_byte_offset] : buffer -> view<32xi32> + %block_word_base = index.mul %block_in_group, %c8 : index + %packed_word_index0 = index.add %block_word_base, %word_in_block : index + %packed_word_index = index.assume %packed_word_index0 [range(%packed_word_index0, 0, 31)] : index + vector.store %packed_word, %group_qs[%packed_word_index] : vector<1xi32>, view<32xi32> + } + %thread_quantized_sum = vector.reduce %rounded_values, %c0_f32 : vector<4xf32>, f32 + view.store %thread_quantized_sum, %scratch_values[%workitem] : f32, view<256xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + %publishes_block_ds = scalar.andi %writes_block_d, %valid_token : i1 + scf.if %publishes_block_ds { + %cohort_base = index.mul %block_in_stripe, %c8 : index + %cohort_sums = vector.load %scratch_values[%cohort_base] : view<256xf32> -> vector<8xf32> + %quantized_sum = vector.reduce %cohort_sums, %c0_f32 : vector<8xf32>, f32 + %s = scalar.mulf %quantized_sum, %d : f32 + %q8_block = index.div %channel, %c32 : index + %physical_group = index.div %q8_block, %c4 : index + %block_in_group = index.rem %q8_block, %c4 : index + %group_byte_add = index.scale %physical_group, %group_bytes : index, offset -> offset + %group_byte_offset = index.add %token_output_byte_base, %group_byte_add : offset + %group_ds = buffer.view %q8_output_noalias[%group_byte_offset] : buffer -> view<8xf16> + %d_f16 = scalar.fptrunc %d : f32 to f16 + %s_f16 = scalar.fptrunc %s : f32 to f16 + %ds_index = index.mul %block_in_group, %c2 : index + %s_index = index.add %ds_index, %c1 : index + view.store %d_f16, %group_ds[%ds_index] : f16, view<8xf16> + view.store %s_f16, %group_ds[%s_index] : f16, view<8xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/quantize_q8_1_x4_group.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/quantize_q8_1_x4_group.loom new file mode 100644 index 000000000000..02f6eff75c9e --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/quantize_q8_1_x4_group.loom @@ -0,0 +1,157 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.quantize_q8_1_x4.group_body(%publish_output: i1, %group_count0: index, %group0: index, %input: buffer, %output: buffer) + +config.decl @ggml.quantize_q8_1_x4.group_capacity : %value: index where [range(%value, 1, 524288)] + +// Loads one logical 32-element Q8_1 block from the four-way physical packing. +// `row_byte_base` addresses the first x4 group for one activation row. +func.def inline @ggml_q8_1_x4_block(%q8_input: buffer, %row_byte_base: offset, %q8_block: index) -> (vector<32xi8>, f32, f32) { + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %group_bytes = index.constant 144 : offset + %payload_byte_add = index.constant 16 : offset + %group = index.div %q8_block, %c4 : index + %inner0 = index.rem %q8_block, %c4 : index + %inner = index.assume %inner0 [range(%inner0, 0, 3)] : index + %group_byte_add = index.scale %group, %group_bytes : index, offset -> offset + %group_byte_base = index.add %row_byte_base, %group_byte_add : offset + %payload_byte_base = index.add %group_byte_base, %payload_byte_add : offset + %ds_view = buffer.view %q8_input[%group_byte_base] : buffer -> view<8xf16> + %payload_view = buffer.view %q8_input[%payload_byte_base] : buffer -> view<32xi32> + %d_index = index.mul %inner, %c2 : index + %s_index = index.add %d_index, %c1 : index + %word_index = index.mul %inner, %c8 : index + %d_f16 = view.load %ds_view[%d_index] : view<8xf16> -> f16 + %s_f16 = view.load %ds_view[%s_index] : view<8xf16> -> f16 + %words = vector.load %payload_view[%word_index] : view<32xi32> -> vector<8xi32> + %values = vector.bitcast %words : vector<8xi32> to vector<32xi8> + %d = scalar.extf %d_f16 : f16 to f32 + %s = scalar.extf %s_f16 : f16 to f32 + func.return %values, %d, %s : vector<32xi8>, f32, f32 +} + +// Loads one packed four-value word and its block scale. Dot-product kernels +// use this narrower form so a lane does not load the other seven words owned +// by its subgroup peers. +func.def inline @ggml_q8_1_x4_word(%q8_input: buffer, %row_byte_base: offset, %q8_block: index, %word_in_block: index) -> (vector<4xi8>, f32) { + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %group_bytes = index.constant 144 : offset + %payload_byte_add = index.constant 16 : offset + %group = index.div %q8_block, %c4 : index + %inner0 = index.rem %q8_block, %c4 : index + %inner = index.assume %inner0 [range(%inner0, 0, 3)] : index + %word0 = index.assume %word_in_block [range(%word_in_block, 0, 7)] : index + %group_byte_add = index.scale %group, %group_bytes : index, offset -> offset + %group_byte_base = index.add %row_byte_base, %group_byte_add : offset + %payload_byte_base = index.add %group_byte_base, %payload_byte_add : offset + %ds_view = buffer.view %q8_input[%group_byte_base] : buffer -> view<8xf16> + %payload_view = buffer.view %q8_input[%payload_byte_base] : buffer -> view<32xi32> + %d_index = index.mul %inner, %c2 : index + %inner_word_base = index.mul %inner, %c8 : index + %word_index = index.add %inner_word_base, %word0 : index + %d_f16 = view.load %ds_view[%d_index] : view<8xf16> -> f16 + %packed = vector.load %payload_view[%word_index] : view<32xi32> -> vector<1xi32> + %values = vector.bitcast %packed : vector<1xi32> to vector<4xi8> + %d = scalar.extf %d_f16 : f16 to f32 + func.return %values, %d : vector<4xi8>, f32 +} + +// Packs one explicit physical group. Callers provide the complete physical +// group domain and an in-range ordinal, and guarantee that one complete +// 32-lane wave executes after all 128 source values are visible. The uniform +// publication predicate suppresses only destination writes so every lane still +// reaches the workgroup barriers. +template.def<@ggml.quantize_q8_1_x4.group_body> device @ggml_quantize_q8_1_x4_group_body(%publish_output: i1, %group_count0: index, %group0: index, %input: buffer, %output: buffer) { + %group_count, %group = index.assume %group_count0, %group0 [range(%group_count0, 1, 524288), lt(%group0, %group_count0)] : index, index + %lane = kernel.workitem.id : index + %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 + %c32 = index.constant 32 : index + %c128 = index.constant 128 : index + %group_bytes = index.constant 144 : offset + %payload_byte_add = index.constant 16 : offset + %scratch_d_byte_add = index.constant 128 : offset + %scratch_bytes = index.constant 144 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %c1_f32 = scalar.constant 1.0 : f32 + %c127 = scalar.constant 127.0 : f32 + %c0_offset = index.constant 0 : offset + %launched_element_count = index.mul %group_count, %c128 : index + %block_in_group0 = index.div %lane, %c8 : index + %block_in_group = index.assume %block_in_group0 [range(%block_in_group0, 0, 3)] : index + %word_in_block0 = index.rem %lane, %c8 : index + %word_in_block = index.assume %word_in_block0 [range(%word_in_block0, 0, 7)] : index + %group_element_base = index.mul %group, %c128 : index + %block_element_add = index.mul %block_in_group, %c32 : index + %block_word_add = index.mul %block_in_group, %c8 : index + %word_element_add = index.mul %word_in_block, %c4 : index + %input_block_base = index.add %group_element_base, %block_element_add : index + %input_index = index.add %input_block_base, %word_element_add : index + %input_noalias, %output_noalias = buffer.assume.noalias %input, %output : buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launched_element_count]xf32> + %input_values = vector.load %input_view[%input_index] : view<[%launched_element_count]xf32> -> vector<4xf32> + %absolute_values = vector.absf %input_values : vector<4xf32> + %thread_max = vector.reduce %absolute_values, %c0_f32 : vector<4xf32>, f32 + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_values = buffer.view %scratch[%c0_offset] : buffer -> view<32xf32> + %scratch_d = buffer.view %scratch[%scratch_d_byte_add] : buffer -> view<4xf32> + view.store %thread_max, %scratch_values[%lane] : f32, view<32xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + %is_cohort_leader = index.cmp eq, %word_in_block, %c0 : index + scf.if %is_cohort_leader { + %cohort_base = index.mul %block_in_group, %c8 : index + %cohort_maxima = vector.load %scratch_values[%cohort_base] : view<32xf32> -> vector<8xf32> + %amax = vector.reduce %cohort_maxima, %c0_f32 : vector<8xf32>, f32 + %d = scalar.divf %amax, %c127 : f32 + view.store %d, %scratch_d[%block_in_group] : f32, view<4xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %d = view.load %scratch_d[%block_in_group] : view<4xf32> -> f32 + %d_nonzero = scalar.cmpf one, %d, %c0_f32 : f32 + %d_inverse = scf.if %d_nonzero -> (f32) { + %inverse = scalar.divf %c1_f32, %d : f32 + scf.yield %inverse : f32 + } else { + scf.yield %c0_f32 : f32 + } + %d_inverse_vector = vector.splat %d_inverse : vector<4xf32> + %scaled_values = vector.mulf %input_values, %d_inverse_vector : vector<4xf32> + %rounded_values = vector.roundf %scaled_values : vector<4xf32> + %quantized_values = vector.fptosi %rounded_values : vector<4xf32> to vector<4xi8> + %packed_word = vector.bitcast %quantized_values : vector<4xi8> to vector<1xi32> + %group_byte_offset = index.scale %group, %group_bytes : index, offset -> offset + %payload_byte_offset = index.add %group_byte_offset, %payload_byte_add : offset + %group_ds = buffer.view %output_noalias[%group_byte_offset] : buffer -> view<8xf16> + %group_qs = buffer.view %output_noalias[%payload_byte_offset] : buffer -> view<32xi32> + %packed_word_index0 = index.add %block_word_add, %word_in_block : index + %packed_word_index = index.assume %packed_word_index0 [range(%packed_word_index0, 0, 31)] : index + scf.if %publish_output { + vector.store %packed_word, %group_qs[%packed_word_index] : vector<1xi32>, view<32xi32> + } + %thread_sum = vector.reduce %rounded_values, %c0_f32 : vector<4xf32>, f32 + view.store %thread_sum, %scratch_values[%lane] : f32, view<32xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + %publishes_metadata = scalar.andi %is_cohort_leader, %publish_output : i1 + scf.if %publishes_metadata { + %cohort_base = index.mul %block_in_group, %c8 : index + %cohort_sums = vector.load %scratch_values[%cohort_base] : view<32xf32> -> vector<8xf32> + %quantized_sum = vector.reduce %cohort_sums, %c0_f32 : vector<8xf32>, f32 + %s = scalar.mulf %quantized_sum, %d : f32 + %d_f16 = scalar.fptrunc %d : f32 to f16 + %s_f16 = scalar.fptrunc %s : f32 to f16 + %ds_index = index.mul %block_in_group, %c2 : index + view.store %d_f16, %group_ds[%ds_index] : f16, view<8xf16> + %s_index = index.add %ds_index, %c1 : index + view.store %s_f16, %group_ds[%s_index] : f16, view<8xf16> + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/rope/rope_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/rope/rope_f32.loom index 2b53e3da3848..b3e926e61b2d 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/rope/rope_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/rope/rope_f32.loom @@ -33,6 +33,12 @@ template.def<@ggml.rope_f32.rotate_scalar_pair> device @ggml_rope_f32_rotate_sca template.decl @ggml.rope_f32.body(%token_count: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %output: buffer) +template.decl @ggml.rope_scale_f32.body(%token_count: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %output: buffer) + +template.decl @ggml.rope_token_scale_f32.body(%token_count: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %scales: buffer, %output: buffer) + +template.decl @ggml.rope_f32.scaled_body(%token_count: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %output_scale: f32, %output: buffer) + template.decl @ggml.rope_set_rows_f32.body(%token_count: index, %cache_row_count: index, %positions: buffer, %indices: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %cache: buffer) template.decl @ggml.rope_f32.rotate_neox_packet(%position: f32, %input_base: index, %pair_channel: index, %half_n_dims: index, %bounded_input_span: index, %mscale: f32, %input: buffer, %theta_packet: vector<2xf32>, %freq_factors_packet: vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) @@ -65,6 +71,8 @@ config.decl @ggml.rope_f32.mscale : f32 config.decl @ggml.rope_f32.mode : %value: index where [range(%value, 0, 2)] +config.decl @ggml.rope_scale_f32.scale : f32 + config.decl @ggml.rope_set_rows_f32.head_size : %value: index where [range(%value, 4, 1024), mul(%value, 4)] config.decl @ggml.rope_set_rows_f32.n_dims : %value: index where [range(%value, 4, 1024), mul(%value, 4)] @@ -143,7 +151,7 @@ template.def<@ggml.rope_f32.copy_packet> device @ggml_rope_f32_copy_packet(%inpu template.return %values : vector<4xf32> } -template.def<@ggml.rope_f32.body> device @ggml_rope_f32_body(%token_count: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %output: buffer) { +template.def<@ggml.rope_f32.scaled_body> device @ggml_rope_f32_scaled_body(%token_count: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %output_scale: f32, %output: buffer) { %token_capacity = config.get @ggml.rope_f32.token_capacity : index %head_count0 = config.get @ggml.rope_f32.head_count : index %head_size0 = config.get @ggml.rope_f32.head_size : index @@ -165,6 +173,8 @@ template.def<@ggml.rope_f32.body> device @ggml_rope_f32_body(%token_count: index %c4 = index.constant 4 : index %mode_neox = index.constant 2 : index %c0_offset = index.constant 0 : offset + %output_scale2 = vector.splat %output_scale : vector<2xf32> + %output_scale4 = vector.splat %output_scale : vector<4xf32> %half_n_dims = index.div %n_dims, %c2 : index %pair_packet_count = index.div %head_size, %c4 : index %rotary_packet_count = index.div %n_dims, %c4 : index @@ -199,12 +209,15 @@ template.def<@ggml.rope_f32.body> device @ggml_rope_f32_body(%token_count: index scf.if %is_neox { %paired_channel = index.add %pair_channel, %half_n_dims : index %scaled_low, %scaled_high = template.apply<@ggml.rope_f32.rotate_neox_packet>(%position, %input_base, %pair_channel, %half_n_dims, %bounded_input_span, %mscale, %input_noalias, %theta_packet, %freq_factors_packet) : (f32, index, index, index, index, f32, buffer, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) - vector.store %scaled_low, %output_view[%token, %head, %pair_channel] : vector<2xf32>, view<[%bounded_token_count]x[%head_count]x[%head_size]xf32> - vector.store %scaled_high, %output_view[%token, %head, %paired_channel] : vector<2xf32>, view<[%bounded_token_count]x[%head_count]x[%head_size]xf32> + %output_low = vector.mulf %scaled_low, %output_scale2 : vector<2xf32> + %output_high = vector.mulf %scaled_high, %output_scale2 : vector<2xf32> + vector.store %output_low, %output_view[%token, %head, %pair_channel] : vector<2xf32>, view<[%bounded_token_count]x[%head_count]x[%head_size]xf32> + vector.store %output_high, %output_view[%token, %head, %paired_channel] : vector<2xf32>, view<[%bounded_token_count]x[%head_count]x[%head_size]xf32> } scf.if %is_normal { %rotated = template.apply<@ggml.rope_f32.rotate_normal_packet>(%position, %input_base, %normal_channel, %bounded_input_span, %mscale, %input_noalias, %theta_packet, %freq_factors_packet) : (f32, index, index, index, f32, buffer, vector<2xf32>, vector<2xf32>) -> (vector<4xf32>) - vector.store %rotated, %output_view[%token, %head, %normal_channel] : vector<4xf32>, view<[%bounded_token_count]x[%head_count]x[%head_size]xf32> + %output_rotated = vector.mulf %rotated, %output_scale4 : vector<4xf32> + vector.store %output_rotated, %output_view[%token, %head, %normal_channel] : vector<4xf32>, view<[%bounded_token_count]x[%head_count]x[%head_size]xf32> } } scf.if %publish_copy { @@ -215,11 +228,38 @@ template.def<@ggml.rope_f32.body> device @ggml_rope_f32_body(%token_count: index %input_head_offset = index.mul %head, %input_stride1 : index %input_base = index.add %input_token_offset, %input_head_offset : index %values = template.apply<@ggml.rope_f32.copy_packet>(%input_base, %normal_channel, %bounded_input_span, %input_noalias) : (index, index, index, buffer) -> (vector<4xf32>) - vector.store %values, %output_view[%token, %head, %normal_channel] : vector<4xf32>, view<[%bounded_token_count]x[%head_count]x[%head_size]xf32> + %output_values = vector.mulf %values, %output_scale4 : vector<4xf32> + vector.store %output_values, %output_view[%token, %head, %normal_channel] : vector<4xf32>, view<[%bounded_token_count]x[%head_count]x[%head_size]xf32> } template.return } +template.def<@ggml.rope_f32.body> device @ggml_rope_f32_body(%token_count: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %output: buffer) { + %one = scalar.constant 1.0 : f32 + template.apply<@ggml.rope_f32.scaled_body>(%token_count, %positions, %input, %theta, %freq_factors, %one, %output) : (index, buffer, buffer, buffer, buffer, f32, buffer) + template.return +} + +template.def<@ggml.rope_scale_f32.body> device @ggml_rope_scale_f32_body(%token_count: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %output: buffer) { + %output_scale = config.get @ggml.rope_scale_f32.scale : f32 + template.apply<@ggml.rope_f32.scaled_body>(%token_count, %positions, %input, %theta, %freq_factors, %output_scale, %output) : (index, buffer, buffer, buffer, buffer, f32, buffer) + template.return +} + +template.def<@ggml.rope_token_scale_f32.body> device @ggml_rope_token_scale_f32_body(%token_count: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %scales: buffer, %output: buffer) { + %token_capacity = config.get @ggml.rope_f32.token_capacity : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %token0 = kernel.workgroup.id : index + %c0 = index.constant 0 : index + %c0_offset = index.constant 0 : offset + %valid_token = index.cmp ult, %token0, %bounded_token_count : index + %token = scf.select %valid_token, %token0, %c0 : index + %scale_view = buffer.view %scales[%c0_offset] : buffer -> view<[%bounded_token_count]xf32> + %output_scale = view.load %scale_view[%token] : view<[%bounded_token_count]xf32> -> f32 + template.apply<@ggml.rope_f32.scaled_body>(%token_count, %positions, %input, %theta, %freq_factors, %output_scale, %output) : (index, buffer, buffer, buffer, buffer, f32, buffer) + template.return +} + template.def<@ggml.rope_set_rows_f32.body> device @ggml_rope_set_rows_f32_body(%token_count: index, %cache_row_count: index, %positions: buffer, %indices: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %cache: buffer) { %token_capacity = config.get @ggml.rope_set_rows_f32.token_capacity : index %head_count0 = config.get @ggml.rope_set_rows_f32.head_count : index diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/router_projection_fused.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/router_projection_fused.loom new file mode 100644 index 000000000000..1d95f9fae09b --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/router_projection_fused.loom @@ -0,0 +1,91 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @qwen3_moe.router.fused.experts_per_wave() -> (index) + +template.decl @qwen3_moe.router.fused.projection(%hidden_size: index, %expert_count: index, %token_count: index, %token: index, %expert_tile: index, %lane: index, %writes_projection: i1, %input: buffer, %weight: buffer, %logits: buffer) + +amdgpu.target @qwen3_moe_router_fused_gfx1151 + +template.def<@qwen3_moe.router.fused.experts_per_wave> target(@qwen3_moe_router_fused_gfx1151) requires [#target.subgroup.size<64>] priority(20) @qwen3_moe_router_fused_two_experts_per_wave() -> (index) { + %c2 = index.constant 2 : index + template.return %c2 : index +} + +template.def<@qwen3_moe.router.fused.experts_per_wave> requires [#target.subgroup.size<64>] priority(1) @qwen3_moe_router_fused_four_experts_per_wave() -> (index) { + %c4 = index.constant 4 : index + template.return %c4 : index +} + +template.def<@qwen3_moe.router.fused.projection> device target(@qwen3_moe_router_fused_gfx1151) requires [#target.subgroup.size<64>] priority(20) @qwen3_moe_router_fused_two_expert_projection(%hidden_size: index, %expert_count: index, %token_count: index, %token: index, %expert_tile: index, %lane: index, %writes_projection: i1, %input: buffer, %weight: buffer, %logits: buffer) { + %c0_offset = index.constant 0 : offset + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c256 = index.constant 256 : index + %c0_f32 = scalar.constant 0.0 : f32 + %expert_base = index.mul %expert_tile, %c2 : index + %expert1 = index.add %expert_base, %c1 : index + %lane_channel = index.mul %lane, %c4 : index + %input_view = buffer.view %input[%c0_offset] : buffer -> view<[%token_count]x[%hidden_size]xf32> + %weight_view = buffer.view %weight[%c0_offset] : buffer -> view<[%expert_count]x[%hidden_size]xf32> + %logits_view = buffer.view %logits[%c0_offset] : buffer -> view<[%token_count]x[%expert_count]xf32> + %c0_f32x2 = vector.splat %c0_f32 : vector<2xf32> + %lane_sums = scf.for %channel = [%lane_channel to %hidden_size step %c256](%accumulators = %c0_f32x2 : vector<2xf32>) -> (vector<2xf32>) unroll(%c4) schedule(interleaved) { + %input_values = vector.load %input_view[%token, %channel] : view<[%token_count]x[%hidden_size]xf32> -> vector<4xf32> + %weight_values0 = vector.load %weight_view[%expert_base, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> + %weight_values1 = vector.load %weight_view[%expert1, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> + %accumulator0 = vector.extract %accumulators[0] : vector<2xf32> -> f32 + %accumulator1 = vector.extract %accumulators[1] : vector<2xf32> -> f32 + %next0 = vector.dotf %input_values, %weight_values0, %accumulator0 : vector<4xf32>, vector<4xf32>, f32 + %next1 = vector.dotf %input_values, %weight_values1, %accumulator1 : vector<4xf32>, vector<4xf32>, f32 + %next_accumulators = vector.from_elements %next0, %next1 : vector<2xf32> + scf.yield %next_accumulators : vector<2xf32> + } + %sums = kernel.subgroup.reduce %lane_sums : vector<2xf32> + scf.if %writes_projection { + vector.store %sums, %logits_view[%token, %expert_base] : vector<2xf32>, view<[%token_count]x[%expert_count]xf32> + } + template.return +} + +template.def<@qwen3_moe.router.fused.projection> device requires [#target.subgroup.size<64>] priority(1) @qwen3_moe_router_fused_four_expert_projection(%hidden_size: index, %expert_count: index, %token_count: index, %token: index, %expert_tile: index, %lane: index, %writes_projection: i1, %input: buffer, %weight: buffer, %logits: buffer) { + %c0_offset = index.constant 0 : offset + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c256 = index.constant 256 : index + %c0_f32 = scalar.constant 0.0 : f32 + %expert_base = index.mul %expert_tile, %c4 : index + %expert1 = index.add %expert_base, %c1 : index + %expert2 = index.add %expert_base, %c2 : index + %expert3 = index.add %expert_base, %c3 : index + %lane_channel = index.mul %lane, %c4 : index + %input_view = buffer.view %input[%c0_offset] : buffer -> view<[%token_count]x[%hidden_size]xf32> + %weight_view = buffer.view %weight[%c0_offset] : buffer -> view<[%expert_count]x[%hidden_size]xf32> + %logits_view = buffer.view %logits[%c0_offset] : buffer -> view<[%token_count]x[%expert_count]xf32> + %c0_f32x4 = vector.splat %c0_f32 : vector<4xf32> + %lane_sums = scf.for %channel = [%lane_channel to %hidden_size step %c256](%accumulators = %c0_f32x4 : vector<4xf32>) -> (vector<4xf32>) unroll(%c4) schedule(interleaved) { + %input_values = vector.load %input_view[%token, %channel] : view<[%token_count]x[%hidden_size]xf32> -> vector<4xf32> + %weight_values0 = vector.load %weight_view[%expert_base, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> + %weight_values1 = vector.load %weight_view[%expert1, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> + %weight_values2 = vector.load %weight_view[%expert2, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> + %weight_values3 = vector.load %weight_view[%expert3, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> + %accumulator0 = vector.extract %accumulators[0] : vector<4xf32> -> f32 + %accumulator1 = vector.extract %accumulators[1] : vector<4xf32> -> f32 + %accumulator2 = vector.extract %accumulators[2] : vector<4xf32> -> f32 + %accumulator3 = vector.extract %accumulators[3] : vector<4xf32> -> f32 + %next0 = vector.dotf %input_values, %weight_values0, %accumulator0 : vector<4xf32>, vector<4xf32>, f32 + %next1 = vector.dotf %input_values, %weight_values1, %accumulator1 : vector<4xf32>, vector<4xf32>, f32 + %next2 = vector.dotf %input_values, %weight_values2, %accumulator2 : vector<4xf32>, vector<4xf32>, f32 + %next3 = vector.dotf %input_values, %weight_values3, %accumulator3 : vector<4xf32>, vector<4xf32>, f32 + %next_accumulators = vector.from_elements %next0, %next1, %next2, %next3 : vector<4xf32> + scf.yield %next_accumulators : vector<4xf32> + } + %sums = kernel.subgroup.reduce %lane_sums : vector<4xf32> + scf.if %writes_projection { + vector.store %sums, %logits_view[%token, %expert_base] : vector<4xf32>, view<[%token_count]x[%expert_count]xf32> + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/router_projection_storage.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/router_projection_storage.loom new file mode 100644 index 000000000000..39cd1a9bea21 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/router_projection_storage.loom @@ -0,0 +1,19 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @qwen3_moe.router_projection.storage(%input: buffer, %weight: buffer, %output: buffer) -> (buffer, buffer, buffer) + +amdgpu.target @qwen3_moe_router_projection_gfx1151 + +// Selects the target's preferred storage contract without exposing addressing +// policy at the projection call site. +template.def<@qwen3_moe.router_projection.storage> target(@qwen3_moe_router_projection_gfx1151) priority(20) @qwen3_moe_router_projection_descriptor_storage(%input: buffer, %weight: buffer, %output: buffer) -> (buffer, buffer, buffer) { + %input_descriptor = buffer.assume.memory_space %input : buffer + %weight_descriptor = buffer.assume.memory_space %weight : buffer + %output_descriptor = buffer.assume.memory_space %output : buffer + template.return %input_descriptor, %weight_descriptor, %output_descriptor : buffer, buffer, buffer +} + +template.def<@qwen3_moe.router_projection.storage> priority(1) @qwen3_moe_router_projection_global_storage(%input: buffer, %weight: buffer, %output: buffer) -> (buffer, buffer, buffer) { + template.return %input, %weight, %output : buffer, buffer, buffer +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/router_topk_row.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/router_topk_row.loom new file mode 100644 index 000000000000..45aadcd71549 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/router_topk_row.loom @@ -0,0 +1,100 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @qwen3_moe.router.top8.row(%valid_token: i1, %token: index, %token_count: index, %route_id_stride: index, %logits: buffer, %route_ids: buffer, %route_weights: buffer) + +config.decl @qwen3_moe.router.expert_count : %value: index where [range(%value, 32, 512), mul(%value, 32)] + +config.decl @qwen3_moe.router.route_count : %value: index where [range(%value, 1, 32)] + +// Selects and normalizes one logical router row. The caller owns mapping a +// subgroup to a safe physical row and tells the helper whether that row should +// publish. Keeping row semantics independent of launch geometry lets fused +// producers consume this exact tie-breaking and normalization contract. +template.def<@qwen3_moe.router.top8.row> device requires [#target.subgroup.size<64>] @qwen3_moe_router_top8_row_f32(%valid_token: i1, %token: index, %token_count: index, %route_id_stride: index, %logits: buffer, %route_ids: buffer, %route_weights: buffer) { + %expert_count0 = config.get @qwen3_moe.router.expert_count : index + %route_count0 = config.get @qwen3_moe.router.route_count : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_route_id_stride0 = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index + %expert_count, %route_count, %bounded_route_id_stride = index.assume %expert_count0, %route_count0, %bounded_route_id_stride0 [range(%expert_count0, 32, 512), mul(%expert_count0, 32), range(%route_count0, 1, 32), le(%route_count0, %expert_count0), le(%route_count0, %bounded_route_id_stride0)] : index, index, index + %safe_token, %launch_token_count = index.assume %token, %bounded_token_count [lt(%token, %bounded_token_count)] : index, index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %wave_size = target.subgroup.size : index + %negative_large = scalar.constant -3.4028234663852886e+38 : f32 + %largest_i32 = scalar.constant 2147483647 : i32 + %c0_offset = index.constant 0 : offset + %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 + %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> + %route_ids_view = buffer.view %route_ids_noalias[%c0_offset] : buffer -> view<[%route_id_storage_count]xi32> + %route_weights_view = buffer.view %route_weights_noalias[%c0_offset] : buffer -> view<[%route_weight_storage_count]xf32> + %initial_logits = vector.load %logits_view[%safe_token, %lane_expert_base] : view<[%launch_token_count]x[%expert_count]xf32> -> vector<[%experts_per_lane]xf32> + %remaining_final, %selected_logit = scf.for %route = [%c0 to %route_count step %c1](%remaining_logits = %initial_logits : vector<[%experts_per_lane]xf32>, %lane_selected_logit = %negative_large : f32) -> (vector<[%experts_per_lane]xf32>, f32) { + %local_value, %local_id = scf.for %slot = [%c0 to %experts_per_lane step %c1](%best_value = %negative_large : f32, %best_id = %largest_i32 : i32) -> (f32, i32) unroll { + %candidate_value = vector.extract %remaining_logits[%slot] : vector<[%experts_per_lane]xf32> -> f32 + %candidate_expert = index.add %lane_expert_base, %slot : index + %candidate_id = index.cast %candidate_expert : index to i32 + %is_greater = scalar.cmpf ogt, %candidate_value, %best_value : f32 + %is_equal = scalar.cmpf oeq, %candidate_value, %best_value : f32 + %is_lower_id = scalar.cmpi ult, %candidate_id, %best_id : i32 + %is_lower_tie = scalar.andi %is_equal, %is_lower_id : i1 + %is_better = scalar.ori %is_greater, %is_lower_tie : i1 + %next_value, %next_id = scf.if %is_better -> (f32, i32) { + scf.yield %candidate_value, %candidate_id : f32, i32 + } else { + scf.yield %best_value, %best_id : f32, i32 + } + scf.yield %next_value, %next_id : f32, i32 + } + %winner_value = kernel.subgroup.reduce %local_value : f32 + %matches_winner = scalar.cmpf oeq, %local_value, %winner_value : f32 + %winner_lane_mask = kernel.subgroup.vote.ballot %matches_winner : i1 -> i64 + %nonzero_winner_lane_mask = scalar.assume %winner_lane_mask [ne(%winner_lane_mask, 0)] : i64 + %winner_lane_i64 = scalar.cttzi %nonzero_winner_lane_mask : i64 + %winner_lane0 = index.cast %winner_lane_i64 : i64 to index + %winner_lane = index.assume %winner_lane0 [range(%winner_lane0, 0, 63)] : index + %local_expert0 = index.cast %local_id : i32 to index + %local_expert = index.assume %local_expert0 [range(%local_expert0, 0, 511)] : index + %winner_slot = index.rem %local_expert, %experts_per_lane : index + %owns_winner = index.cmp eq, %lane, %winner_lane : index + %next_remaining_logits = scf.if %owns_winner -> (vector<[%experts_per_lane]xf32>) { + %removed = vector.insert %negative_large into %remaining_logits[%winner_slot] : f32, vector<[%experts_per_lane]xf32> + scf.yield %removed : vector<[%experts_per_lane]xf32> + } else { + scf.yield %remaining_logits : vector<[%experts_per_lane]xf32> + } + %publishes_route_id = scalar.andi %valid_token, %owns_winner : i1 + scf.if %publishes_route_id { + %route_id_token_base = index.mul %safe_token, %bounded_route_id_stride : index + %route_id_index = index.add %route_id_token_base, %route : index + view.store %local_id, %route_ids_view[%route_id_index] : i32, view<[%route_id_storage_count]xi32> + } + %lane_publishes_route = index.cmp eq, %lane, %route : index + %publishes_route = scalar.andi %valid_token, %lane_publishes_route : i1 + %next_lane_selected_logit = scf.if %publishes_route -> (f32) { + scf.yield %winner_value : f32 + } else { + scf.yield %lane_selected_logit : f32 + } + scf.yield %next_remaining_logits, %next_lane_selected_logit : vector<[%experts_per_lane]xf32>, f32 + } + %selected_max = kernel.subgroup.reduce %selected_logit : f32 + %selected_delta = scalar.subf %selected_logit, %selected_max : f32 + %unnormalized_weight = scalar.expf %selected_delta : f32 + %selected_sum = kernel.subgroup.reduce %unnormalized_weight : f32 + %route_weight = scalar.divf %unnormalized_weight, %selected_sum : f32 + %lane_publishes_weight = index.cmp ult, %lane, %route_count : index + %publishes_weight = scalar.andi %valid_token, %lane_publishes_weight : i1 + scf.if %publishes_weight { + %route_weight_token_base = index.mul %safe_token, %route_count : index + %route_weight_index = index.add %route_weight_token_base, %lane : index + view.store %route_weight, %route_weights_view[%route_weight_index] : f32, view<[%route_weight_storage_count]xf32> + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/weighted_reduce_f16_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/weighted_reduce_f16_f32.loom new file mode 100644 index 000000000000..f3f92446428d --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/weighted_reduce_f16_f32.loom @@ -0,0 +1,57 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.routing.weighted_reduce_f16_f32(%token_count: index, %route_weights: buffer, %routed_output: buffer, %residual_input: buffer, %output: buffer) + +config.decl @qwen3_moe.routed_down.route_count : %value: index where [range(%value, 1, 8)] + +config.decl @qwen3_moe.routed_down.output_size : %value: index where [range(%value, 1, 4096)] + +// Preserves route-order FP32 accumulation before loading and adding the +// residual. The caller owns launch geometry and the public ABI. +template.def<@ggml.routing.weighted_reduce_f16_f32> device @ggml_routing_weighted_reduce_f16_f32(%token_count: index, %route_weights: buffer, %routed_output: buffer, %residual_input: buffer, %output: buffer) { + %route_count = config.get @qwen3_moe.routed_down.route_count : index + %output_size = config.get @qwen3_moe.routed_down.output_size : index + %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 4096)] : index + %channel_tile = kernel.workgroup.id : index + %token0 = kernel.workgroup.id : index + %lane = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %c0_f16x4 = vector.constant 0.0 : vector<4xf16> + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %assignment_count = index.mul %token_count, %bounded_route_count : index + %channel_tile_base = index.mul %channel_tile, %c256 : index + %lane_channel_add = index.mul %lane, %c4 : index + %channel = index.add %channel_tile_base, %lane_channel_add : index + %active_token = index.cmp ult, %token0, %token_count : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %publishes_output = scalar.andi %active_token, %valid_channel : i1 + %route_weights_noalias, %routed_output_noalias, %residual_input_noalias, %output_noalias = buffer.assume.noalias %route_weights, %routed_output, %residual_input, %output : buffer, buffer, buffer, buffer + %route_weight_view = buffer.view %route_weights_noalias[%c0_offset] : buffer -> view<[%token_count]x[%bounded_route_count]xf32> + %routed_output_view = buffer.view %routed_output_noalias[%c0_offset] : buffer -> view<[%assignment_count]x[%bounded_output_size]xf16> + %residual_input_view = buffer.view %residual_input_noalias[%c0_offset] : buffer -> view<[%token_count]x[%bounded_output_size]xf32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%token_count]x[%bounded_output_size]xf32> + scf.if %publishes_output { + %token, %output_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %mask = vector.mask.range [%channel to %bounded_output_size step %c1] : index -> vector<4xi1> + %weighted_sum = scf.for %route = [%c0 to %bounded_route_count step %c1](%sum = %c0_f32x4 : vector<4xf32>) -> (vector<4xf32>) unroll { + %assignment = index.madd %token, %bounded_route_count, %route : index + %routed = vector.load.mask %routed_output_view[%assignment, %channel], %mask, %c0_f16x4 : view<[%assignment_count]x[%bounded_output_size]xf16>, vector<4xi1>, vector<4xf16> + %wide = vector.extf %routed : vector<4xf16> to vector<4xf32> + %route_weight = view.load %route_weight_view[%token, %route] : view<[%token_count]x[%bounded_route_count]xf32> -> f32 + %route_weight_x4 = vector.splat %route_weight : vector<4xf32> + %weighted = vector.mulf %wide, %route_weight_x4 : vector<4xf32> + %next = vector.addf %sum, %weighted : vector<4xf32> + scf.yield %next : vector<4xf32> + } + %residual = vector.load.mask %residual_input_view[%token, %channel], %mask, %c0_f32x4 : view<[%token_count]x[%bounded_output_size]xf32>, vector<4xi1>, vector<4xf32> + %result = vector.addf %residual, %weighted_sum : vector<4xf32> + vector.store.mask %result, %output_view[%token, %channel], %mask : vector<4xf32>, view<[%token_count]x[%bounded_output_size]xf32>, vector<4xi1> + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/weighted_reduce_next_rmsnorm_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/weighted_reduce_next_rmsnorm_f32.loom new file mode 100644 index 000000000000..6d68753088a9 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/weighted_reduce_next_rmsnorm_f32.loom @@ -0,0 +1,104 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.routing.weighted_reduce_next_rmsnorm_f32(%token_count: index, %route_weights: buffer, %routed_output: buffer, %hidden_state: buffer, %next_norm_weight: buffer, %next_projection_input: buffer) + +config.decl @qwen3_moe.model.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] + +config.decl @qwen3_moe.model.rms_epsilon : f32 + +config.decl @qwen3_moe.routed_down.route_count : %value: index where [range(%value, 1, 8)] + +config.decl @qwen3_moe.routed_down.output_size : %value: index where [range(%value, 1, 4096)] + +// Preserves the specialized route-order reduction, residual publication, and +// row RMSNorm sequence. The caller owns launch geometry and the public ABI. +template.def<@ggml.routing.weighted_reduce_next_rmsnorm_f32> device requires [#target.subgroup.size<32>] @ggml_routing_weighted_reduce_next_rmsnorm_f32(%token_count: index, %route_weights: buffer, %routed_output: buffer, %hidden_state: buffer, %next_norm_weight: buffer, %next_projection_input: buffer) { + %route_count0 = config.get @qwen3_moe.routed_down.route_count : index + %output_size0 = config.get @qwen3_moe.routed_down.output_size : index + %hidden_size0 = config.get @qwen3_moe.model.hidden_size : index + %epsilon = config.get @qwen3_moe.model.rms_epsilon : f32 + %route_count = index.assume %route_count0 [range(%route_count0, 8, 8)] : index + %output_size, %hidden_size = index.assume %output_size0, %hidden_size0 [range(%output_size0, 2048, 2048), range(%hidden_size0, 2048, 2048), eq(%output_size0, %hidden_size0)] : index, index + %token0 = kernel.workgroup.id : index + %workitem0 = kernel.workitem.id : index + %workitem = index.assume %workitem0 [range(%workitem0, 0, 255)] : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c8 = index.constant 8 : index + %channel = index.mul %workitem, %c8 : index + %scratch_bytes = index.constant 32 : offset + %c0_offset = index.constant 0 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %c0_f32x8 = vector.constant 0.0 : vector<8xf32> + %active_token = index.cmp ult, %token0, %token_count : index + %safe_token0 = scf.select %active_token, %token0, %c0 : index + %token, %launch_token_count = index.assume %safe_token0, %token_count [lt(%safe_token0, %token_count)] : index, index + %hidden_size_i32 = index.cast %hidden_size : index to i32 + %hidden_size_f32 = scalar.sitofp %hidden_size_i32 : i32 to f32 + %assignment_count = index.mul %launch_token_count, %route_count : index + %route_weights_noalias, %routed_output_noalias, %hidden_state_noalias, %next_norm_weight_noalias, %next_projection_input_noalias = buffer.assume.noalias %route_weights, %routed_output, %hidden_state, %next_norm_weight, %next_projection_input : buffer, buffer, buffer, buffer, buffer + %route_weights_view = buffer.view %route_weights_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%route_count]xf32> + %routed_output_view = buffer.view %routed_output_noalias[%c0_offset] : buffer -> view<[%assignment_count]x[%output_size]xf16> + %hidden_state_view = buffer.view %hidden_state_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + %next_norm_weight_view = buffer.view %next_norm_weight_noalias[%c0_offset] : buffer -> view<[%hidden_size]xf32> + %next_projection_input_view = buffer.view %next_projection_input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_view = buffer.view %scratch[%c0_offset] : buffer -> view<8xf32> + + // The predicate depends only on the workgroup ID, so every workitem either + // executes the complete barrier-bearing row reduction or skips it together. + scf.if %active_token { + %weighted_sum = scf.for %route = [%c0 to %route_count step %c1](%sum = %c0_f32x8 : vector<8xf32>) -> (vector<8xf32>) unroll { + %assignment = index.madd %token, %route_count, %route : index + %routed_f16 = vector.load %routed_output_view[%assignment, %channel] : view<[%assignment_count]x[%output_size]xf16> -> vector<8xf16> + %routed = vector.extf %routed_f16 : vector<8xf16> to vector<8xf32> + %route_weight = view.load %route_weights_view[%token, %route] : view<[%launch_token_count]x[%route_count]xf32> -> f32 + %route_weight_x8 = vector.splat %route_weight : vector<8xf32> + %weighted = vector.mulf %routed, %route_weight_x8 : vector<8xf32> + %next = vector.addf %sum, %weighted : vector<8xf32> + scf.yield %next : vector<8xf32> + } + %residual = vector.load %hidden_state_view[%token, %channel] : view<[%launch_token_count]x[%hidden_size]xf32> -> vector<8xf32> + %result = vector.addf %residual, %weighted_sum : vector<8xf32> + %squares = vector.mulf %result, %result : vector<8xf32> + %thread_sum = vector.reduce %squares, %c0_f32 : vector<8xf32>, f32 + %subgroup_sum = kernel.subgroup.reduce %thread_sum : f32 + + %is_subgroup_leader = index.cmp eq, %lane, %c0 : index + scf.if %is_subgroup_leader { + view.store %subgroup_sum, %scratch_view[%subgroup] : f32, view<8xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %is_reduction_subgroup = index.cmp eq, %subgroup, %c0 : index + %is_reduction_lane = index.cmp ult, %lane, %c8 : index + %loads_subgroup_sum = scalar.andi %is_reduction_subgroup, %is_reduction_lane : i1 + %subgroup_partial = scf.if %loads_subgroup_sum -> (f32) { + %value = view.load %scratch_view[%lane] : view<8xf32> -> f32 + scf.yield %value : f32 + } else { + scf.yield %c0_f32 : f32 + } + %row_sum = kernel.subgroup.reduce %subgroup_partial : f32 + %writes_scale = scalar.andi %is_reduction_subgroup, %is_subgroup_leader : i1 + scf.if %writes_scale { + %mean = scalar.divf %row_sum, %hidden_size_f32 : f32 + %biased_mean = scalar.addf %mean, %epsilon : f32 + %scale = scalar.rsqrtf %biased_mean : f32 + view.store %scale, %scratch_view[%c0] : f32, view<8xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + + %scale = view.load %scratch_view[%c0] : view<8xf32> -> f32 + %scale_x8 = vector.splat %scale : vector<8xf32> + %learned_weight = vector.load %next_norm_weight_view[%channel] : view<[%hidden_size]xf32> -> vector<8xf32> + %normalized = vector.mulf %result, %scale_x8 : vector<8xf32> + %next_projection = vector.mulf %normalized, %learned_weight : vector<8xf32> + vector.store %result, %hidden_state_view[%token, %channel] : vector<8xf32>, view<[%launch_token_count]x[%hidden_size]xf32> + vector.store %next_projection, %next_projection_input_view[%token, %channel] : vector<8xf32>, view<[%launch_token_count]x[%hidden_size]xf32> + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/scale_add_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/scale_add_f32.loom new file mode 100644 index 000000000000..dfd639bac97a --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/scale_add_f32.loom @@ -0,0 +1,96 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.elementwise.scale_apply(%value: f32, %scale: f32, %bias: f32) -> (f32) +template.decl @ggml.binary_f32.apply(%op: index, %lhs: f32, %rhs: f32) -> (f32) + +amdgpu.target @ggml_scale_add_f32_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.scale_add_f32.ne0 : %value: index where [range(%value, 1, 134217728)] +config.decl @ggml.scale_add_f32.ne1 : %value: index where [range(%value, 1, 134217728)] +config.decl @ggml.scale_add_f32.ne2 : %value: index where [range(%value, 1, 134217728)] +config.decl @ggml.scale_add_f32.input_stride1 : %value: index where [range(%value, 1, 134217728)] +config.decl @ggml.scale_add_f32.input_stride2 : %value: index where [range(%value, 1, 134217728)] +config.decl @ggml.scale_add_f32.input_stride3 : %value: index where [range(%value, 1, 134217728)] +config.decl @ggml.scale_add_f32.input_span : %value: index where [range(%value, 1, 134217728)] +config.decl @ggml.scale_add_f32.residual_stride1 : %value: index where [range(%value, 1, 134217728)] +config.decl @ggml.scale_add_f32.residual_stride2 : %value: index where [range(%value, 1, 134217728)] +config.decl @ggml.scale_add_f32.residual_stride3 : %value: index where [range(%value, 1, 134217728)] +config.decl @ggml.scale_add_f32.residual_span : %value: index where [range(%value, 1, 134217728)] +config.decl @ggml.scale_add_f32.scale : f32 +config.decl @ggml.scale_add_f32.bias : f32 +config.decl @ggml.scale_add_f32.binary_op : %value: index where [range(%value, 0, 3)] +config.decl @ggml.scale_add_f32.scaled_lhs : %value: index where [range(%value, 0, 1)] + +kernel.def target(@ggml_scale_add_f32_gfx11_wave64) export("ggml_scale_add_f32") @ggml_scale_add_f32(%element_count: index) { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + %c255 = index.constant 255 : index + %rounded = index.add %element_count, %c255 : index + %workgroup_count = index.div %rounded, %c256 : index + kernel.launch.config workgroups(%workgroup_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%element_count: index, %input: buffer, %residual: buffer, %output: buffer) { + %count = index.assume %element_count [range(%element_count, 1, 134217728)] : index + %n0 = config.get @ggml.scale_add_f32.ne0 : index + %n1 = config.get @ggml.scale_add_f32.ne1 : index + %n2 = config.get @ggml.scale_add_f32.ne2 : index + %input_stride1 = config.get @ggml.scale_add_f32.input_stride1 : index + %input_stride2 = config.get @ggml.scale_add_f32.input_stride2 : index + %input_stride3 = config.get @ggml.scale_add_f32.input_stride3 : index + %input_count = config.get @ggml.scale_add_f32.input_span : index + %residual_stride1 = config.get @ggml.scale_add_f32.residual_stride1 : index + %residual_stride2 = config.get @ggml.scale_add_f32.residual_stride2 : index + %residual_stride3 = config.get @ggml.scale_add_f32.residual_stride3 : index + %residual_count = config.get @ggml.scale_add_f32.residual_span : index + %scale = config.get @ggml.scale_add_f32.scale : f32 + %bias = config.get @ggml.scale_add_f32.bias : f32 + %binary_op = config.get @ggml.scale_add_f32.binary_op : index + %scaled_lhs = config.get @ggml.scale_add_f32.scaled_lhs : index + %workgroup = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %c256 = index.constant 256 : index + %base0 = index.mul %workgroup, %c256 : index + %linear0 = index.add %base0, %workitem : index + %linear = index.assume %linear0 [range(%linear0, 0, 134217983)] : index + %in_bounds = index.cmp ult, %linear, %count : index + %c0_offset = index.constant 0 : offset + %input_view = buffer.view %input[%c0_offset] : buffer -> view<[%input_count]xf32> + %residual_view = buffer.view %residual[%c0_offset] : buffer -> view<[%residual_count]xf32> + %output_view = buffer.view %output[%c0_offset] : buffer -> view<[%count]xf32> + scf.if %in_bounds { + %coord0 = index.rem %linear, %n0 : index + %linear_div_n0 = index.div %linear, %n0 : index + %coord1 = index.rem %linear_div_n0, %n1 : index + %linear_div_n01 = index.div %linear_div_n0, %n1 : index + %coord2 = index.rem %linear_div_n01, %n2 : index + %coord3 = index.div %linear_div_n01, %n2 : index + %input_term1 = index.mul %coord1, %input_stride1 : index + %input_term2 = index.mul %coord2, %input_stride2 : index + %input_term3 = index.mul %coord3, %input_stride3 : index + %input_linear01 = index.add %coord0, %input_term1 : index + %input_linear012 = index.add %input_linear01, %input_term2 : index + %input_linear = index.add %input_linear012, %input_term3 : index + %residual_term1 = index.mul %coord1, %residual_stride1 : index + %residual_term2 = index.mul %coord2, %residual_stride2 : index + %residual_term3 = index.mul %coord3, %residual_stride3 : index + %residual_linear01 = index.add %coord0, %residual_term1 : index + %residual_linear012 = index.add %residual_linear01, %residual_term2 : index + %residual_linear = index.add %residual_linear012, %residual_term3 : index + %input_in_bounds = index.cmp ult, %input_linear, %input_count : index + scf.if %input_in_bounds { + %residual_in_bounds = index.cmp ult, %residual_linear, %residual_count : index + scf.if %residual_in_bounds { + %input_value = view.load %input_view[%input_linear] : view<[%input_count]xf32> -> f32 + %scaled = template.apply<@ggml.elementwise.scale_apply>(%input_value, %scale, %bias) : (f32, f32, f32) -> (f32) + %residual_value = view.load %residual_view[%residual_linear] : view<[%residual_count]xf32> -> f32 + %c1 = index.constant 1 : index + %is_scaled_lhs = index.cmp eq, %scaled_lhs, %c1 : index + %lhs = scf.select %is_scaled_lhs, %scaled, %residual_value : f32 + %rhs = scf.select %is_scaled_lhs, %residual_value, %scaled : f32 + %result = template.apply<@ggml.binary_f32.apply>(%binary_op, %lhs, %rhs) : (index, f32, f32) -> (f32) + view.store %result, %output_view[%linear] : f32, view<[%count]xf32> + } + } + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_v_matmul_set_rows_tiled_f32_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_v_matmul_set_rows_tiled_f32_f32.loom index cc6bb6d32be7..2b72a2feb93b 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_v_matmul_set_rows_tiled_f32_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_v_matmul_set_rows_tiled_f32_f32.loom @@ -14,6 +14,8 @@ config.decl @llm.attention_qkv.output_size : %value: index where [range(%value, config.decl @llm.attention_qkv.weight_format : %value: index where [range(%value, 4, 81)] +config.decl @ggml.mul_mat.activation_format : %value: index where [range(%value, 0, 16)] + config.decl @llm.attention_qkv.cache_row_count : %value: index where [range(%value, 1, 1048576)] config.decl @llm.attention_qkv.cache_output_format : %value: index where [range(%value, 16, 32)] diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/skinny_input_f32_publish_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/skinny_input_f32_publish_f32.loom index 5135b662819c..7f5839752649 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/skinny_input_f32_publish_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/skinny_input_f32_publish_f32.loom @@ -36,6 +36,8 @@ config.decl @ggml.mul_mat_postops.output_size : %value: index where [range(%valu config.decl @ggml.mul_mat_postops.weight_format : %value: index where [range(%value, 4, 81)] +config.decl @ggml.mul_mat_postops.epilogue : %value: index where [range(%value, 1, 3)] + config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 2, 5)] template.def<@ggml.matmul.skinny_input_f32_dequant_accumulate_f32.publish_vector4> device @ggml_matmul_skinny_input_f32_publish_f32_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { @@ -115,53 +117,7 @@ kernel.def target(@ggml_matmul_skinny_input_f32_publish_f32_gfx11_wave64) @ggml_ kernel.return } -kernel.def target(@ggml_matmul_skinny_input_f32_publish_f32_gfx11_wave64) @ggml_mul_mat_skinny_input_f32_bias_publish_f32(%token_count: index) { - %output_size = config.get @ggml.mul_mat_postops.output_size : index - %token_capacity = config.get @ggml.workload.token_capacity : index - %c1 = index.constant 1 : index - %input_size = config.get @ggml.mul_mat_postops.input_size : index - %weight_format = config.get @ggml.mul_mat_postops.weight_format : index - %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.matmul.skinny_input_f32_dequant_accumulate_f32.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) - kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index -} launch(%token_count: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer) { - %input_size = config.get @ggml.mul_mat_postops.input_size : index - %output_size = config.get @ggml.mul_mat_postops.output_size : index - %token_capacity = config.get @ggml.workload.token_capacity : index - %weight_format = config.get @ggml.mul_mat_postops.weight_format : index - %c0 = index.constant 0 : index - %identity = index.constant 23 : index - %epilogue = index.constant 1 : index - %c0_f32 = scalar.constant 0.0 : f32 - %bounded_token_count = index.assume %token_count [range(%token_count, 2, 5), le(%token_count, %token_capacity)] : index - %channel_tile = kernel.workgroup.id : index - %token_tile = kernel.workgroup.id : index - template.apply<@ggml.matmul.skinny_input_f32_dequant_accumulate_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %identity, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %bias, %output, %output, %output, %output, %output, %output) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) - kernel.return -} -kernel.def target(@ggml_matmul_skinny_input_f32_publish_f32_gfx11_wave64) @ggml_mul_mat_skinny_input_f32_residual_publish_f32(%token_count: index) { - %output_size = config.get @ggml.mul_mat_postops.output_size : index - %token_capacity = config.get @ggml.workload.token_capacity : index - %c1 = index.constant 1 : index - %input_size = config.get @ggml.mul_mat_postops.input_size : index - %weight_format = config.get @ggml.mul_mat_postops.weight_format : index - %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.matmul.skinny_input_f32_dequant_accumulate_f32.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) - kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index -} launch(%token_count: index, %input: buffer, %weight: buffer, %residual_input: buffer, %residual_output: buffer) { - %input_size = config.get @ggml.mul_mat_postops.input_size : index - %output_size = config.get @ggml.mul_mat_postops.output_size : index - %token_capacity = config.get @ggml.workload.token_capacity : index - %weight_format = config.get @ggml.mul_mat_postops.weight_format : index - %c0 = index.constant 0 : index - %identity = index.constant 23 : index - %epilogue = index.constant 2 : index - %c0_f32 = scalar.constant 0.0 : f32 - %bounded_token_count = index.assume %token_count [range(%token_count, 2, 5), le(%token_count, %token_capacity)] : index - %channel_tile = kernel.workgroup.id : index - %token_tile = kernel.workgroup.id : index - template.apply<@ggml.matmul.skinny_input_f32_dequant_accumulate_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %identity, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %residual_output, %residual_output, %residual_input, %residual_output, %residual_output, %residual_output, %residual_output) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) - kernel.return -} kernel.def target(@ggml_matmul_skinny_input_f32_publish_f32_gfx11_wave64) @ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32(%token_count: index) { %output_size = config.get @ggml.mul_mat_postops.output_size : index @@ -178,7 +134,7 @@ kernel.def target(@ggml_matmul_skinny_input_f32_publish_f32_gfx11_wave64) @ggml_ %weight_format = config.get @ggml.mul_mat_postops.weight_format : index %c0 = index.constant 0 : index %identity = index.constant 23 : index - %epilogue = index.constant 3 : index + %epilogue = config.get @ggml.mul_mat_postops.epilogue : index %c0_f32 = scalar.constant 0.0 : f32 %bounded_token_count = index.assume %token_count [range(%token_count, 2, 5), le(%token_count, %token_capacity)] : index %channel_tile = kernel.workgroup.id : index diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_input_f32_publish_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_input_f32_publish_f32.loom index 9e7cde4571b4..217155372655 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_input_f32_publish_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_input_f32_publish_f32.loom @@ -38,6 +38,8 @@ config.decl @ggml.mul_mat_postops.output_size : %value: index where [range(%valu config.decl @ggml.mul_mat_postops.weight_format : %value: index where [range(%value, 4, 81)] +config.decl @ggml.mul_mat_postops.epilogue : %value: index where [range(%value, 1, 3)] + config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] template.def<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4> device @ggml_matmul_tiled_input_f32_publish_f32_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { @@ -172,65 +174,7 @@ kernel.def target(@ggml_matmul_tiled_input_f32_publish_f32_gfx11_wave64) @ggml_m kernel.return } -kernel.def target(@ggml_matmul_tiled_input_f32_publish_f32_gfx11_wave64) @ggml_mul_mat_tiled_input_f32_bias_publish_f32(%token_count: index) { - %output_size = config.get @ggml.mul_mat_postops.output_size : index - %token_capacity = config.get @ggml.workload.token_capacity : index - %c1 = index.constant 1 : index - %c31 = index.constant 31 : index - %c32 = index.constant 32 : index - %c63 = index.constant 63 : index - %c64 = index.constant 64 : index - %c128 = index.constant 128 : index - %padded_output_size = index.add %output_size, %c63 : index - %output_tiles = index.div %padded_output_size, %c64 : index - %padded_token_count = index.add %token_capacity, %c31 : index - %token_tiles = index.div %padded_token_count, %c32 : index - kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index -} launch(%token_count: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer) { - %input_size = config.get @ggml.mul_mat_postops.input_size : index - %output_size = config.get @ggml.mul_mat_postops.output_size : index - %token_capacity = config.get @ggml.workload.token_capacity : index - %weight_format = config.get @ggml.mul_mat_postops.weight_format : index - %c0 = index.constant 0 : index - %identity = index.constant 23 : index - %epilogue = index.constant 1 : index - %c0_f32 = scalar.constant 0.0 : f32 - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index - %channel_tile = kernel.workgroup.id : index - %token_tile = kernel.workgroup.id : index - template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %identity, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %bias, %output, %output, %output, %output, %output, %output) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) - kernel.return -} -kernel.def target(@ggml_matmul_tiled_input_f32_publish_f32_gfx11_wave64) @ggml_mul_mat_tiled_input_f32_residual_publish_f32(%token_count: index) { - %output_size = config.get @ggml.mul_mat_postops.output_size : index - %token_capacity = config.get @ggml.workload.token_capacity : index - %c1 = index.constant 1 : index - %c31 = index.constant 31 : index - %c32 = index.constant 32 : index - %c63 = index.constant 63 : index - %c64 = index.constant 64 : index - %c128 = index.constant 128 : index - %padded_output_size = index.add %output_size, %c63 : index - %output_tiles = index.div %padded_output_size, %c64 : index - %padded_token_count = index.add %token_capacity, %c31 : index - %token_tiles = index.div %padded_token_count, %c32 : index - kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index -} launch(%token_count: index, %input: buffer, %weight: buffer, %residual_input: buffer, %residual_output: buffer) { - %input_size = config.get @ggml.mul_mat_postops.input_size : index - %output_size = config.get @ggml.mul_mat_postops.output_size : index - %token_capacity = config.get @ggml.workload.token_capacity : index - %weight_format = config.get @ggml.mul_mat_postops.weight_format : index - %c0 = index.constant 0 : index - %identity = index.constant 23 : index - %epilogue = index.constant 2 : index - %c0_f32 = scalar.constant 0.0 : f32 - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index - %channel_tile = kernel.workgroup.id : index - %token_tile = kernel.workgroup.id : index - template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %identity, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %residual_output, %residual_output, %residual_input, %residual_output, %residual_output, %residual_output, %residual_output) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) - kernel.return -} kernel.def target(@ggml_matmul_tiled_input_f32_publish_f32_gfx11_wave64) @ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32(%token_count: index) { %output_size = config.get @ggml.mul_mat_postops.output_size : index @@ -253,7 +197,7 @@ kernel.def target(@ggml_matmul_tiled_input_f32_publish_f32_gfx11_wave64) @ggml_m %weight_format = config.get @ggml.mul_mat_postops.weight_format : index %c0 = index.constant 0 : index %identity = index.constant 23 : index - %epilogue = index.constant 3 : index + %epilogue = config.get @ggml.mul_mat_postops.epilogue : index %c0_f32 = scalar.constant 0.0 : f32 %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index %channel_tile = kernel.workgroup.id : index diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom index 0e41b7bccf9f..7f5843424db5 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom @@ -25,6 +25,8 @@ config.decl @ggml.mul_mat_swiglu.up_weight_format : %value: index where [range(% config.decl @ggml.mul_mat_swiglu.op : %value: index where [range(%value, 0, 8)] +config.decl @ggml.mul_mat_swiglu.epilogue : %value: index where [range(%value, 1, 3)] + config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] template.def<@ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.publish_vector4> device @ggml_matmul_tiled_pair_input_f32_binary_publish_f32_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %op: index, %epilogue: index, %lhs: vector<4xf32>, %rhs: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer) { @@ -88,63 +90,7 @@ kernel.def target(@ggml_matmul_tiled_pair_input_f32_binary_publish_f32_gfx11_wav kernel.return } -kernel.def target(@ggml_matmul_tiled_pair_input_f32_binary_publish_f32_gfx11_wave64) @ggml_mul_mat_tiled_pair_input_f32_binary_bias_publish_f32(%token_count: index) { - %output_size = config.get @ggml.mul_mat_swiglu.output_size : index - %token_capacity = config.get @ggml.workload.token_capacity : index - %c1 = index.constant 1 : index - %c31 = index.constant 31 : index - %c32 = index.constant 32 : index - %c63 = index.constant 63 : index - %c64 = index.constant 64 : index - %c128 = index.constant 128 : index - %padded_output_size = index.add %output_size, %c63 : index - %output_tiles = index.div %padded_output_size, %c64 : index - %padded_token_count = index.add %token_capacity, %c31 : index - %token_tiles = index.div %padded_token_count, %c32 : index - kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index -} launch(%token_count: index, %input: buffer, %gate_weight: buffer, %up_weight: buffer, %bias: buffer, %output: buffer) where [range(%token_count, 2, 2048)] { - %input_size = config.get @ggml.mul_mat_swiglu.input_size : index - %output_size = config.get @ggml.mul_mat_swiglu.output_size : index - %token_capacity = config.get @ggml.workload.token_capacity : index - %gate_weight_format = config.get @ggml.mul_mat_swiglu.gate_weight_format : index - %up_weight_format = config.get @ggml.mul_mat_swiglu.up_weight_format : index - %op = config.get @ggml.mul_mat_swiglu.op : index - %epilogue = index.constant 1 : index - %bounded_token_count = index.assume %token_count [range(%token_count, 2, 2048), le(%token_count, %token_capacity)] : index - %channel_tile = kernel.workgroup.id : index - %token_tile = kernel.workgroup.id : index - template.apply<@ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.core>(%op, %epilogue, %gate_weight_format, %up_weight_format, %bounded_token_count, %input_size, %output_size, %channel_tile, %token_tile, %input, %gate_weight, %up_weight, %bias, %output, %output, %output) : (index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) - kernel.return -} -kernel.def target(@ggml_matmul_tiled_pair_input_f32_binary_publish_f32_gfx11_wave64) @ggml_mul_mat_tiled_pair_input_f32_binary_residual_publish_f32(%token_count: index) { - %output_size = config.get @ggml.mul_mat_swiglu.output_size : index - %token_capacity = config.get @ggml.workload.token_capacity : index - %c1 = index.constant 1 : index - %c31 = index.constant 31 : index - %c32 = index.constant 32 : index - %c63 = index.constant 63 : index - %c64 = index.constant 64 : index - %c128 = index.constant 128 : index - %padded_output_size = index.add %output_size, %c63 : index - %output_tiles = index.div %padded_output_size, %c64 : index - %padded_token_count = index.add %token_capacity, %c31 : index - %token_tiles = index.div %padded_token_count, %c32 : index - kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index -} launch(%token_count: index, %input: buffer, %gate_weight: buffer, %up_weight: buffer, %residual_input: buffer, %residual_output: buffer) where [range(%token_count, 2, 2048)] { - %input_size = config.get @ggml.mul_mat_swiglu.input_size : index - %output_size = config.get @ggml.mul_mat_swiglu.output_size : index - %token_capacity = config.get @ggml.workload.token_capacity : index - %gate_weight_format = config.get @ggml.mul_mat_swiglu.gate_weight_format : index - %up_weight_format = config.get @ggml.mul_mat_swiglu.up_weight_format : index - %op = config.get @ggml.mul_mat_swiglu.op : index - %epilogue = index.constant 2 : index - %bounded_token_count = index.assume %token_count [range(%token_count, 2, 2048), le(%token_count, %token_capacity)] : index - %channel_tile = kernel.workgroup.id : index - %token_tile = kernel.workgroup.id : index - template.apply<@ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.core>(%op, %epilogue, %gate_weight_format, %up_weight_format, %bounded_token_count, %input_size, %output_size, %channel_tile, %token_tile, %input, %gate_weight, %up_weight, %residual_output, %residual_output, %residual_input, %residual_output) : (index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) - kernel.return -} kernel.def target(@ggml_matmul_tiled_pair_input_f32_binary_publish_f32_gfx11_wave64) @ggml_mul_mat_tiled_pair_input_f32_binary_bias_residual_publish_f32(%token_count: index) { %output_size = config.get @ggml.mul_mat_swiglu.output_size : index @@ -167,7 +113,7 @@ kernel.def target(@ggml_matmul_tiled_pair_input_f32_binary_publish_f32_gfx11_wav %gate_weight_format = config.get @ggml.mul_mat_swiglu.gate_weight_format : index %up_weight_format = config.get @ggml.mul_mat_swiglu.up_weight_format : index %op = config.get @ggml.mul_mat_swiglu.op : index - %epilogue = index.constant 3 : index + %epilogue = config.get @ggml.mul_mat_swiglu.epilogue : index %bounded_token_count = index.assume %token_count [range(%token_count, 2, 2048), le(%token_count, %token_capacity)] : index %channel_tile = kernel.workgroup.id : index %token_tile = kernel.workgroup.id : index diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_postops_f32_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_postops_f32_f32.loom index 625812e33558..65402e073a27 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_postops_f32_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_postops_f32_f32.loom @@ -110,49 +110,7 @@ template.def<@ggml.matmul.vector_activation.finish_output_group> device @ggml_mu template.return } -kernel.def target(@ggml_matmul_vector_postops_gfx11_wave64) @ggml_mul_mat_vector_bias_f32_f32(%token_count: index) { - %input_size = config.get @ggml.matmul.vector_postops.input_size : index - %output_size = config.get @ggml.matmul.vector_postops.output_size : index - %token_capacity = config.get @ggml.workload.token_capacity : index - %weight_format = config.get @ggml.matmul.vector_postops.weight_format : index - %c1 = index.constant 1 : index - %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.matmul.vector_activation.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) - kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index -} launch(%token_count: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %next_output: buffer) { - %input_size = config.get @ggml.matmul.vector_postops.input_size : index - %output_size = config.get @ggml.matmul.vector_postops.output_size : index - %token_capacity = config.get @ggml.workload.token_capacity : index - %weight_format = config.get @ggml.matmul.vector_postops.weight_format : index - %c0 = index.constant 0 : index - %c0_f32 = scalar.constant 0.0 : f32 - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1), le(%token_count, %token_capacity)] : index - %channel_group = kernel.workgroup.id : index - %token_group = kernel.workgroup.id : index - template.apply<@ggml.matmul.vector_activation.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %c0, %c0_f32, %channel_group, %token_group, %input, %weight, %bias, %output, %output, %output, %output, %output, %output, %next_output) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) - kernel.return -} -kernel.def target(@ggml_matmul_vector_postops_gfx11_wave64) @ggml_mul_mat_vector_residual_f32_f32(%token_count: index) { - %input_size = config.get @ggml.matmul.vector_postops.input_size : index - %output_size = config.get @ggml.matmul.vector_postops.output_size : index - %token_capacity = config.get @ggml.workload.token_capacity : index - %weight_format = config.get @ggml.matmul.vector_postops.weight_format : index - %c1 = index.constant 1 : index - %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.matmul.vector_activation.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) - kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index -} launch(%token_count: index, %input: buffer, %weight: buffer, %residual_input: buffer, %residual_output: buffer, %next_output: buffer) { - %input_size = config.get @ggml.matmul.vector_postops.input_size : index - %output_size = config.get @ggml.matmul.vector_postops.output_size : index - %token_capacity = config.get @ggml.workload.token_capacity : index - %weight_format = config.get @ggml.matmul.vector_postops.weight_format : index - %c0 = index.constant 0 : index - %c0_f32 = scalar.constant 0.0 : f32 - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1), le(%token_count, %token_capacity)] : index - %channel_group = kernel.workgroup.id : index - %token_group = kernel.workgroup.id : index - template.apply<@ggml.matmul.vector_activation.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %c0, %c0_f32, %channel_group, %token_group, %input, %weight, %residual_output, %residual_output, %residual_input, %residual_output, %residual_output, %residual_output, %residual_output, %next_output) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) - kernel.return -} kernel.def target(@ggml_matmul_vector_postops_gfx11_wave64) @ggml_mul_mat_vector_bias_residual_f32_f32(%token_count: index) { %input_size = config.get @ggml.matmul.vector_postops.input_size : index diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_binary_strided_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_binary_strided_f32.loom new file mode 100644 index 000000000000..3ae10cd93821 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_binary_strided_f32.loom @@ -0,0 +1,34 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.rmsnorm_binary_strided_f32.body(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer) + +amdgpu.target @ggml_rmsnorm_binary_gfx11_wave32 {subgroup_size = 32} + +config.decl @ggml.rmsnorm_binary_f32.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] + +config.decl @ggml.rmsnorm_binary_f32.input_stride : %value: index where [range(%value, 128, 1048576)] + +config.decl @ggml.rmsnorm_binary_f32.rms_epsilon : f32 + +config.decl @ggml.rmsnorm_binary_f32.op : %value: index where [range(%value, 0, 3)] + +kernel.def target(@ggml_rmsnorm_binary_gfx11_wave32) export("ggml_rmsnorm_binary_strided_f32") @ggml_rmsnorm_binary_strided_f32(%token_count: index) { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%token_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer) where [range(%token_count, 1, 1048576)] { + template.apply<@ggml.rmsnorm_binary_strided_f32.body>(%token_count, %input, %rhs, %output) : (index, buffer, buffer, buffer) + kernel.return +} + +check.case public @ggml_rmsnorm_binary_strided_f32_mul_case { + %token_count = check.literal value(2) : index + %input = check.generate.fill value(2.0) : tensor<2x160xf32> + %rhs = check.generate.fill value(3.0) : tensor<128xf32> + %output = check.generate.fill value(0.0) : tensor<2x128xf32> + %expected = check.generate.fill value(3.0) : tensor<2x128xf32> + kernel.launch @ggml_rmsnorm_binary_strided_f32[%token_count](%token_count, %input, %rhs, %output) : [index](index, tensor<2x160xf32>, tensor<128xf32>, tensor<2x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(1.0000000000000001e-05) rtol(1.0000000000000001e-05) nan(same) : tensor<2x128xf32> + check.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_mul_add_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_mul_add_f32.loom new file mode 100644 index 000000000000..55ccbe655c28 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_mul_add_f32.loom @@ -0,0 +1,69 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.rmsnorm_f32.row_scale(%token_count: index, %token: index, %hidden_size: index, %input_stride: index, %epsilon: f32, %input: buffer) -> (f32, index, index, i1) + +template.decl @ggml.rmsnorm_f32.apply(%value: f32, %scale: f32) -> (f32) +template.decl @ggml.binary_f32.apply(%op: index, %lhs: f32, %rhs: f32) -> (f32) + +amdgpu.target @ggml_rmsnorm_mul_add_gfx11_wave32 {subgroup_size = 32} + +config.decl @ggml.rmsnorm_mul_add_f32.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] + +config.decl @ggml.rmsnorm_mul_add_f32.rms_epsilon : f32 +config.decl @ggml.rmsnorm_mul_add_f32.normalized_op : %value: index where [range(%value, 0, 3)] +config.decl @ggml.rmsnorm_mul_add_f32.output_op : %value: index where [range(%value, 0, 3)] +config.decl @ggml.rmsnorm_mul_add_f32.binary_lhs : %value: index where [range(%value, 0, 1)] + +kernel.def target(@ggml_rmsnorm_mul_add_gfx11_wave32) export("ggml_rmsnorm_mul_add_f32") @ggml_rmsnorm_mul_add_f32(%token_count: index) { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%token_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %residual: buffer, %output: buffer) where [range(%token_count, 1, 1048576)] { + %hidden0 = config.get @ggml.rmsnorm_mul_add_f32.hidden_size : index + %hidden = index.assume %hidden0 [range(%hidden0, 128, 32768), mul(%hidden0, 128)] : index + %epsilon = config.get @ggml.rmsnorm_mul_add_f32.rms_epsilon : f32 + %normalized_op = config.get @ggml.rmsnorm_mul_add_f32.normalized_op : index + %output_op = config.get @ggml.rmsnorm_mul_add_f32.output_op : index + %binary_lhs = config.get @ggml.rmsnorm_mul_add_f32.binary_lhs : index + %token0 = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %c0_offset = index.constant 0 : offset + %c256 = index.constant 256 : index + %input_noalias, %weight_noalias, %residual_noalias, %output_noalias = buffer.assume.noalias %input, %weight, %residual, %output : buffer, buffer, buffer, buffer + %row_scale, %token, %launch_token_count, %valid_token = template.apply<@ggml.rmsnorm_f32.row_scale>(%token_count, %token0, %hidden, %hidden, %epsilon, %input_noalias) : (index, index, index, index, f32, buffer) -> (f32, index, index, i1) + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden]xf32> + %weight_view = buffer.view %weight_noalias[%c0_offset] : buffer -> view<[%hidden]xf32> + %residual_view = buffer.view %residual_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden]xf32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden]xf32> + scf.if %valid_token { + scf.for %channel = [%workitem to %hidden step %c256] { + %value = view.load %input_view[%token, %channel] : view<[%launch_token_count]x[%hidden]xf32> -> f32 + %weight_value = view.load %weight_view[%channel] : view<[%hidden]xf32> -> f32 + %residual_value = view.load %residual_view[%token, %channel] : view<[%launch_token_count]x[%hidden]xf32> -> f32 + %normalized = template.apply<@ggml.rmsnorm_f32.apply>(%value, %row_scale) : (f32, f32) -> (f32) + %scaled = template.apply<@ggml.binary_f32.apply>(%normalized_op, %normalized, %weight_value) : (index, f32, f32) -> (f32) + %c1 = index.constant 1 : index + %is_binary_lhs = index.cmp eq, %binary_lhs, %c1 : index + %lhs = scf.select %is_binary_lhs, %scaled, %residual_value : f32 + %rhs = scf.select %is_binary_lhs, %residual_value, %scaled : f32 + %result = template.apply<@ggml.binary_f32.apply>(%output_op, %lhs, %rhs) : (index, f32, f32) -> (f32) + view.store %result, %output_view[%token, %channel] : f32, view<[%launch_token_count]x[%hidden]xf32> + } + } + kernel.return +} + +check.case public @ggml_rmsnorm_mul_add_f32_case { + %token_count = check.literal value(1) : index + %input = check.generate.fill value(2.0) : tensor<128xf32> + %weight = check.generate.iota offset(0.25) step(0.01) : tensor<128xf32> + %residual = check.generate.fill value(3.0) : tensor<128xf32> + %output = check.generate.fill value(0.0) : tensor<128xf32> + %expected = check.generate.iota offset(3.25) step(0.01) : tensor<128xf32> + kernel.launch @ggml_rmsnorm_mul_add_f32[%token_count](%token_count, %input, %weight, %residual, %output) : [index](index, tensor<128xf32>, tensor<128xf32>, tensor<128xf32>, tensor<128xf32>) + check.expect.close actual(%output) expected(%expected) atol(9.9999999999999991e-06) rtol(9.9999999999999991e-06) nan(same) : tensor<128xf32> + check.return +} + +check.benchmark<@ggml_rmsnorm_mul_add_f32_case> @ggml_rmsnorm_mul_add_f32_decode diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_mul_rope_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_mul_rope_f32.loom index 08b4b4001ef9..e12c1dc82681 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_mul_rope_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_mul_rope_f32.loom @@ -21,7 +21,7 @@ config.decl @ggml.rmsnorm_mul_rope.section0 : %value: index where [range(%value, config.decl @ggml.rmsnorm_mul_rope.section1 : %value: index where [range(%value, 0, 65536)] config.decl @ggml.rmsnorm_mul_rope.section2 : %value: index where [range(%value, 0, 65536)] config.decl @ggml.rmsnorm_mul_rope.section3 : %value: index where [range(%value, 0, 65536)] -config.decl @ggml.rmsnorm_mul_rope.mode : %value: index where [range(%value, 8, 40)] +config.decl @ggml.rmsnorm_mul_rope.mode : %value: index where [range(%value, 2, 40)] config.decl @ggml.rmsnorm_mul_rope.workgroup_size : %value: index where [range(%value, 32, 1024), mul(%value, 32)] config.decl @ggml.rmsnorm_mul_rope.epsilon : f32 config.decl @ggml.rmsnorm_mul_rope.freq_base : f32 diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rope/rope_concat_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rope/rope_concat_f32.loom new file mode 100644 index 000000000000..b7e03b4ab1ad --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rope/rope_concat_f32.loom @@ -0,0 +1,114 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +// Publishes an unrotated prefix and a partially rotated suffix directly into +// one packed head. The source remains in its original strided full-head layout. +template.decl @ggml.rope_f32.rotate_neox_packet(%position: f32, %input_base: index, %pair_channel: index, %half_n_dims: index, %bounded_input_span: index, %mscale: f32, %input: buffer, %theta_packet: vector<2xf32>, %freq_factors_packet: vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) +template.decl @ggml.rope_f32.rotate_normal_packet(%position: f32, %input_base: index, %normal_channel: index, %bounded_input_span: index, %mscale: f32, %input: buffer, %theta_packet: vector<2xf32>, %freq_factors_packet: vector<2xf32>) -> (vector<4xf32>) +template.decl @ggml.rope_f32.copy_packet(%input_base: index, %normal_channel: index, %bounded_input_span: index, %input: buffer) -> (vector<4xf32>) + +amdgpu.target @ggml_rope_concat_f32_gfx11_wave32 {subgroup_size = 32} + +config.decl @ggml.rope_concat_f32.prefix_size : %value: index where [range(%value, 4, 1024), mul(%value, 4)] +config.decl @ggml.rope_concat_f32.rope_size : %value: index where [range(%value, 4, 1024), mul(%value, 4)] +config.decl @ggml.rope_concat_f32.n_dims : %value: index where [range(%value, 4, 1024), mul(%value, 4)] +config.decl @ggml.rope_concat_f32.head_count : %value: index where [range(%value, 1, 64)] +config.decl @ggml.rope_concat_f32.token_capacity : %value: index where [range(%value, 1, 2048)] +config.decl @ggml.rope_concat_f32.input_stride1 : %value: index where [range(%value, 4, 1073741824)] +config.decl @ggml.rope_concat_f32.input_stride2 : %value: index where [range(%value, 4, 1073741824)] +config.decl @ggml.rope_concat_f32.input_span : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.rope_concat_f32.mscale : f32 +config.decl @ggml.rope_concat_f32.mode : %value: index where [range(%value, 0, 2)] + +kernel.def target(@ggml_rope_concat_f32_gfx11_wave32) export("ggml_rope_concat_f32") @ggml_rope_concat_f32(%token_count: index) { + %token_capacity = config.get @ggml.rope_concat_f32.token_capacity : index + %head_count = config.get @ggml.rope_concat_f32.head_count : index + %prefix_size = config.get @ggml.rope_concat_f32.prefix_size : index + %rope_size = config.get @ggml.rope_concat_f32.rope_size : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %output_size = index.add %prefix_size, %rope_size : index + %packet_count = index.div %output_size, %c4 : index + kernel.launch.config workgroups(%head_count, %token_capacity, %c1) workgroup_size(%packet_count, %c1, %c1) : index +} launch(%token_count: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + %token_capacity = config.get @ggml.rope_concat_f32.token_capacity : index + %head_count0 = config.get @ggml.rope_concat_f32.head_count : index + %prefix_size0 = config.get @ggml.rope_concat_f32.prefix_size : index + %rope_size0 = config.get @ggml.rope_concat_f32.rope_size : index + %n_dims0 = config.get @ggml.rope_concat_f32.n_dims : index + %input_stride1_0 = config.get @ggml.rope_concat_f32.input_stride1 : index + %input_stride2_0 = config.get @ggml.rope_concat_f32.input_stride2 : index + %input_span0 = config.get @ggml.rope_concat_f32.input_span : index + %mscale = config.get @ggml.rope_concat_f32.mscale : f32 + %mode = config.get @ggml.rope_concat_f32.mode : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %bounded_input_span = index.assume %input_span0 [range(%input_span0, 1, 1073741824)] : index + %head_count, %prefix_size, %rope_size, %n_dims = index.assume %head_count0, %prefix_size0, %rope_size0, %n_dims0 [range(%head_count0, 1, 64), range(%prefix_size0, 4, 1024), range(%rope_size0, 4, 1024), range(%n_dims0, 4, 1024), mul(%prefix_size0, 4), mul(%rope_size0, 4), mul(%n_dims0, 4), le(%n_dims0, %rope_size0)] : index, index, index, index + %input_stride1, %input_stride2 = index.assume %input_stride1_0, %input_stride2_0 [range(%input_stride1_0, 4, 1073741824), range(%input_stride2_0, 4, 1073741824)] : index, index + %output_size0 = index.add %prefix_size, %rope_size : index + %output_size = index.assume %output_size0 [range(%output_size0, 8, 1024), mul(%output_size0, 4)] : index + %head = kernel.workgroup.id : index + %token0 = kernel.workgroup.id : index + %packet0 = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %mode_neox = index.constant 2 : index + %c0_offset = index.constant 0 : offset + %prefix_packet_count = index.div %prefix_size, %c4 : index + %rotary_packet_count = index.div %n_dims, %c4 : index + %packet_count = index.div %output_size, %c4 : index + %valid_token = index.cmp ult, %token0, %bounded_token_count : index + %active_packet = index.cmp ult, %packet0, %packet_count : index + %publish = scalar.andi %valid_token, %active_packet : i1 + %token = scf.select %valid_token, %token0, %c0 : index + %is_prefix = index.cmp ult, %packet0, %prefix_packet_count : index + %rope_packet0 = index.sub %packet0, %prefix_packet_count : index + %rope_packet = scf.select %is_prefix, %c0, %rope_packet0 : index + %rotary_packet = index.cmp ult, %rope_packet, %rotary_packet_count : index + %tail_packet = index.cmp uge, %rope_packet, %rotary_packet_count : index + %is_neox = index.cmp eq, %mode, %mode_neox : index + %positions_noalias, %input_noalias, %theta_noalias, %freq_factors_noalias, %output_noalias = buffer.assume.noalias %positions, %input, %theta, %freq_factors, %output : buffer, buffer, buffer, buffer, buffer + scf.if %publish { + %packet = index.assume %packet0 [lt(%packet0, %packet_count)] : index + %positions_view = buffer.view %positions_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]xi32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%head_count]x[%output_size]xf32> + %input_token_offset = index.mul %token, %input_stride2 : index + %input_head_offset = index.mul %head, %input_stride1 : index + %input_base = index.add %input_token_offset, %input_head_offset : index + %output_channel = index.mul %packet, %c4 : index + scf.if %is_prefix { + %values = template.apply<@ggml.rope_f32.copy_packet>(%input_base, %output_channel, %bounded_input_span, %input_noalias) : (index, index, index, buffer) -> (vector<4xf32>) + vector.store %values, %output_view[%token, %head, %output_channel] : vector<4xf32>, view<[%bounded_token_count]x[%head_count]x[%output_size]xf32> + } else { + %rope_input_base = index.add %input_base, %prefix_size : index + %rope_channel = index.mul %rope_packet, %c4 : index + scf.if %rotary_packet { + %half_n_dims = index.div %n_dims, %c2 : index + %pair_channel = index.mul %rope_packet, %c2 : index + %theta_view = buffer.view %theta_noalias[%c0_offset] : buffer -> view<[%half_n_dims]xf32> + %freq_factors_view = buffer.view %freq_factors_noalias[%c0_offset] : buffer -> view<[%half_n_dims]xf32> + %position_i32 = view.load %positions_view[%token] : view<[%bounded_token_count]xi32> -> i32 + %position = scalar.sitofp %position_i32 : i32 to f32 + %theta_packet = vector.load %theta_view[%pair_channel] : view<[%half_n_dims]xf32> -> vector<2xf32> + %freq_factors_packet = vector.load %freq_factors_view[%pair_channel] : view<[%half_n_dims]xf32> -> vector<2xf32> + scf.if %is_neox { + %paired_channel = index.add %pair_channel, %half_n_dims : index + %scaled_low, %scaled_high = template.apply<@ggml.rope_f32.rotate_neox_packet>(%position, %rope_input_base, %pair_channel, %half_n_dims, %bounded_input_span, %mscale, %input_noalias, %theta_packet, %freq_factors_packet) : (f32, index, index, index, index, f32, buffer, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) + %low_output_channel = index.add %prefix_size, %pair_channel : index + %high_output_channel = index.add %prefix_size, %paired_channel : index + vector.store %scaled_low, %output_view[%token, %head, %low_output_channel] : vector<2xf32>, view<[%bounded_token_count]x[%head_count]x[%output_size]xf32> + vector.store %scaled_high, %output_view[%token, %head, %high_output_channel] : vector<2xf32>, view<[%bounded_token_count]x[%head_count]x[%output_size]xf32> + } else { + %rotated = template.apply<@ggml.rope_f32.rotate_normal_packet>(%position, %rope_input_base, %rope_channel, %bounded_input_span, %mscale, %input_noalias, %theta_packet, %freq_factors_packet) : (f32, index, index, index, f32, buffer, vector<2xf32>, vector<2xf32>) -> (vector<4xf32>) + vector.store %rotated, %output_view[%token, %head, %output_channel] : vector<4xf32>, view<[%bounded_token_count]x[%head_count]x[%output_size]xf32> + } + } + scf.if %tail_packet { + %values = template.apply<@ggml.rope_f32.copy_packet>(%rope_input_base, %rope_channel, %bounded_input_span, %input_noalias) : (index, index, index, buffer) -> (vector<4xf32>) + vector.store %values, %output_view[%token, %head, %output_channel] : vector<4xf32>, view<[%bounded_token_count]x[%head_count]x[%output_size]xf32> + } + } + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rope/rope_scale_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rope/rope_scale_f32.loom new file mode 100644 index 000000000000..0f059de07ff7 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rope/rope_scale_f32.loom @@ -0,0 +1,37 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.rope_scale_f32.body(%token_count: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %output: buffer) +template.decl @ggml.rope_token_scale_f32.body(%token_count: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %scales: buffer, %output: buffer) + +amdgpu.target @ggml_rope_scale_f32_gfx11_wave32 {subgroup_size = 32} + +config.decl @ggml.rope_f32.head_size : %value: index where [range(%value, 4, 1024), mul(%value, 4)] +config.decl @ggml.rope_f32.head_count : %value: index where [range(%value, 1, 64)] +config.decl @ggml.rope_f32.token_capacity : %value: index where [range(%value, 1, 2048)] + +kernel.def target(@ggml_rope_scale_f32_gfx11_wave32) export("ggml_rope_scale_f32") @ggml_rope_scale_f32(%token_count: index) { + %token_capacity = config.get @ggml.rope_f32.token_capacity : index + %head_count = config.get @ggml.rope_f32.head_count : index + %head_size = config.get @ggml.rope_f32.head_size : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %pair_packet_count = index.div %head_size, %c4 : index + kernel.launch.config workgroups(%head_count, %token_capacity, %c1) workgroup_size(%pair_packet_count, %c1, %c1) : index +} launch(%token_count: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + template.apply<@ggml.rope_scale_f32.body>(%token_count, %positions, %input, %theta, %freq_factors, %output) : (index, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@ggml_rope_scale_f32_gfx11_wave32) export("ggml_rope_token_scale_f32") @ggml_rope_token_scale_f32(%token_count: index) { + %token_capacity = config.get @ggml.rope_f32.token_capacity : index + %head_count = config.get @ggml.rope_f32.head_count : index + %head_size = config.get @ggml.rope_f32.head_size : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %pair_packet_count = index.div %head_size, %c4 : index + kernel.launch.config workgroups(%head_count, %token_capacity, %c1) workgroup_size(%pair_packet_count, %c1, %c1) : index +} launch(%token_count: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %scales: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + template.apply<@ggml.rope_token_scale_f32.body>(%token_count, %positions, %input, %theta, %freq_factors, %scales, %output) : (index, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/gather_add_rmsnorm_binary_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/gather_add_rmsnorm_binary_f32.loom new file mode 100644 index 000000000000..7434e3648561 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/gather_add_rmsnorm_binary_f32.loom @@ -0,0 +1,133 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +// Publishes both the indexed residual and its RMS-normalized learned-scale +// output. Each output row may select any valid source row independently. +template.decl @ggml.binary_f32.apply(%op: index, %lhs: f32, %rhs: f32) -> (f32) + +amdgpu.target @ggml_gather_add_rmsnorm_gfx11_wave32 {subgroup_size = 32} + +config.decl @ggml.gather_add_rmsnorm_binary_f32.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] +config.decl @ggml.gather_add_rmsnorm_binary_f32.rms_epsilon : f32 +config.decl @ggml.gather_add_rmsnorm_binary_f32.gather_op : %value: index where [range(%value, 0, 3)] + +kernel.def target(@ggml_gather_add_rmsnorm_gfx11_wave32) export("ggml_gather_add_rmsnorm_binary_f32") @ggml_gather_add_rmsnorm_binary_f32(%source_token_count: index, %output_token_count: index) { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%output_token_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%source_token_count: index, %output_token_count: index, %attention: buffer, %residual: buffer, %output_ids: buffer, %raw_output: buffer, %weight: buffer, %normalized_output: buffer) where [range(%source_token_count, 1, 2048), range(%output_token_count, 1, 2048)] { + %source_count = index.assume %source_token_count [range(%source_token_count, 1, 2048)] : index + %output_count = index.assume %output_token_count [range(%output_token_count, 1, 2048)] : index + %hidden0 = config.get @ggml.gather_add_rmsnorm_binary_f32.hidden_size : index + %hidden = index.assume %hidden0 [range(%hidden0, 128, 32768), mul(%hidden0, 128)] : index + %epsilon = config.get @ggml.gather_add_rmsnorm_binary_f32.rms_epsilon : f32 + %gather_op = config.get @ggml.gather_add_rmsnorm_binary_f32.gather_op : index + %output_row0 = kernel.workgroup.id : index + %output_row = index.assume %output_row0 [lt(%output_row0, %output_count)] : index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c8 = index.constant 8 : index + %c256 = index.constant 256 : index + %base = index.constant 0 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %scratch_bytes = index.constant 1024 : offset + + %attention_noalias, %residual_noalias, %ids_noalias, %raw_noalias, %weight_noalias, %normalized_noalias = buffer.assume.noalias %attention, %residual, %output_ids, %raw_output, %weight, %normalized_output : buffer, buffer, buffer, buffer, buffer, buffer + %attention_view = buffer.view %attention_noalias[%base] : buffer -> view<[%source_count]x[%hidden]xf32> + %residual_view = buffer.view %residual_noalias[%base] : buffer -> view<[%source_count]x[%hidden]xf32> + %ids_view = buffer.view %ids_noalias[%base] : buffer -> view<[%output_count]xi32> + %raw_view = buffer.view %raw_noalias[%base] : buffer -> view<[%output_count]x[%hidden]xf32> + %weight_view = buffer.view %weight_noalias[%base] : buffer -> view<[%hidden]xf32> + %normalized_view = buffer.view %normalized_noalias[%base] : buffer -> view<[%output_count]x[%hidden]xf32> + %source_row_i32 = view.load %ids_view[%output_row] : view<[%output_count]xi32> -> i32 + %source_row0 = index.cast %source_row_i32 : i32 to index + %source_row = index.assume %source_row0 [range(%source_row0, 0, 2047), lt(%source_row0, %source_count)] : index + + %thread_sum = scf.for %channel = [%workitem to %hidden step %c256](%running_sum = %c0_f32 : f32) -> (f32) { + %attention_value = view.load %attention_view[%source_row, %channel] : view<[%source_count]x[%hidden]xf32> -> f32 + %residual_value = view.load %residual_view[%source_row, %channel] : view<[%source_count]x[%hidden]xf32> -> f32 + %combined = template.apply<@ggml.binary_f32.apply>(%gather_op, %attention_value, %residual_value) : (index, f32, f32) -> (f32) + view.store %combined, %raw_view[%output_row, %channel] : f32, view<[%output_count]x[%hidden]xf32> + %square = scalar.mulf %combined, %combined : f32 + %next_sum = scalar.addf %running_sum, %square : f32 + scf.yield %next_sum : f32 + } + %subgroup_sum = kernel.subgroup.reduce %thread_sum : f32 + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_view = buffer.view %scratch[%base] : buffer -> view<256xf32> + %is_subgroup_leader = index.cmp eq, %lane, %c0 : index + scf.if %is_subgroup_leader { + view.store %subgroup_sum, %scratch_view[%subgroup] : f32, view<256xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %is_reduction_subgroup = index.cmp eq, %subgroup, %c0 : index + %is_reduction_lane = index.cmp ult, %lane, %c8 : index + %loads_subgroup_sum = scalar.andi %is_reduction_subgroup, %is_reduction_lane : i1 + %subgroup_partial = scf.if %loads_subgroup_sum -> (f32) { + %value = view.load %scratch_view[%lane] : view<256xf32> -> f32 + scf.yield %value : f32 + } else { + scf.yield %c0_f32 : f32 + } + %row_sum = kernel.subgroup.reduce %subgroup_partial : f32 + %writes_scale = scalar.andi %is_reduction_subgroup, %is_subgroup_leader : i1 + scf.if %writes_scale { + %hidden_i32 = index.cast %hidden : index to i32 + %hidden_f32 = scalar.sitofp %hidden_i32 : i32 to f32 + %mean = scalar.divf %row_sum, %hidden_f32 : f32 + %biased_mean = scalar.addf %mean, %epsilon : f32 + %scale = scalar.rsqrtf %biased_mean : f32 + view.store %scale, %scratch_view[%c0] : f32, view<256xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %row_scale = view.load %scratch_view[%c0] : view<256xf32> -> f32 + scf.for %channel = [%workitem to %hidden step %c256] { + %value = view.load %raw_view[%output_row, %channel] : view<[%output_count]x[%hidden]xf32> -> f32 + %learned_weight = view.load %weight_view[%channel] : view<[%hidden]xf32> -> f32 + %normalized = scalar.mulf %value, %row_scale : f32 + %result = scalar.mulf %normalized, %learned_weight : f32 + view.store %result, %normalized_view[%output_row, %channel] : f32, view<[%output_count]x[%hidden]xf32> + } + kernel.return +} + +check.case public @ggml_gather_add_rmsnorm_binary_f32_index_case { + %source_count = check.literal value(3) : index + %output_count = check.literal value(3) : index + %attention = check.generate.fill value(1.0) : tensor<3x128xf32> + %residual = check.generate.fill value(3.0) : tensor<3x128xf32> + %ids = check.generate.iota offset(2) step(-1) : tensor<3xi32> + %raw = check.generate.fill value(0.0) : tensor<3x128xf32> + %weight = check.generate.fill value(2.0) : tensor<128xf32> + %normalized = check.generate.fill value(0.0) : tensor<3x128xf32> + %expected_raw = check.generate.fill value(4.0) : tensor<3x128xf32> + %expected_normalized = check.generate.fill value(2.0) : tensor<3x128xf32> + kernel.launch @ggml_gather_add_rmsnorm_binary_f32[%source_count, %output_count](%source_count, %output_count, %attention, %residual, %ids, %raw, %weight, %normalized) : [index, index](index, index, tensor<3x128xf32>, tensor<3x128xf32>, tensor<3xi32>, tensor<3x128xf32>, tensor<128xf32>, tensor<3x128xf32>) + check.expect.close actual(%raw) expected(%expected_raw) atol(0.0) rtol(0.0) nan(same) : tensor<3x128xf32> + check.expect.close actual(%normalized) expected(%expected_normalized) atol(0.000001) rtol(0.000001) nan(same) : tensor<3x128xf32> + check.return +} + +check.benchmark<@ggml_gather_add_rmsnorm_binary_f32_index_case> @ggml_gather_add_rmsnorm_binary_f32_index + +check.case public @ggml_gather_add_rmsnorm_binary_f32_repeated_index_case { + %source_count = check.literal value(3) : index + %output_count = check.literal value(3) : index + %attention = check.generate.fill value(1.0) : tensor<3x128xf32> + %residual = check.generate.fill value(3.0) : tensor<3x128xf32> + %ids = check.generate.fill value(2) : tensor<3xi32> + %raw = check.generate.fill value(0.0) : tensor<3x128xf32> + %weight = check.generate.fill value(2.0) : tensor<128xf32> + %normalized = check.generate.fill value(0.0) : tensor<3x128xf32> + %expected_raw = check.generate.fill value(4.0) : tensor<3x128xf32> + %expected_normalized = check.generate.fill value(2.0) : tensor<3x128xf32> + kernel.launch @ggml_gather_add_rmsnorm_binary_f32[%source_count, %output_count](%source_count, %output_count, %attention, %residual, %ids, %raw, %weight, %normalized) : [index, index](index, index, tensor<3x128xf32>, tensor<3x128xf32>, tensor<3xi32>, tensor<3x128xf32>, tensor<128xf32>, tensor<3x128xf32>) + check.expect.close actual(%raw) expected(%expected_raw) atol(0.0) rtol(0.0) nan(same) : tensor<3x128xf32> + check.expect.close actual(%normalized) expected(%expected_normalized) atol(0.000001) rtol(0.000001) nan(same) : tensor<3x128xf32> + check.return +} + +check.benchmark<@ggml_gather_add_rmsnorm_binary_f32_repeated_index_case> @ggml_gather_add_rmsnorm_binary_f32_repeated_index diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/get_rows_scale_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/get_rows_scale_f32.loom new file mode 100644 index 000000000000..b229d9968b9b --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/get_rows_scale_f32.loom @@ -0,0 +1,90 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.get_rows_f32.load_vector4_common(%weight_format: index, %publish: i1, %valid_token_id: i1, %token: index, %row_count: index, %hidden_size: index, %channel: index, %weight: buffer) -> (vector<4xf32>) + +amdgpu.target @ggml_get_rows_scale_f32_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.get_rows_f32.token_capacity : %value: index where [range(%value, 1, 2048)] +config.decl @ggml.get_rows_f32.hidden_capacity : %value: index where [range(%value, 4, 1073741824), mul(%value, 4)] +config.decl @ggml.get_rows_f32.weight_format : %value: index where [range(%value, 4, 81)] +config.decl @ggml.get_rows_scale_f32.scale : f32 +config.decl @ggml.get_rows_scale_f32.bias : f32 + +kernel.def target(@ggml_get_rows_scale_f32_gfx11_wave64) export("ggml_get_rows_scale_f32") @ggml_get_rows_scale_f32(%token_count: index, %row_count: index, %hidden_size: index) { + %token_capacity = config.get @ggml.get_rows_f32.token_capacity : index + %hidden_capacity = config.get @ggml.get_rows_f32.hidden_capacity : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %packet_count = index.div %hidden_capacity, %c4 : index + %padded_packet_count = index.add %packet_count, %c255 : index + %packet_tiles = index.div %padded_packet_count, %c256 : index + kernel.launch.config workgroups(%packet_tiles, %token_capacity, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %row_count: index, %hidden_size: index, %token_ids: buffer, %weight: buffer, %raw_output: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + %token_capacity = config.get @ggml.get_rows_f32.token_capacity : index + %hidden_capacity = config.get @ggml.get_rows_f32.hidden_capacity : index + %weight_format = config.get @ggml.get_rows_f32.weight_format : index + %scale = config.get @ggml.get_rows_scale_f32.scale : f32 + %bias = config.get @ggml.get_rows_scale_f32.bias : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %bounded_row_count = index.assume %row_count [range(%row_count, 1, 262208)] : index + %bounded_hidden_size = index.assume %hidden_size [range(%hidden_size, 4, 1073741824), mul(%hidden_size, 4), le(%hidden_size, %hidden_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token0 = kernel.workgroup.id : index + %workitem0 = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c4 = index.constant 4 : index + %c256 = index.constant 256 : index + %c0_i32 = scalar.constant 0 : i32 + %c0_offset = index.constant 0 : offset + %workitem = index.assume %workitem0 [range(%workitem0, 0, 255)] : index + %packet_tile_base = index.mul %channel_tile, %c256 : index + %packet0 = index.add %packet_tile_base, %workitem : index + %packet_count = index.div %bounded_hidden_size, %c4 : index + %valid_token = index.cmp ult, %token0, %bounded_token_count : index + %valid_packet = index.cmp ult, %packet0, %packet_count : index + %publish = scalar.andi %valid_token, %valid_packet : i1 + %token = scf.select %valid_token, %token0, %c0 : index + %packet = scf.select %valid_packet, %packet0, %c0 : index + %token_ids_na, %weight_na, %raw_output_na, %output_na = buffer.assume.noalias %token_ids, %weight, %raw_output, %output : buffer, buffer, buffer, buffer + %token_ids_view = buffer.view %token_ids_na[%c0_offset] : buffer -> view<[%bounded_token_count]xi32> + %token_id_raw = view.load %token_ids_view[%token] : view<[%bounded_token_count]xi32> -> i32 + %token_id_nonnegative = scalar.cmpi sge, %token_id_raw, %c0_i32 : i32 + %safe_token_id_i32 = scf.select %token_id_nonnegative, %token_id_raw, %c0_i32 : i32 + %token_id = index.cast %safe_token_id_i32 : i32 to index + %channel0 = index.mul %packet, %c4 : index + %channel_end0 = index.add %channel0, %c4 : index + %channel, %channel_end = index.assume %channel0, %channel_end0 [le(%channel_end0, %bounded_hidden_size)] : index, index + %values = template.apply<@ggml.get_rows_f32.load_vector4_common>(%weight_format, %publish, %token_id_nonnegative, %token_id, %bounded_row_count, %bounded_hidden_size, %channel, %weight_na) : (index, i1, i1, index, index, index, index, buffer) -> (vector<4xf32>) + %scale_vector = vector.splat %scale : vector<4xf32> + %bias_vector = vector.splat %bias : vector<4xf32> + %scaled = vector.mulf %values, %scale_vector : vector<4xf32> + %result = vector.addf %scaled, %bias_vector : vector<4xf32> + scf.if %publish { + %raw_output_view = buffer.view %raw_output_na[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_hidden_size]xf32> + %output_view = buffer.view %output_na[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_hidden_size]xf32> + vector.store %values, %raw_output_view[%token, %channel] : vector<4xf32>, view<[%bounded_token_count]x[%bounded_hidden_size]xf32> + vector.store %result, %output_view[%token, %channel] : vector<4xf32>, view<[%bounded_token_count]x[%bounded_hidden_size]xf32> + } + kernel.return +} + +check.case public @ggml_get_rows_scale_f32_f32_case { + %token_count = check.literal value(4) : index + %row_count = check.literal value(4) : index + %hidden_size = check.literal value(4) : index + %token_ids = check.generate.iota offset(3) step(-1) : tensor<4xi32> + %weight = check.generate.fill value(2.0) : tensor<4x4xf32> + %raw_output = check.generate.fill value(0.0) : tensor<4x4xf32> + %output = check.generate.fill value(0.0) : tensor<4x4xf32> + %expected_raw = check.generate.fill value(2.0) : tensor<4x4xf32> + %expected = check.generate.fill value(4.25) : tensor<4x4xf32> + kernel.launch @ggml_get_rows_scale_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %raw_output, %output) : [index, index, index](index, index, index, tensor<4xi32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>) + check.expect.close actual(%raw_output) expected(%expected_raw) atol(0.0) rtol(0.0) nan(same) : tensor<4x4xf32> + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<4x4xf32> + check.return +} + +check.benchmark<@ggml_get_rows_scale_f32_f32_case> @ggml_get_rows_scale_f32_f32 diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/ssm_conv/lfm_dconv3_state_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/ssm_conv/lfm_dconv3_state_f32.loom new file mode 100644 index 000000000000..a0ffcb1a08cd --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/ssm_conv/lfm_dconv3_state_f32.loom @@ -0,0 +1,128 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +amdgpu.target @llm_ssm_conv_lfm_dconv3_gfx11_wave32 {subgroup_size = 32} + +config.decl @llm.ssm_conv.lfm_dconv3.n_t : %value: index where [range(%value, 1, 64)] + +kernel.def target(@llm_ssm_conv_lfm_dconv3_gfx11_wave32) export("llm_ssm_conv_lfm_dconv3_state_f32") @llm_ssm_conv_lfm_dconv3_state_f32() { + %c1 = index.constant 1 : index + %c8 = index.constant 8 : index + %c256 = index.constant 256 : index + %n_t = config.get @llm.ssm_conv.lfm_dconv3.n_t : index + %groups = index.mul %n_t, %c8 : index + kernel.launch.config workgroups(%groups, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%state: buffer, %x: buffer, %filter: buffer, %output: buffer, %cache: buffer) { + %base = index.constant 0 : offset + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c8 = index.constant 8 : index + %c2048 = index.constant 2048 : index + %c256 = index.constant 256 : index + %n_t0 = config.get @llm.ssm_conv.lfm_dconv3.n_t : index + %n_t = index.assume %n_t0 [range(%n_t0, 1, 64)] : index + %group0 = kernel.workgroup.id : index + %lane0 = kernel.workitem.id : index + %group = index.assume %group0 [range(%group0, 0, 511)] : index + %lane = index.assume %lane0 [range(%lane0, 0, 255)] : index + %token0 = index.div %group, %c8 : index + %token = index.assume %token0 [range(%token0, 0, 63)] : index + %channel_group = index.rem %group, %c8 : index + %channel_base = index.mul %channel_group, %c256 : index + %channel0 = index.add %channel_base, %lane : index + %channel = index.assume %channel0 [range(%channel0, 0, 2047)] : index + %token_elements0 = index.mul %n_t, %c2048 : index + %token_elements = index.assume %token_elements0 [range(%token_elements0, 2048, 131072), mul(%token_elements0, 2048)] : index + + %state_v = buffer.view %state[%base] : buffer -> view<4096xf32> + %filter_v = buffer.view %filter[%base] : buffer -> view<6144xf32> + %cache_v = buffer.view %cache[%base] : buffer -> view<4096xf32> + %state_base0 = index.mul %channel, %c2 : index + %state_base = index.assume %state_base0 [range(%state_base0, 0, 4094), mul(%state_base0, 2)] : index + %filter_base0 = index.mul %channel, %c3 : index + %filter_base = index.assume %filter_base0 [range(%filter_base0, 0, 6141), mul(%filter_base0, 3)] : index + %filter_one0 = index.add %filter_base, %c1 : index + %filter_one = index.assume %filter_one0 [range(%filter_one0, 1, 6142)] : index + %filter_two0 = index.add %filter_base, %c2 : index + %filter_two = index.assume %filter_two0 [range(%filter_two0, 2, 6143)] : index + %w0 = view.load %filter_v[%filter_base] : view<6144xf32> -> f32 + %w1 = view.load %filter_v[%filter_one] : view<6144xf32> -> f32 + %w2 = view.load %filter_v[%filter_two] : view<6144xf32> -> f32 + + %from_x0 = index.cmp uge, %token, %c2 : index + %v0 = scf.if %from_x0 -> (f32) { + %x_token = index.sub %token, %c2 : index + %x_base = index.mul %x_token, %c2048 : index + %x_index0 = index.add %x_base, %channel : index + %x_index, %x_extent = index.assume %x_index0, %token_elements [lt(%x_index0, %token_elements)] : index, index + %x_view = buffer.view %x[%base] : buffer -> view<[%x_extent]xf32> + %loaded = view.load %x_view[%x_index] : view<[%x_extent]xf32> -> f32 + scf.yield %loaded : f32 + } else { + %state_index0 = index.add %state_base, %token : index + %state_index = index.assume %state_index0 [range(%state_index0, 0, 4095)] : index + %loaded = view.load %state_v[%state_index] : view<4096xf32> -> f32 + scf.yield %loaded : f32 + } + %position1 = index.add %token, %c1 : index + %from_x1 = index.cmp uge, %position1, %c2 : index + %v1 = scf.if %from_x1 -> (f32) { + %x_token = index.sub %position1, %c2 : index + %x_base = index.mul %x_token, %c2048 : index + %x_index0 = index.add %x_base, %channel : index + %x_index, %x_extent = index.assume %x_index0, %token_elements [lt(%x_index0, %token_elements)] : index, index + %x_view = buffer.view %x[%base] : buffer -> view<[%x_extent]xf32> + %loaded = view.load %x_view[%x_index] : view<[%x_extent]xf32> -> f32 + scf.yield %loaded : f32 + } else { + %state_index0 = index.add %state_base, %position1 : index + %state_index = index.assume %state_index0 [range(%state_index0, 1, 4095)] : index + %loaded = view.load %state_v[%state_index] : view<4096xf32> -> f32 + scf.yield %loaded : f32 + } + %x_base2 = index.mul %token, %c2048 : index + %x_index20 = index.add %x_base2, %channel : index + %x_index2, %x_extent2 = index.assume %x_index20, %token_elements [lt(%x_index20, %token_elements)] : index, index + %x_view2 = buffer.view %x[%base] : buffer -> view<[%x_extent2]xf32> + %v2 = view.load %x_view2[%x_index2] : view<[%x_extent2]xf32> -> f32 + %product0 = scalar.mulf %v0, %w0 : f32 + %product1 = scalar.mulf %v1, %w1 : f32 + %sum0 = scalar.addf %product0, %product1 : f32 + %product2 = scalar.mulf %v2, %w2 : f32 + %sum1 = scalar.addf %sum0, %product2 : f32 + %output_index, %output_extent = index.assume %x_index20, %token_elements [lt(%x_index20, %token_elements)] : index, index + %output_view = buffer.view %output[%base] : buffer -> view<[%output_extent]xf32> + view.store %sum1, %output_view[%output_index] : f32, view<[%output_extent]xf32> + + %publish = index.cmp eq, %token, %c0 : index + scf.if %publish { + %one_token = index.cmp eq, %n_t, %c1 : index + %tail0, %tail1 = scf.if %one_token -> (f32, f32) { + %state_one0 = index.add %state_base, %c1 : index + %state_one = index.assume %state_one0 [range(%state_one0, 1, 4095)] : index + %old1 = view.load %state_v[%state_one] : view<4096xf32> -> f32 + scf.yield %old1, %v2 : f32, f32 + } else { + %tail_token0 = index.sub %n_t, %c2 : index + %tail_token1 = index.sub %n_t, %c1 : index + %tail_base0 = index.mul %tail_token0, %c2048 : index + %tail_base1 = index.mul %tail_token1, %c2048 : index + %tail_index00 = index.add %tail_base0, %channel : index + %tail_index10 = index.add %tail_base1, %channel : index + %tail_index0, %tail_extent0 = index.assume %tail_index00, %token_elements [lt(%tail_index00, %token_elements)] : index, index + %tail_index1, %tail_extent1 = index.assume %tail_index10, %token_elements [lt(%tail_index10, %token_elements)] : index, index + %tail_view0 = buffer.view %x[%base] : buffer -> view<[%tail_extent0]xf32> + %tail_view1 = buffer.view %x[%base] : buffer -> view<[%tail_extent1]xf32> + %loaded0 = view.load %tail_view0[%tail_index0] : view<[%tail_extent0]xf32> -> f32 + %loaded1 = view.load %tail_view1[%tail_index1] : view<[%tail_extent1]xf32> -> f32 + scf.yield %loaded0, %loaded1 : f32, f32 + } + %cache_one0 = index.add %state_base, %c1 : index + %cache_one = index.assume %cache_one0 [range(%cache_one0, 1, 4095)] : index + view.store %tail0, %cache_v[%state_base] : f32, view<4096xf32> + view.store %tail1, %cache_v[%cache_one] : f32, view<4096xf32> + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen/attention_metadata.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen/attention_metadata.loom index a75a3fa0059d..4eebda0ccda2 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen/attention_metadata.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen/attention_metadata.loom @@ -78,7 +78,7 @@ kernel.def target(@qwen_attention_metadata_gfx11_wave64) export("qwen_attention_ // Publishes the single position and cache row needed by one decode issue. // Decode attention bounds itself with the same request control word and does // not need a materialized causal mask: every prior row is visible. -kernel.def target(@qwen_attention_metadata_gfx11_wave64) export("qwen_decode_attention_metadata") @qwen_decode_attention_metadata() { +kernel.def target(@qwen_attention_metadata_gfx11_wave64) @qwen_decode_attention_metadata() { %c1 = index.constant 1 : index kernel.launch.config workgroups(%c1, %c1, %c1) workgroup_size(%c1, %c1, %c1) : index } launch(%control: buffer, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer) { @@ -182,3 +182,5 @@ check.benchmark<@qwen_attention_metadata_benchmark_case> @qwen_attention_metadat check.benchmark<@qwen_attention_metadata_benchmark_case> @qwen_attention_metadata_prefill_128 {token_count = 128} check.benchmark<@qwen_attention_metadata_benchmark_case> @qwen_attention_metadata_prefill_512 {token_count = 512} + +check.benchmark<@qwen_decode_attention_metadata_sequence_case> @qwen_decode_attention_metadata_sequence diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/ggml/linear_q6k_q8_1_x4.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/ggml/linear_q6k_q8_1_x4.loom index 2cc4cd97a3d8..40218bb61a75 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/ggml/linear_q6k_q8_1_x4.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/ggml/linear_q6k_q8_1_x4.loom @@ -25,275 +25,7 @@ kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count$15: index, %input_size$16: i func.decl @ggml_q8_1_x4_word(%q8_input: buffer, %row_byte_base: offset, %q8_block: index, %word_in_block: index) -> (vector<4xi8>, f32) -// Sign-extends four packed six-bit values without scalar lane extraction. -// Each byte enters with bits [5:0] populated and leaves as signed i8. -func.def inline @ggml_q6k_sign_extend_dot4(%code: vector<1xi32>) -> (vector<4xi8>) { - %c1_i32v = vector.constant 1 : vector<1xi32> - %c2_i32v = vector.constant 2 : vector<1xi32> - %low5_mask = vector.constant 522133279 : vector<1xi32> - %bit5_mask = vector.constant 538976288 : vector<1xi32> - %sign_mask = vector.constant -522133280 : vector<1xi32> - %low5 = vector.andi %code, %low5_mask : vector<1xi32> - %bit5 = vector.andi %code, %bit5_mask : vector<1xi32> - %bit6 = vector.shli %bit5, %c1_i32v : vector<1xi32> - %bit7 = vector.shli %bit5, %c2_i32v : vector<1xi32> - %high01 = vector.ori %bit5, %bit6 : vector<1xi32> - %high = vector.ori %high01, %bit7 : vector<1xi32> - %sign = vector.xori %high, %sign_mask : vector<1xi32> - %signed_i32 = vector.ori %low5, %sign : vector<1xi32> - %signed = vector.bitcast %signed_i32 : vector<1xi32> to vector<4xi8> - func.return %signed : vector<4xi8> -} - -// Decodes four adjacent values from one 32-value Q6_K group for FP16 matrix -// staging. The group and packet coordinates match the contiguous K dimension -// consumed by WMMA tiles, unlike the lane/part mapping used by the Q8_1 dot -// contraction below. -func.def inline @ggml_q6k_f16_vector4(%weight: buffer, %weight_row_byte_base: offset, %q6_block: index, %q6_group: index, %packet: index) -> (vector<4xf16>) { - %c2 = index.constant 2 : index - %c4 = index.constant 4 : index - %c8 = index.constant 8 : index - %c16 = index.constant 16 : index - %block_bytes = index.constant 210 : offset - %qh_byte_add = index.constant 128 : offset - %scale_byte_add = index.constant 192 : offset - %d_byte_add = index.constant 208 : 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 - %block_byte_add = index.scale %q6_block, %block_bytes : index, offset -> offset - %block_byte_base = index.add %weight_row_byte_base, %block_byte_add : offset - %qh_byte_base = index.add %block_byte_base, %qh_byte_add : offset - %scale_byte_base = index.add %block_byte_base, %scale_byte_add : offset - %d_byte_base = index.add %block_byte_base, %d_byte_add : offset - %ql_view = buffer.view %weight[%block_byte_base] : buffer -> view<32xi32> - %qh_view = buffer.view %weight[%qh_byte_base] : buffer -> view<16xi32> - %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<16xi8> - %d_view = buffer.view %weight[%d_byte_base] : buffer -> view<1xf16> - %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 - %ql_half_word_base = index.mul %half, %c16 : index - %ql_side_word_add = index.mul %ql_side, %c8 : index - %ql_word_base = index.add %ql_half_word_base, %ql_side_word_add : index - %ql_word_index = index.add %ql_word_base, %bounded_packet : index - %qh_half_word_base = index.mul %half, %c8 : index - %qh_word_index = index.add %qh_half_word_base, %bounded_packet : 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> - %scale_packet_half = index.div %bounded_packet, %c4 : index - %scale_group_base = index.mul %bounded_group, %c2 : index - %scale_index = index.add %scale_group_base, %scale_packet_half : index - %ql_word = vector.load %ql_view[%ql_word_index] : view<32xi32> -> vector<1xi32> - %qh_word = vector.load %qh_view[%qh_word_index] : view<16xi32> -> 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[%scale_index] : view<16xi8> -> i8 - %d_f16 = view.load %d_view[0] : 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_f32 = vector.mulf %centered, %combined_scale_vector : vector<4xf32> - %values = vector.fptrunc %values_f32 : vector<4xf32> to vector<4xf16> - func.return %values : vector<4xf16> -} - -// Contracts both four-value packets assigned to one lane in a Q6_K block. -// QL, QH, and the block scale are shared by the two packed parts, and their -// contributions are summed before entering the row-wide recurrence. -func.def inline @ggml_q6k_q8_1_x4_block_lane(%weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %q6_block: index, %lane: index) -> (f32) { - %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 - %block_bytes = index.constant 210 : offset - %qh_byte_add = index.constant 128 : offset - %scale_byte_add = index.constant 192 : offset - %d_byte_add = index.constant 208 : offset - %c4_i32v = vector.constant 4 : vector<1xi32> - %c0_i32v = vector.constant 0 : vector<1xi32> - %nibble_mask = vector.constant 252645135 : vector<1xi32> - %high_mask = vector.constant 808464432 : vector<1xi32> - %c0_f32 = scalar.constant 0.0 : f32 - %bounded_lane = index.assume %lane [range(%lane, 0, 31)] : index - %block_byte_add = index.scale %q6_block, %block_bytes : index, offset -> offset - %block_byte_base = index.add %weight_row_byte_base, %block_byte_add : offset - %qh_byte_base = index.add %block_byte_base, %qh_byte_add : offset - %scale_byte_base = index.add %block_byte_base, %scale_byte_add : offset - %d_byte_base = index.add %block_byte_base, %d_byte_add : offset - %ql_view = buffer.view %weight[%block_byte_base] : buffer -> view<32xi32> - %qh_view = buffer.view %weight[%qh_byte_base] : buffer -> view<16xi32> - %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<16xi8> - %d_view = buffer.view %weight[%d_byte_base] : buffer -> view<1xf16> - %lane_mod8 = index.rem %bounded_lane, %c8 : index - %lane_mod16 = index.rem %bounded_lane, %c16 : index - %lane_div16 = index.div %bounded_lane, %c16 : index - %lane_div8_in_16 = index.div %lane_mod16, %c8 : index - %lane_div4_in_16 = index.div %lane_mod16, %c4 : index - %qh_high_base = index.mul %lane_div16, %c8 : index - %qh_index0 = index.add %qh_high_base, %lane_mod8 : index - %qh_index = index.assume %qh_index0 [range(%qh_index0, 0, 15)] : index - %ql_word = vector.load %ql_view[%bounded_lane] : view<32xi32> -> vector<1xi32> - %qh_word = vector.load %qh_view[%qh_index] : view<16xi32> -> vector<1xi32> - %qh_base_shift_index = index.mul %lane_div8_in_16, %c2 : index - %qh_base_shift_i32 = index.cast %qh_base_shift_index : index to i32 - %q8_block_base = index.mul %q6_block, %c8 : index - %q8_high_add = index.mul %lane_div16, %c4 : index - %q8_quadrant = index.add %q8_high_add, %lane_div8_in_16 : index - %scale_high_base = index.mul %lane_div16, %c8 : index - %scale_lane0 = index.add %scale_high_base, %lane_div4_in_16 : index - %d_f16 = view.load %d_view[0] : view<1xf16> -> f16 - %d = scalar.extf %d_f16 : f16 to f32 - %sum = scf.for %part = [%c0 to %c2 step %c1](%accumulator = %c0_f32 : f32) -> (f32) unroll { - %bounded_part = index.assume %part [range(%part, 0, 1)] : index - %part_shift_index = index.mul %bounded_part, %c4 : index - %part_shift_i32 = index.cast %part_shift_index : index to i32 - %part_shift = vector.splat %part_shift_i32 : vector<1xi32> - %ql_shifted = vector.shrui %ql_word, %part_shift : vector<1xi32> - %ql = vector.andi %ql_shifted, %nibble_mask : vector<1xi32> - %qh_shift_i32 = scalar.addi %qh_base_shift_i32, %part_shift_i32 : i32 - %qh_shift = vector.splat %qh_shift_i32 : vector<1xi32> - %qh_shifted = vector.shrui %qh_word, %qh_shift : vector<1xi32> - %qh_positioned = vector.shli %qh_shifted, %c4_i32v : vector<1xi32> - %qh = vector.andi %qh_positioned, %high_mask : vector<1xi32> - %code = vector.ori %ql, %qh : vector<1xi32> - %signed_weight = func.call @ggml_q6k_sign_extend_dot4(%code) : (vector<1xi32>) -> (vector<4xi8>) - %q8_part_add = index.mul %bounded_part, %c2 : index - %q8_block_part = index.add %q8_block_base, %q8_part_add : index - %q8_block = index.add %q8_block_part, %q8_quadrant : index - %q8_values, %q8_d = func.call @ggml_q8_1_x4_word(%q8_input, %q8_row_byte_base, %q8_block, %lane_mod8) : (buffer, offset, index, index) -> (vector<4xi8>, f32) - %scale_lane = index.add %scale_lane0, %part_shift_index : index - %scale_i8 = view.load %scale_view[%scale_lane] : view<16xi8> -> i8 - %scale = scalar.sitofp %scale_i8 : i8 to f32 - %dot = vector.dot4i %signed_weight, %q8_values, %c0_i32v : vector<4xi8>, vector<4xi8>, vector<1xi32> - %dot_i32 = vector.extract %dot[0] : vector<1xi32> -> i32 - %dot_f32 = scalar.sitofp %dot_i32 : i32 to f32 - %scaled0 = scalar.mulf %dot_f32, %scale : f32 - %scaled1 = scalar.mulf %scaled0, %d : f32 - %contribution = scalar.mulf %scaled1, %q8_d : f32 - %next = scalar.addf %accumulator, %contribution : f32 - scf.yield %next : f32 - } - func.return %sum : f32 -} - -// Computes one lane's partial for a complete Q6_K row. Callers choose how -// rows and activations are tiled, then reduce the returned value by subgroup. -func.def inline @ggml_q6k_q8_1_x4_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32) { - %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c256 = index.constant 256 : index - %c0_f32 = scalar.constant 0.0 : f32 - %block_count = index.div %bounded_input_size, %c256 : index - %result = scf.for %block = [%c0 to %block_count step %c1](%block_acc = %c0_f32 : f32) -> (f32) { - %contribution = func.call @ggml_q6k_q8_1_x4_block_lane(%weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %block, %lane) : (buffer, offset, buffer, offset, index, index) -> (f32) - %next = scalar.addf %block_acc, %contribution : f32 - scf.yield %next : f32 - } - func.return %result : f32 -} - -// Maps four Q6_K blocks across the four 16-lane partitions of a wave64. Each -// physical lane contracts the pair of virtual wave32 lanes that own the same -// packed positions in the low and high halves of one block. The subgroup -// reduction therefore combines four complete blocks per loop iteration while -// preserving the canonical wave32 unpacking and dot-product primitive. -func.def inline @ggml_q6k_q8_1_x4_row_lane_wave64_block4(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32) { - %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index - %bounded_lane = index.assume %lane [range(%lane, 0, 63)] : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c3 = index.constant 3 : index - %c4 = index.constant 4 : index - %c16 = index.constant 16 : index - %c256 = index.constant 256 : index - %c0_f32 = scalar.constant 0.0 : f32 - %block_count = index.div %bounded_input_size, %c256 : index - %padded_block_count = index.add %block_count, %c3 : index - %iteration_count = index.div %padded_block_count, %c4 : index - %block_in_iteration = index.div %bounded_lane, %c16 : index - %virtual_lane0 = index.rem %bounded_lane, %c16 : index - %virtual_lane1 = index.add %virtual_lane0, %c16 : index - %result = scf.for %iteration = [%c0 to %iteration_count step %c1](%row_acc = %c0_f32 : f32) -> (f32) { - %block_base = index.mul %iteration, %c4 : index - %block0 = index.add %block_base, %block_in_iteration : index - %valid_block = index.cmp ult, %block0, %block_count : index - %block_acc = scf.if %valid_block -> (f32) { - %low = func.call @ggml_q6k_q8_1_x4_block_lane(%weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %block0, %virtual_lane0) : (buffer, offset, buffer, offset, index, index) -> (f32) - %high = func.call @ggml_q6k_q8_1_x4_block_lane(%weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %block0, %virtual_lane1) : (buffer, offset, buffer, offset, index, index) -> (f32) - %sum = scalar.addf %low, %high : f32 - scf.yield %sum : f32 - } else { - scf.yield %c0_f32 : f32 - } - %next = scalar.addf %row_acc, %block_acc : f32 - scf.yield %next : f32 - } - func.return %result : f32 -} - -// Contracts one output channel for the calling Q6_K projection kernel. The -// caller owns publication so the same canonical contraction can feed either a -// dense logits tensor or an endpoint reduction. -template.def<@ggml.linear_q6k_q8_1_x4.body> device @ggml_linear_q6k_q8_1_x4_body(%publish_output: i1, %token_count: index, %token0: index, %input_size: index, %output_size: index, %q8_input: buffer, %weight: buffer) -> (f32, index, index, i1, index, index) { - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index - %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index - %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144)] : index - %channel_tile = kernel.workgroup.id : index - %subgroup0 = kernel.subgroup.id : index - %lane = kernel.subgroup.lane.id : index - %c0 = index.constant 0 : index - %c8 = index.constant 8 : index - %c128 = index.constant 128 : index - %c210_bytes = index.constant 210 : offset - %c256 = index.constant 256 : index - %c144_bytes = index.constant 144 : offset - %token, %launch_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index - %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index - %channel_base = index.mul %channel_tile, %c8 : index - %channel = index.add %channel_base, %subgroup : index - %valid_channel = index.cmp ult, %channel, %bounded_output_size : index - %q6_block_count = index.div %bounded_input_size, %c256 : index - %weight_row_bytes = index.scale %q6_block_count, %c210_bytes : index, offset -> offset - %weight_row_byte_base = index.scale %channel, %weight_row_bytes : index, offset -> offset - %q8_group_count = index.div %bounded_input_size, %c128 : index - %q8_row_bytes = index.scale %q8_group_count, %c144_bytes : index, offset -> offset - %q8_row_byte_base = index.scale %token, %q8_row_bytes : index, offset -> offset - %lane_acc = scf.if %publish_output -> (f32) { - %channel_acc = scf.if %valid_channel -> (f32) { - %value = func.call @ggml_q6k_q8_1_x4_row_lane(%bounded_input_size, %weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) - scf.yield %value : f32 - } else { - %c0_f32 = scalar.constant 0.0 : f32 - scf.yield %c0_f32 : f32 - } - scf.yield %channel_acc : f32 - } else { - %c0_f32 = scalar.constant 0.0 : f32 - scf.yield %c0_f32 : f32 - } - %dot = kernel.subgroup.reduce %lane_acc : f32 - template.return %dot, %token, %channel, %valid_channel, %launch_token_count, %bounded_output_size : f32, index, index, i1, index, index -} +func.decl @ggml_q6k_q8_1_x4_row_lane_wave64_block4(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32) // Dense raw-layout baseline. A 256-thread workgroup carries eight independent // wave32 output rows so decode launches enough waves without duplicating Q6_K diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/ggml/quantize_q8_1_x4.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/ggml/quantize_q8_1_x4.loom index 95c96b299276..f0da9d9cf4dd 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/ggml/quantize_q8_1_x4.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/ggml/quantize_q8_1_x4.loom @@ -10,156 +10,6 @@ template.decl @ggml.quantize_q8_1_x4.group_body(%publish_output: i1, %group_coun config.decl @ggml.quantize_q8_1_x4.group_capacity : %value: index where [range(%value, 1, 524288)] -// Loads one logical 32-element Q8_1 block from the four-way physical packing. -// `row_byte_base` addresses the first x4 group for one activation row. -func.def inline @ggml_q8_1_x4_block(%q8_input: buffer, %row_byte_base: offset, %q8_block: index) -> (vector<32xi8>, f32, f32) { - %c1 = index.constant 1 : index - %c2 = index.constant 2 : index - %c4 = index.constant 4 : index - %c8 = index.constant 8 : index - %group_bytes = index.constant 144 : offset - %payload_byte_add = index.constant 16 : offset - %group = index.div %q8_block, %c4 : index - %inner0 = index.rem %q8_block, %c4 : index - %inner = index.assume %inner0 [range(%inner0, 0, 3)] : index - %group_byte_add = index.scale %group, %group_bytes : index, offset -> offset - %group_byte_base = index.add %row_byte_base, %group_byte_add : offset - %payload_byte_base = index.add %group_byte_base, %payload_byte_add : offset - %ds_view = buffer.view %q8_input[%group_byte_base] : buffer -> view<8xf16> - %payload_view = buffer.view %q8_input[%payload_byte_base] : buffer -> view<32xi32> - %d_index = index.mul %inner, %c2 : index - %s_index = index.add %d_index, %c1 : index - %word_index = index.mul %inner, %c8 : index - %d_f16 = view.load %ds_view[%d_index] : view<8xf16> -> f16 - %s_f16 = view.load %ds_view[%s_index] : view<8xf16> -> f16 - %words = vector.load %payload_view[%word_index] : view<32xi32> -> vector<8xi32> - %values = vector.bitcast %words : vector<8xi32> to vector<32xi8> - %d = scalar.extf %d_f16 : f16 to f32 - %s = scalar.extf %s_f16 : f16 to f32 - func.return %values, %d, %s : vector<32xi8>, f32, f32 -} - -// Loads one packed four-value word and its block scale. Dot-product kernels -// use this narrower form so a lane does not load the other seven words owned -// by its subgroup peers. -func.def inline @ggml_q8_1_x4_word(%q8_input: buffer, %row_byte_base: offset, %q8_block: index, %word_in_block: index) -> (vector<4xi8>, f32) { - %c2 = index.constant 2 : index - %c4 = index.constant 4 : index - %c8 = index.constant 8 : index - %group_bytes = index.constant 144 : offset - %payload_byte_add = index.constant 16 : offset - %group = index.div %q8_block, %c4 : index - %inner0 = index.rem %q8_block, %c4 : index - %inner = index.assume %inner0 [range(%inner0, 0, 3)] : index - %word0 = index.assume %word_in_block [range(%word_in_block, 0, 7)] : index - %group_byte_add = index.scale %group, %group_bytes : index, offset -> offset - %group_byte_base = index.add %row_byte_base, %group_byte_add : offset - %payload_byte_base = index.add %group_byte_base, %payload_byte_add : offset - %ds_view = buffer.view %q8_input[%group_byte_base] : buffer -> view<8xf16> - %payload_view = buffer.view %q8_input[%payload_byte_base] : buffer -> view<32xi32> - %d_index = index.mul %inner, %c2 : index - %inner_word_base = index.mul %inner, %c8 : index - %word_index = index.add %inner_word_base, %word0 : index - %d_f16 = view.load %ds_view[%d_index] : view<8xf16> -> f16 - %packed = vector.load %payload_view[%word_index] : view<32xi32> -> vector<1xi32> - %values = vector.bitcast %packed : vector<1xi32> to vector<4xi8> - %d = scalar.extf %d_f16 : f16 to f32 - func.return %values, %d : vector<4xi8>, f32 -} - -// Packs one explicit physical group. Callers provide the complete physical -// group domain and an in-range ordinal, and guarantee that one complete -// 32-lane wave executes after all 128 source values are visible. The uniform -// publication predicate suppresses only destination writes so every lane still -// reaches the workgroup barriers. -template.def<@ggml.quantize_q8_1_x4.group_body> device @ggml_quantize_q8_1_x4_group_body(%publish_output: i1, %group_count0: index, %group0: index, %input: buffer, %output: buffer) { - %group_count, %group = index.assume %group_count0, %group0 [range(%group_count0, 1, 524288), lt(%group0, %group_count0)] : index, index - %lane = kernel.workitem.id : index - %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 - %c32 = index.constant 32 : index - %c128 = index.constant 128 : index - %group_bytes = index.constant 144 : offset - %payload_byte_add = index.constant 16 : offset - %scratch_d_byte_add = index.constant 128 : offset - %scratch_bytes = index.constant 144 : offset - %c0_f32 = scalar.constant 0.0 : f32 - %c1_f32 = scalar.constant 1.0 : f32 - %c127 = scalar.constant 127.0 : f32 - %c0_offset = index.constant 0 : offset - %launched_element_count = index.mul %group_count, %c128 : index - %block_in_group0 = index.div %lane, %c8 : index - %block_in_group = index.assume %block_in_group0 [range(%block_in_group0, 0, 3)] : index - %word_in_block0 = index.rem %lane, %c8 : index - %word_in_block = index.assume %word_in_block0 [range(%word_in_block0, 0, 7)] : index - %group_element_base = index.mul %group, %c128 : index - %block_element_add = index.mul %block_in_group, %c32 : index - %block_word_add = index.mul %block_in_group, %c8 : index - %word_element_add = index.mul %word_in_block, %c4 : index - %input_block_base = index.add %group_element_base, %block_element_add : index - %input_index = index.add %input_block_base, %word_element_add : index - %input_noalias, %output_noalias = buffer.assume.noalias %input, %output : buffer, buffer - %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launched_element_count]xf32> - %input_values = vector.load %input_view[%input_index] : view<[%launched_element_count]xf32> -> vector<4xf32> - %absolute_values = vector.absf %input_values : vector<4xf32> - %thread_max = vector.reduce %absolute_values, %c0_f32 : vector<4xf32>, f32 - %scratch = buffer.alloca align(16) %scratch_bytes : buffer - %scratch_values = buffer.view %scratch[%c0_offset] : buffer -> view<32xf32> - %scratch_d = buffer.view %scratch[%scratch_d_byte_add] : buffer -> view<4xf32> - view.store %thread_max, %scratch_values[%lane] : f32, view<32xf32> - kernel.barrier scope(workgroup) ordering(acq_rel) - %is_cohort_leader = index.cmp eq, %word_in_block, %c0 : index - scf.if %is_cohort_leader { - %cohort_base = index.mul %block_in_group, %c8 : index - %cohort_maxima = vector.load %scratch_values[%cohort_base] : view<32xf32> -> vector<8xf32> - %amax = vector.reduce %cohort_maxima, %c0_f32 : vector<8xf32>, f32 - %d = scalar.divf %amax, %c127 : f32 - view.store %d, %scratch_d[%block_in_group] : f32, view<4xf32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %d = view.load %scratch_d[%block_in_group] : view<4xf32> -> f32 - %d_nonzero = scalar.cmpf one, %d, %c0_f32 : f32 - %d_inverse = scf.if %d_nonzero -> (f32) { - %inverse = scalar.divf %c1_f32, %d : f32 - scf.yield %inverse : f32 - } else { - scf.yield %c0_f32 : f32 - } - %d_inverse_vector = vector.splat %d_inverse : vector<4xf32> - %scaled_values = vector.mulf %input_values, %d_inverse_vector : vector<4xf32> - %rounded_values = vector.roundf %scaled_values : vector<4xf32> - %quantized_values = vector.fptosi %rounded_values : vector<4xf32> to vector<4xi8> - %packed_word = vector.bitcast %quantized_values : vector<4xi8> to vector<1xi32> - %group_byte_offset = index.scale %group, %group_bytes : index, offset -> offset - %payload_byte_offset = index.add %group_byte_offset, %payload_byte_add : offset - %group_ds = buffer.view %output_noalias[%group_byte_offset] : buffer -> view<8xf16> - %group_qs = buffer.view %output_noalias[%payload_byte_offset] : buffer -> view<32xi32> - %packed_word_index0 = index.add %block_word_add, %word_in_block : index - %packed_word_index = index.assume %packed_word_index0 [range(%packed_word_index0, 0, 31)] : index - scf.if %publish_output { - vector.store %packed_word, %group_qs[%packed_word_index] : vector<1xi32>, view<32xi32> - } - %thread_sum = vector.reduce %rounded_values, %c0_f32 : vector<4xf32>, f32 - view.store %thread_sum, %scratch_values[%lane] : f32, view<32xf32> - kernel.barrier scope(workgroup) ordering(acq_rel) - %publishes_metadata = scalar.andi %is_cohort_leader, %publish_output : i1 - scf.if %publishes_metadata { - %cohort_base = index.mul %block_in_group, %c8 : index - %cohort_sums = vector.load %scratch_values[%cohort_base] : view<32xf32> -> vector<8xf32> - %quantized_sum = vector.reduce %cohort_sums, %c0_f32 : vector<8xf32>, f32 - %s = scalar.mulf %quantized_sum, %d : f32 - %d_f16 = scalar.fptrunc %d : f32 to f16 - %s_f16 = scalar.fptrunc %s : f32 to f16 - %ds_index = index.mul %block_in_group, %c2 : index - view.store %d_f16, %group_ds[%ds_index] : f16, view<8xf16> - %s_index = index.add %ds_index, %c1 : index - view.store %s_f16, %group_ds[%s_index] : f16, view<8xf16> - } - template.return -} // A 32-lane workgroup owns one physical group. Each eight-lane cohort loads // four values, computes one Q8_1 block, and writes disjoint metadata and packed diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/manifest.json b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/manifest.json index 65382c5e4f09..224ed9a5c8b6 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/manifest.json +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/manifest.json @@ -3,7 +3,43 @@ "upstream_revision": "local", "files": [ { - "path": "../loom-libs/motifs/matmul/q6_f32_decode.loom" + "path": "../loom-libs/motifs/q6_k_f16.loom" + }, + { + "path": "../loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom" + }, + { + "path": "../loom-libs/motifs/matmul/q4_q8_1_x4_decode.loom" + }, + { + "path": "../loom-libs/motifs/matmul/q4_q8_1_x4_dense.loom" + }, + { + "path": "../loom-libs/motifs/matmul/q4_q6_q8_1_x4_qkv.loom" + }, + { + "path": "../loom-libs/motifs/routing/router_projection_storage.loom" + }, + { + "path": "../loom-libs/motifs/routing/router_topk_row.loom" + }, + { + "path": "../loom-libs/motifs/routing/router_projection_fused.loom" + }, + { + "path": "../loom-libs/motifs/routing/weighted_reduce_f16_f32.loom" + }, + { + "path": "../loom-libs/motifs/routing/weighted_reduce_next_rmsnorm_f32.loom" + }, + { + "path": "../loom-libs/motifs/attention/rmsnorm_neox_cache_f16.loom" + }, + { + "path": "../loom-libs/motifs/normalization/rmsnorm_q8_1_x4.loom" + }, + { + "path": "../loom-libs/motifs/quantize_q8_1_x4_group.loom" }, { "path": "ggml/linear_q6k_q8_1_x4.loom" @@ -27,25 +63,22 @@ "path": "qwen3_moe/dense_linear_quantized_f16_wmma.loom" }, { - "path": "qwen3_moe/expert_table_partition_fused.loom" - }, - { - "path": "qwen3_moe/flash_attention_decode_f32_f16_wmma.loom" + "path": "qwen3_moe/motifs/dense_linear_q4k_next_q8.loom" }, { - "path": "qwen3_moe/flash_attention_decode_q128_f32_f16_wmma.loom" + "path": "qwen3_moe/motifs/attention_prepare_target.loom" }, { - "path": "qwen3_moe/model_config.loom" + "path": "qwen3_moe/motifs/attention_qkv_postprocess_fused_decode.loom" }, { - "path": "qwen3_moe/routed_down_q4k.loom" + "path": "qwen3_moe/expert_table_partition_fused.loom" }, { - "path": "qwen3_moe/routed_down_q6k.loom" + "path": "qwen3_moe/motifs/expert_table_partition_prefill_512.loom" }, { - "path": "qwen3_moe/routed_down_next_q8.loom" + "path": "qwen3_moe/model_config.loom" }, { "path": "qwen3_moe/routed_down_quantized_f16_wmma.loom" @@ -53,18 +86,9 @@ { "path": "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom" }, - { - "path": "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_q8_1_x4.loom" - }, - { - "path": "qwen3_moe/routed_gate_up_swiglu_q4k.loom" - }, { "path": "../loom-libs/motifs/matmul_id/vector_pair.loom" }, - { - "path": "../loom-libs/motifs/matmul_id/vector_weighted.loom" - }, { "path": "../loom-libs/motifs/matmul_id/vector_q4k_q8.loom" }, @@ -74,15 +98,9 @@ { "path": "../loom-libs/ops/matmul_id/vector_pair_binary.loom" }, - { - "path": "../loom-libs/ops/matmul_id/vector_weighted.loom" - }, { "path": "../loom-libs/ops/moe_routing_tables.loom" }, - { - "path": "qwen3_moe/routed_linear_q4k_f16_wmma.loom" - }, { "path": "qwen3_moe/router_projection_f32.loom" }, @@ -148,11 +166,17 @@ "ggml/linear_q6k_q8_1_x4.loom" ], "library_sources": [ - "ggml/quantize_q8_1_x4.loom" + "../loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom", + "../loom-libs/motifs/q6_k_f16.loom", + "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ] }, "compile_dependencies": [ - "ggml/quantize_q8_1_x4.loom" + "../loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom", + "../loom-libs/motifs/q6_k_f16.loom", + "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ] }, { @@ -205,41 +229,19 @@ "ggml/linear_q6k_q8_1_x4.loom" ], "library_sources": [ - "ggml/quantize_q8_1_x4.loom" + "../loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom", + "../loom-libs/motifs/q6_k_f16.loom", + "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ] }, "compile_dependencies": [ - "ggml/quantize_q8_1_x4.loom" + "../loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom", + "../loom-libs/motifs/q6_k_f16.loom", + "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ] }, - { - "name": "ggml_q8_1_x4_inspect_one_group", - "symbol": "ggml_q8_1_x4_inspect_one_group", - "source": "ggml/quantize_q8_1_x4.loom", - "workload_parameters": [], - "launch_parameters": [], - "bindings": [ - "packed", - "words", - "d_values", - "s_values" - ], - "binding_access": [ - "read", - "write", - "write", - "write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "ggml/quantize_q8_1_x4.loom" - ], - "library_sources": [] - }, - "compile_dependencies": [] - }, { "name": "ggml_quantize_q8_1_x4_f32", "symbol": "ggml_quantize_q8_1_x4_f32", @@ -278,9 +280,13 @@ "primary_sources": [ "ggml/quantize_q8_1_x4.loom" ], - "library_sources": [] + "library_sources": [ + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" + ] }, - "compile_dependencies": [] + "compile_dependencies": [ + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" + ] }, { "name": "qwen3_moe_attention_postprocess_f32_f16", @@ -342,11 +348,13 @@ "qwen3_moe/attention_postprocess_f32_f16.loom" ], "library_sources": [ - "qwen3_moe/model_config.loom" + "qwen3_moe/model_config.loom", + "../loom-libs/motifs/attention/rmsnorm_neox_cache_f16.loom" ] }, "compile_dependencies": [ - "qwen3_moe/model_config.loom" + "qwen3_moe/model_config.loom", + "../loom-libs/motifs/attention/rmsnorm_neox_cache_f16.loom" ] }, { @@ -419,137 +427,29 @@ "qwen3_moe/attention_qkv_postprocess_fused.loom" ], "library_sources": [ - "qwen3_moe/attention_postprocess_f32_f16.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/attention_qkv_quantized.loom", - "qwen3_moe/dense_linear_quantized_f16_wmma.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/attention_postprocess_f32_f16.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/attention_qkv_quantized.loom", - "qwen3_moe/dense_linear_quantized_f16_wmma.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_attention_qkv_quantized", - "symbol": "qwen3_moe_attention_qkv_quantized", - "source": "qwen3_moe/attention_qkv_quantized.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "q8_input", - "query_weight", - "key_weight", - "value_weight", - "query_output", - "key_output", - "value_output" - ], - "binding_access": [ - "read", - "read", - "read", - "read", - "read_write", - "read_write", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "attention_qkv_quantized_linked", - "primary_sources": [ - "qwen3_moe/attention_qkv_quantized.loom" - ], - "library_sources": [ - "qwen3_moe/dense_linear_quantized_f16_wmma.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/dense_linear_quantized_f16_wmma.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_attention_rmsnorm_quantize_q8_1_x4", - "symbol": "qwen3_moe_attention_rmsnorm_quantize_q8_1_x4", - "source": "qwen3_moe/attention_prepare_quantized.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "input", - "weight", - "output" - ], - "binding_access": [ - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "attention_prepare_quantized_linked", - "primary_sources": [ - "qwen3_moe/attention_prepare_quantized.loom" - ], - "library_sources": [ + "qwen3_moe/motifs/attention_qkv_postprocess_fused_decode.loom", + "../loom-libs/motifs/attention/rmsnorm_neox_cache_f16.loom", "qwen3_moe/model_config.loom", - "ggml/quantize_q8_1_x4.loom" + "../loom-libs/motifs/matmul/q4_q6_q8_1_x4_qkv.loom", + "../loom-libs/motifs/matmul/q4_q8_1_x4_decode.loom", + "../loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ] }, "compile_dependencies": [ + "qwen3_moe/motifs/attention_qkv_postprocess_fused_decode.loom", + "../loom-libs/motifs/attention/rmsnorm_neox_cache_f16.loom", "qwen3_moe/model_config.loom", - "ggml/quantize_q8_1_x4.loom" + "../loom-libs/motifs/matmul/q4_q6_q8_1_x4_qkv.loom", + "../loom-libs/motifs/matmul/q4_q8_1_x4_decode.loom", + "../loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ] }, { "name": "qwen3_moe_build_expert_table_partition_prefill_512", "symbol": "qwen3_moe_build_expert_table_partition_prefill_512", - "source": "qwen3_moe/expert_table_partition_fused.loom", + "source": "qwen3_moe/motifs/expert_table_partition_prefill_512.loom", "workload_parameters": [ { "name": "token_count", @@ -603,22 +503,22 @@ "mode": "archive", "link_module": "expert_table_partition_fused_linked", "primary_sources": [ - "qwen3_moe/expert_table_partition_fused.loom" + "qwen3_moe/motifs/expert_table_partition_prefill_512.loom" ], "library_sources": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" + "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ] }, "compile_dependencies": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" + "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ] }, { - "name": "qwen3_moe_dense_linear_q4k_q8_1_x4", - "symbol": "qwen3_moe_dense_linear_q4k_q8_1_x4", - "source": "qwen3_moe/dense_linear_quantized_f16_wmma.loom", + "name": "qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8", + "symbol": "qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8", + "source": "qwen3_moe/motifs/dense_linear_q4k_next_q8.loom", "workload_parameters": [ { "name": "token_count", @@ -634,11 +534,19 @@ "bindings": [ "q8_input", "weight", - "output" + "output", + "norm_weight", + "normalized_output", + "completion_counter", + "next_q8_output" ], "binding_access": [ "read", "read", + "read_write", + "read", + "read_write", + "read_write", "read_write" ], "target_selector": "", @@ -646,30 +554,28 @@ "mode": "archive", "link_module": "dense_linear_quantized_f16_wmma_linked", "primary_sources": [ - "qwen3_moe/dense_linear_quantized_f16_wmma.loom" + "qwen3_moe/motifs/dense_linear_q4k_next_q8.loom" ], "library_sources": [ - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" + "qwen3_moe/motifs/attention_prepare_target.loom", + "../loom-libs/motifs/matmul/q4_q8_1_x4_dense.loom", + "../loom-libs/motifs/matmul/q4_q8_1_x4_decode.loom", + "../loom-libs/motifs/normalization/rmsnorm_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ] }, "compile_dependencies": [ - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" + "qwen3_moe/motifs/attention_prepare_target.loom", + "../loom-libs/motifs/matmul/q4_q8_1_x4_dense.loom", + "../loom-libs/motifs/matmul/q4_q8_1_x4_decode.loom", + "../loom-libs/motifs/normalization/rmsnorm_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ] }, { - "name": "qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8", - "symbol": "qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8", - "source": "qwen3_moe/dense_linear_quantized_f16_wmma.loom", + "name": "qwen3_moe_rmsnorm_f32_quantize_q8_1_x4", + "symbol": "qwen3_moe_rmsnorm_f32_quantize_q8_1_x4", + "source": "qwen3_moe/attention_prepare_quantized.loom", "workload_parameters": [ { "name": "token_count", @@ -683,1322 +589,131 @@ } ], "bindings": [ - "q8_input", + "input", "weight", - "output", - "norm_weight", "normalized_output", - "completion_counter", - "next_q8_output" + "q8_output" ], "binding_access": [ "read", "read", "read_write", - "read", - "read_write", - "read_write", "read_write" ], "target_selector": "", "compile_recipe": { "mode": "archive", - "link_module": "dense_linear_quantized_f16_wmma_linked", + "link_module": "attention_prepare_quantized_linked", "primary_sources": [ - "qwen3_moe/dense_linear_quantized_f16_wmma.loom" + "qwen3_moe/attention_prepare_quantized.loom" ], "library_sources": [ - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/motifs/attention_prepare_target.loom", + "../loom-libs/motifs/normalization/rmsnorm_q8_1_x4.loom", "qwen3_moe/model_config.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" + "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ] }, "compile_dependencies": [ - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/motifs/attention_prepare_target.loom", + "../loom-libs/motifs/normalization/rmsnorm_q8_1_x4.loom", "qwen3_moe/model_config.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" + "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ] }, { - "name": "qwen3_moe_flash_attention_decode_f32_f16_wmma", - "symbol": "qwen3_moe_flash_attention_decode_f32_f16_wmma", - "source": "qwen3_moe/flash_attention_decode_f32_f16_wmma.loom", + "name": "qwen3_moe_routed_down_weighted_reduce_f16_f32", + "symbol": "qwen3_moe_routed_down_weighted_reduce_f16_f32", + "source": "qwen3_moe/routed_down_quantized_f16_wmma.loom", "workload_parameters": [ { - "name": "key_value_token_count", + "name": "token_count", "type": "index" } ], "launch_parameters": [ { - "name": "key_value_token_count", + "name": "token_count", "type": "index" } ], "bindings": [ - "query", - "key", - "value", - "mask", + "route_weights", + "routed_output", + "residual_input", "output" ], "binding_access": [ "read", "read", "read", - "read", - "read_write" + "write" ], "target_selector": "", "compile_recipe": { - "mode": "direct", + "mode": "archive", + "link_module": "routed_down_quantized_f16_wmma_linked", "primary_sources": [ - "qwen3_moe/flash_attention_decode_f32_f16_wmma.loom" + "qwen3_moe/routed_down_quantized_f16_wmma.loom" ], - "library_sources": [] + "library_sources": [ + "../loom-libs/motifs/routing/weighted_reduce_f16_f32.loom" + ] }, - "compile_dependencies": [] + "compile_dependencies": [ + "../loom-libs/motifs/routing/weighted_reduce_f16_f32.loom" + ] }, { - "name": "qwen3_moe_flash_attention_decode_q128_fused_f32_f16_wmma", - "symbol": "qwen3_moe_flash_attention_decode_q128_fused_f32_f16_wmma", - "source": "qwen3_moe/flash_attention_decode_q128_f32_f16_wmma.loom", + "name": "qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32", + "symbol": "qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32", + "source": "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom", "workload_parameters": [ { - "name": "key_value_token_count", + "name": "token_count", "type": "index" } ], "launch_parameters": [ { - "name": "key_value_token_count", + "name": "token_count", "type": "index" } ], "bindings": [ - "query", - "key", - "value", - "mask", - "output" - ], - "binding_access": [ - "read", - "read", - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "qwen3_moe/flash_attention_decode_q128_f32_f16_wmma.loom" - ], - "library_sources": [] - }, - "compile_dependencies": [] - }, - { - "name": "qwen3_moe_rmsnorm_f32", - "symbol": "qwen3_moe_rmsnorm_f32", - "source": "qwen3_moe/attention_prepare_quantized.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "input", - "weight", - "output" - ], - "binding_access": [ - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "attention_prepare_quantized_linked", - "primary_sources": [ - "qwen3_moe/attention_prepare_quantized.loom" - ], - "library_sources": [ - "qwen3_moe/model_config.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/model_config.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_rmsnorm_f32_quantize_q8_1_x4", - "symbol": "qwen3_moe_rmsnorm_f32_quantize_q8_1_x4", - "source": "qwen3_moe/attention_prepare_quantized.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "input", - "weight", - "normalized_output", - "q8_output" - ], - "binding_access": [ - "read", - "read", - "read_write", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "attention_prepare_quantized_linked", - "primary_sources": [ - "qwen3_moe/attention_prepare_quantized.loom" - ], - "library_sources": [ - "qwen3_moe/model_config.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/model_config.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_rmsnorm_quantize_q8_1_x4_wave64_production_check", - "symbol": "qwen3_moe_rmsnorm_quantize_q8_1_x4_wave64_production_check", - "source": "qwen3_moe/routed_down_q6k.loom", - "workload_parameters": [], - "launch_parameters": [], - "bindings": [ - "input", - "norm_weight", - "q8_output" - ], - "binding_access": [ - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "routed_down_q6k_linked", - "primary_sources": [ - "qwen3_moe/routed_down_q6k.loom" - ], - "library_sources": [ - "../loom-libs/motifs/matmul/q6_f32_decode.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "ggml/quantize_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_down_next_q8.loom", - "../loom-libs/ops/matmul_id/vector_weighted.loom", - "../loom-libs/motifs/matmul_id/vector_q4k_q8.loom", - "../loom-libs/motifs/matmul_id/vector_weighted.loom" - ] - }, - "compile_dependencies": [ - "../loom-libs/motifs/matmul/q6_f32_decode.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "ggml/quantize_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_down_next_q8.loom", - "../loom-libs/ops/matmul_id/vector_weighted.loom", - "../loom-libs/motifs/matmul_id/vector_q4k_q8.loom", - "../loom-libs/motifs/matmul_id/vector_weighted.loom" - ] - }, - { - "name": "qwen3_moe_routed_down_q4k_f16_wmma_grouped", - "symbol": "qwen3_moe_routed_down_q4k_f16_wmma_grouped", - "source": "qwen3_moe/routed_down_quantized_f16_wmma.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "input", - "expert_table", - "weight", - "output" - ], - "binding_access": [ - "read", - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "routed_down_quantized_f16_wmma_linked", - "primary_sources": [ - "qwen3_moe/routed_down_quantized_f16_wmma.loom" - ], - "library_sources": [ - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_down_q4k.loom", - "qwen3_moe/routed_down_q6k.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_down_q4k.loom", - "qwen3_moe/routed_down_q6k.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_routed_down_q4k_q8_1_x4", - "symbol": "qwen3_moe_routed_down_q4k_q8_1_x4", - "source": "qwen3_moe/routed_down_q4k.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "route_id_stride", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "route_id_stride", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "bindings": [ - "q8_input", - "route_ids", - "route_weights", - "weight", - "output" - ], - "binding_access": [ - "read", - "read", - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "routed_down_q4k_linked", - "primary_sources": [ - "qwen3_moe/routed_down_q4k.loom" - ], - "library_sources": [ - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom", - "../loom-libs/ops/matmul_id/vector_pair_binary.loom", - "../loom-libs/ops/matmul_id/vector_weighted.loom", - "../loom-libs/motifs/matmul_id/vector_weighted.loom", - "../loom-libs/motifs/matmul_id/vector_pair.loom", - "../loom-libs/motifs/matmul_id/vector_q4k_q8.loom", - "../loom-libs/motifs/elementwise/binary_apply.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom", - "../loom-libs/ops/matmul_id/vector_pair_binary.loom", - "../loom-libs/ops/matmul_id/vector_weighted.loom", - "../loom-libs/motifs/matmul_id/vector_weighted.loom", - "../loom-libs/motifs/matmul_id/vector_pair.loom", - "../loom-libs/motifs/matmul_id/vector_q4k_q8.loom", - "../loom-libs/motifs/elementwise/binary_apply.loom" - ] - }, - { - "name": "qwen3_moe_routed_down_q6k_f16_wmma_grouped", - "symbol": "qwen3_moe_routed_down_q6k_f16_wmma_grouped", - "source": "qwen3_moe/routed_down_quantized_f16_wmma.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "input", - "expert_table", - "weight", - "output" - ], - "binding_access": [ - "read", - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "routed_down_quantized_f16_wmma_linked", - "primary_sources": [ - "qwen3_moe/routed_down_quantized_f16_wmma.loom" - ], - "library_sources": [ - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_down_q4k.loom", - "qwen3_moe/routed_down_q6k.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_down_q4k.loom", - "qwen3_moe/routed_down_q6k.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_routed_down_q6k_f32_wave64", - "symbol": "qwen3_moe_routed_down_q6k_f32_wave64", - "source": "qwen3_moe/routed_down_q6k.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "route_id_stride", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "route_id_stride", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "bindings": [ - "input", - "route_ids", - "route_weights", - "weight", - "output" - ], - "binding_access": [ - "read", - "read", - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "routed_down_q6k_linked", - "primary_sources": [ - "qwen3_moe/routed_down_q6k.loom" - ], - "library_sources": [ - "../loom-libs/motifs/matmul/q6_f32_decode.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "ggml/quantize_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_down_next_q8.loom", - "../loom-libs/ops/matmul_id/vector_weighted.loom", - "../loom-libs/motifs/matmul_id/vector_q4k_q8.loom", - "../loom-libs/motifs/matmul_id/vector_weighted.loom" - ] - }, - "compile_dependencies": [ - "../loom-libs/motifs/matmul/q6_f32_decode.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "ggml/quantize_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_down_next_q8.loom", - "../loom-libs/ops/matmul_id/vector_weighted.loom", - "../loom-libs/motifs/matmul_id/vector_q4k_q8.loom", - "../loom-libs/motifs/matmul_id/vector_weighted.loom" - ] - }, - { - "name": "qwen3_moe_routed_down_q6k_q8_1_x4", - "symbol": "qwen3_moe_routed_down_q6k_q8_1_x4", - "source": "qwen3_moe/routed_down_q6k.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "route_id_stride", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "route_id_stride", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "bindings": [ - "q8_input", - "route_ids", - "route_weights", - "weight", - "output" - ], - "binding_access": [ - "read", - "read", - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "routed_down_q6k_linked", - "primary_sources": [ - "qwen3_moe/routed_down_q6k.loom" - ], - "library_sources": [ - "../loom-libs/motifs/matmul/q6_f32_decode.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "ggml/quantize_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_down_next_q8.loom", - "../loom-libs/ops/matmul_id/vector_weighted.loom", - "../loom-libs/motifs/matmul_id/vector_q4k_q8.loom", - "../loom-libs/motifs/matmul_id/vector_weighted.loom" - ] - }, - "compile_dependencies": [ - "../loom-libs/motifs/matmul/q6_f32_decode.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "ggml/quantize_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_down_next_q8.loom", - "../loom-libs/ops/matmul_id/vector_weighted.loom", - "../loom-libs/motifs/matmul_id/vector_q4k_q8.loom", - "../loom-libs/motifs/matmul_id/vector_weighted.loom" - ] - }, - { - "name": "qwen3_moe_routed_down_q6k_q8_1_x4_next_q8", - "symbol": "qwen3_moe_routed_down_q6k_q8_1_x4_next_q8", - "source": "qwen3_moe/routed_down_q6k.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "route_id_stride", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "route_id_stride", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "bindings": [ - "q8_input", - "route_ids", - "route_weights", - "weight", - "output", - "norm_weight", - "completion_counter", - "next_q8_output" - ], - "binding_access": [ - "read", - "read", - "read", - "read", - "read_write", - "read", - "read_write", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "routed_down_q6k_linked", - "primary_sources": [ - "qwen3_moe/routed_down_q6k.loom" - ], - "library_sources": [ - "../loom-libs/motifs/matmul/q6_f32_decode.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "ggml/quantize_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_down_next_q8.loom", - "../loom-libs/ops/matmul_id/vector_weighted.loom", - "../loom-libs/motifs/matmul_id/vector_q4k_q8.loom", - "../loom-libs/motifs/matmul_id/vector_weighted.loom" - ] - }, - "compile_dependencies": [ - "../loom-libs/motifs/matmul/q6_f32_decode.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "ggml/quantize_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_down_next_q8.loom", - "../loom-libs/ops/matmul_id/vector_weighted.loom", - "../loom-libs/motifs/matmul_id/vector_q4k_q8.loom", - "../loom-libs/motifs/matmul_id/vector_weighted.loom" - ] - }, - { - "name": "qwen3_moe_routed_down_weighted_reduce_f16_f32", - "symbol": "qwen3_moe_routed_down_weighted_reduce_f16_f32", - "source": "qwen3_moe/routed_down_quantized_f16_wmma.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "route_weights", - "routed_output", - "residual_input", - "output" - ], - "binding_access": [ - "read", - "read", - "read", - "write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "routed_down_quantized_f16_wmma_linked", - "primary_sources": [ - "qwen3_moe/routed_down_quantized_f16_wmma.loom" - ], - "library_sources": [ - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_down_q4k.loom", - "qwen3_moe/routed_down_q6k.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_down_q4k.loom", - "qwen3_moe/routed_down_q6k.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32", - "symbol": "qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32", - "source": "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "route_weights", - "routed_output", - "hidden_state", - "next_norm_weight", - "next_projection_input" - ], - "binding_access": [ - "read", - "read", - "read_write", - "read", - "write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "routed_down_weighted_reduce_next_rmsnorm_f32_linked", - "primary_sources": [ - "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom" - ], - "library_sources": [ - "qwen3_moe/routed_down_quantized_f16_wmma.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_down_q4k.loom", - "qwen3_moe/routed_down_q6k.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/routed_down_quantized_f16_wmma.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_down_q4k.loom", - "qwen3_moe/routed_down_q6k.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4", - "symbol": "qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4", - "source": "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_q8_1_x4.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "route_weights", - "routed_output", - "hidden_state", - "next_norm_weight", - "next_projection_input" - ], - "binding_access": [ - "read", - "read", - "read", - "read", - "read" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_linked", - "primary_sources": [ - "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_q8_1_x4.loom" - ], - "library_sources": [ - "qwen3_moe/routed_down_quantized_f16_wmma.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_down_q4k.loom", - "qwen3_moe/routed_down_q6k.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/routed_down_quantized_f16_wmma.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_down_q4k.loom", - "qwen3_moe/routed_down_q6k.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma", - "symbol": "qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma", - "source": "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "input", - "expert_table", - "partition_table", - "gate_weight", - "up_weight", - "output" - ], - "binding_access": [ - "read", - "read", - "read", - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "routed_gate_up_swiglu_q4k_f16_wmma_linked", - "primary_sources": [ - "qwen3_moe/routed_linear_q4k_f16_wmma.loom" - ], - "library_sources": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped", - "symbol": "qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped", - "source": "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "bindings": [ - "q8_input", - "expert_table", - "gate_weight", - "up_weight", - "output" - ], - "binding_access": [ - "read", - "read", - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "routed_gate_up_swiglu_q4k_linked", - "primary_sources": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom" - ], - "library_sources": [ - "ggml/quantize_q8_1_x4.loom", - "../loom-libs/ops/matmul_id/vector_pair_binary.loom", - "../loom-libs/motifs/matmul_id/vector_pair.loom", - "../loom-libs/motifs/matmul_id/vector_q4k_q8.loom", - "../loom-libs/motifs/elementwise/binary_apply.loom" - ] - }, - "compile_dependencies": [ - "ggml/quantize_q8_1_x4.loom", - "../loom-libs/ops/matmul_id/vector_pair_binary.loom", - "../loom-libs/motifs/matmul_id/vector_pair.loom", - "../loom-libs/motifs/matmul_id/vector_q4k_q8.loom", - "../loom-libs/motifs/elementwise/binary_apply.loom" - ] - }, - { - "name": "qwen3_moe_routed_linear_q4k_f16_wmma", - "symbol": "qwen3_moe_routed_linear_q4k_f16_wmma", - "source": "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "input", - "expert_table", - "weight", - "output" - ], - "binding_access": [ - "read", - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "routed_gate_up_swiglu_q4k_f16_wmma_linked", - "primary_sources": [ - "qwen3_moe/routed_linear_q4k_f16_wmma.loom" - ], - "library_sources": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_routed_swiglu_f16", - "symbol": "qwen3_moe_routed_swiglu_f16", - "source": "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "gate", - "up", - "output" - ], - "binding_access": [ - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "routed_gate_up_swiglu_q4k_f16_wmma_linked", - "primary_sources": [ - "qwen3_moe/routed_linear_q4k_f16_wmma.loom" - ], - "library_sources": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_routed_swiglu_f32", - "symbol": "qwen3_moe_routed_swiglu_f32", - "source": "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "gate", - "up", - "output" - ], - "binding_access": [ - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "routed_gate_up_swiglu_q4k_f16_wmma_linked", - "primary_sources": [ - "qwen3_moe/routed_linear_q4k_f16_wmma.loom" - ], - "library_sources": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_router_projection_f32_four_row_wave32", - "symbol": "qwen3_moe_router_projection_f32_four_row_wave32", - "source": "qwen3_moe/router_projection_f32.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "input", - "weight", - "output" - ], - "binding_access": [ - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "router_projection_f32_linked", - "primary_sources": [ - "qwen3_moe/router_projection_f32.loom" - ], - "library_sources": [ - "qwen3_moe/model_config.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/model_config.loom" - ] - }, - { - "name": "qwen3_moe_router_projection_f32_one_row_wave64", - "symbol": "qwen3_moe_router_projection_f32_one_row_wave64", - "source": "qwen3_moe/router_projection_f32.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "input", - "weight", - "output" + "route_weights", + "routed_output", + "hidden_state", + "next_norm_weight", + "next_projection_input" ], "binding_access": [ "read", "read", - "read_write" + "read_write", + "read", + "write" ], "target_selector": "", "compile_recipe": { "mode": "archive", - "link_module": "router_projection_f32_linked", + "link_module": "routed_down_weighted_reduce_next_rmsnorm_f32_linked", "primary_sources": [ - "qwen3_moe/router_projection_f32.loom" + "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom" ], "library_sources": [ - "qwen3_moe/model_config.loom" + "../loom-libs/motifs/routing/weighted_reduce_next_rmsnorm_f32.loom" ] }, "compile_dependencies": [ - "qwen3_moe/model_config.loom" + "../loom-libs/motifs/routing/weighted_reduce_next_rmsnorm_f32.loom" ] }, { - "name": "qwen3_moe_router_projection_f32_reference", - "symbol": "qwen3_moe_router_projection_f32_reference", + "name": "qwen3_moe_router_projection_f32_four_row_wave32", + "symbol": "qwen3_moe_router_projection_f32_four_row_wave32", "source": "qwen3_moe/router_projection_f32.loom", "workload_parameters": [ { @@ -2030,10 +745,12 @@ "qwen3_moe/router_projection_f32.loom" ], "library_sources": [ + "../loom-libs/motifs/routing/router_projection_storage.loom", "qwen3_moe/model_config.loom" ] }, "compile_dependencies": [ + "../loom-libs/motifs/routing/router_projection_storage.loom", "qwen3_moe/model_config.loom" ] }, @@ -2086,14 +803,16 @@ ], "library_sources": [ "qwen3_moe/model_config.loom", - "qwen3_moe/router_projection_f32.loom", - "qwen3_moe/router_top8_f32.loom" + "../loom-libs/motifs/routing/router_projection_storage.loom", + "../loom-libs/motifs/routing/router_topk_row.loom", + "../loom-libs/motifs/routing/router_projection_fused.loom" ] }, "compile_dependencies": [ "qwen3_moe/model_config.loom", - "qwen3_moe/router_projection_f32.loom", - "qwen3_moe/router_top8_f32.loom" + "../loom-libs/motifs/routing/router_projection_storage.loom", + "../loom-libs/motifs/routing/router_topk_row.loom", + "../loom-libs/motifs/routing/router_projection_fused.loom" ] }, { @@ -2138,37 +857,12 @@ "qwen3_moe/router_top8_f32.loom" ], "library_sources": [ + "../loom-libs/motifs/routing/router_topk_row.loom", "qwen3_moe/model_config.loom" ] }, "compile_dependencies": [ - "qwen3_moe/model_config.loom" - ] - }, - { - "name": "qwen3_moe_router_top8_wide_stride_reference", - "symbol": "qwen3_moe_router_top8_wide_stride_reference", - "source": "qwen3_moe/router_top8_f32.loom", - "workload_parameters": [], - "launch_parameters": [], - "bindings": [ - "route_ids" - ], - "binding_access": [ - "read" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "router_top8_f32_linked", - "primary_sources": [ - "qwen3_moe/router_top8_f32.loom" - ], - "library_sources": [ - "qwen3_moe/model_config.loom" - ] - }, - "compile_dependencies": [ + "../loom-libs/motifs/routing/router_topk_row.loom", "qwen3_moe/model_config.loom" ] }, @@ -2219,41 +913,14 @@ "library_sources": [] }, "compile_dependencies": [] - }, - { - "name": "qwen_decode_attention_metadata", - "symbol": "qwen_decode_attention_metadata", - "source": "../qwen/attention_metadata.loom", - "workload_parameters": [], - "launch_parameters": [], - "bindings": [ - "control", - "positions", - "key_cache_indices", - "value_cache_indices" - ], - "binding_access": [ - "read", - "read_write", - "read_write", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "../qwen/attention_metadata.loom" - ], - "library_sources": [] - }, - "compile_dependencies": [] } ], "link_modules": [ { "name": "router_projection_f32_linked", "srcs": [ - "qwen3_moe/router_projection_f32.loom" + "qwen3_moe/router_projection_f32.loom", + "../loom-libs/motifs/routing/router_projection_storage.loom" ], "libraries": [ "qwen3_moe/model_config.loom" @@ -2262,7 +929,8 @@ { "name": "router_top8_f32_linked", "srcs": [ - "qwen3_moe/router_top8_f32.loom" + "qwen3_moe/router_top8_f32.loom", + "../loom-libs/motifs/routing/router_topk_row.loom" ], "libraries": [ "qwen3_moe/model_config.loom" @@ -2271,7 +939,10 @@ { "name": "router_projection_top8_fused_f32_linked", "srcs": [ - "qwen3_moe/router_projection_top8_fused_f32.loom" + "qwen3_moe/router_projection_top8_fused_f32.loom", + "../loom-libs/motifs/routing/router_projection_storage.loom", + "../loom-libs/motifs/routing/router_topk_row.loom", + "../loom-libs/motifs/routing/router_projection_fused.loom" ], "libraries": [ "qwen3_moe/model_config.loom", @@ -2282,7 +953,8 @@ { "name": "attention_postprocess_f32_f16_linked", "srcs": [ - "qwen3_moe/attention_postprocess_f32_f16.loom" + "qwen3_moe/attention_postprocess_f32_f16.loom", + "../loom-libs/motifs/attention/rmsnorm_neox_cache_f16.loom" ], "libraries": [ "qwen3_moe/model_config.loom" @@ -2291,11 +963,14 @@ { "name": "attention_prepare_quantized_linked", "srcs": [ - "qwen3_moe/attention_prepare_quantized.loom" + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/motifs/attention_prepare_target.loom" ], "libraries": [ + "../loom-libs/motifs/normalization/rmsnorm_q8_1_x4.loom", "qwen3_moe/model_config.loom", - "ggml/quantize_q8_1_x4.loom" + "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ] }, { @@ -2304,16 +979,19 @@ "qwen3_moe/attention_qkv_quantized.loom" ], "libraries": [ - ":dense_linear_quantized_f16_wmma_linked" + "../loom-libs/motifs/matmul/q4_q6_q8_1_x4_qkv.loom", + ":dense_linear_quantized_f16_wmma_linked", + ":attention_prepare_quantized_linked" ] }, { "name": "attention_qkv_postprocess_fused_linked", "srcs": [ - "qwen3_moe/attention_qkv_postprocess_fused.loom" + "qwen3_moe/attention_qkv_postprocess_fused.loom", + "../loom-libs/motifs/attention/rmsnorm_neox_cache_f16.loom" ], "libraries": [ - ":attention_postprocess_f32_f16_linked", + "qwen3_moe/motifs/attention_qkv_postprocess_fused_decode.loom", ":attention_qkv_quantized_linked" ] }, @@ -2323,7 +1001,10 @@ "ggml/linear_q6k_q8_1_x4.loom" ], "libraries": [ - "ggml/quantize_q8_1_x4.loom" + "../loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom", + "../loom-libs/motifs/q6_k_f16.loom", + "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ] }, { @@ -2332,22 +1013,7 @@ "qwen3_moe/routed_down_quantized_f16_wmma.loom" ], "libraries": [ - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_down_q4k.loom", - "qwen3_moe/routed_down_q6k.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom", - "../loom-libs/ops/matmul_id/vector_pair_binary.loom", - "../loom-libs/ops/matmul_id/vector_weighted.loom", - "../loom-libs/motifs/matmul_id/vector_weighted.loom", - "../loom-libs/motifs/matmul_id/vector_pair.loom", - "../loom-libs/motifs/matmul_id/vector_q4k_q8.loom", - "../loom-libs/motifs/elementwise/binary_apply.loom", - "../loom-libs/ops/moe_routing_tables.loom" + "../loom-libs/motifs/routing/weighted_reduce_f16_f32.loom" ] }, { @@ -2356,124 +1022,18 @@ "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom" ], "libraries": [ - "qwen3_moe/routed_down_quantized_f16_wmma.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_down_q4k.loom", - "qwen3_moe/routed_down_q6k.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom", - "../loom-libs/ops/matmul_id/vector_pair_binary.loom", - "../loom-libs/ops/matmul_id/vector_weighted.loom", - "../loom-libs/motifs/matmul_id/vector_weighted.loom", - "../loom-libs/motifs/matmul_id/vector_pair.loom", - "../loom-libs/motifs/matmul_id/vector_q4k_q8.loom", - "../loom-libs/motifs/elementwise/binary_apply.loom", - "../loom-libs/ops/moe_routing_tables.loom" - ] - }, - { - "name": "routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_linked", - "srcs": [ - "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_q8_1_x4.loom" - ], - "libraries": [ - "qwen3_moe/routed_down_quantized_f16_wmma.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_down_q4k.loom", - "qwen3_moe/routed_down_q6k.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom", - "../loom-libs/ops/matmul_id/vector_pair_binary.loom", - "../loom-libs/ops/matmul_id/vector_weighted.loom", - "../loom-libs/motifs/matmul_id/vector_weighted.loom", - "../loom-libs/motifs/matmul_id/vector_pair.loom", - "../loom-libs/motifs/matmul_id/vector_q4k_q8.loom", - "../loom-libs/motifs/elementwise/binary_apply.loom", - "../loom-libs/ops/moe_routing_tables.loom" - ] - }, - { - "name": "routed_down_q4k_linked", - "srcs": [ - "qwen3_moe/routed_down_q4k.loom" - ], - "libraries": [ - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom", - "../loom-libs/ops/matmul_id/vector_pair_binary.loom", - "../loom-libs/ops/matmul_id/vector_weighted.loom", - "../loom-libs/motifs/matmul_id/vector_weighted.loom", - "../loom-libs/motifs/matmul_id/vector_pair.loom", - "../loom-libs/motifs/matmul_id/vector_q4k_q8.loom", - "../loom-libs/motifs/elementwise/binary_apply.loom", - "../loom-libs/ops/moe_routing_tables.loom" - ] - }, - { - "name": "routed_down_q6k_linked", - "srcs": [ - "qwen3_moe/routed_down_q6k.loom" - ], - "libraries": [ - "../loom-libs/motifs/matmul/q6_f32_decode.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "ggml/quantize_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_down_next_q8.loom", - "../loom-libs/ops/matmul_id/vector_weighted.loom", - "../loom-libs/motifs/matmul_id/vector_q4k_q8.loom", - "../loom-libs/motifs/matmul_id/vector_weighted.loom" - ] - }, - { - "name": "routed_gate_up_swiglu_q4k_linked", - "srcs": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom" - ], - "libraries": [ - "ggml/quantize_q8_1_x4.loom", - "../loom-libs/ops/matmul_id/vector_pair_binary.loom", - "../loom-libs/motifs/matmul_id/vector_pair.loom", - "../loom-libs/motifs/matmul_id/vector_q4k_q8.loom", - "../loom-libs/motifs/elementwise/binary_apply.loom", - "../loom-libs/ops/moe_routing_tables.loom" + "../loom-libs/motifs/routing/weighted_reduce_next_rmsnorm_f32.loom" ] }, { "name": "expert_table_partition_fused_linked", "srcs": [ - "qwen3_moe/expert_table_partition_fused.loom" - ], - "libraries": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom", - "../loom-libs/ops/matmul_id/vector_pair_binary.loom", - "../loom-libs/motifs/matmul_id/vector_pair.loom", - "../loom-libs/motifs/matmul_id/vector_q4k_q8.loom", - "../loom-libs/motifs/elementwise/binary_apply.loom", - "../loom-libs/ops/moe_routing_tables.loom" - ] - }, - { - "name": "routed_gate_up_swiglu_q4k_f16_wmma_linked", - "srcs": [ - "qwen3_moe/routed_linear_q4k_f16_wmma.loom" + "qwen3_moe/expert_table_partition_fused.loom", + "qwen3_moe/motifs/expert_table_partition_prefill_512.loom" ], "libraries": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom", "../loom-libs/ops/matmul_id/vector_pair_binary.loom", "../loom-libs/motifs/matmul_id/vector_pair.loom", "../loom-libs/motifs/matmul_id/vector_q4k_q8.loom", @@ -2484,15 +1044,20 @@ { "name": "dense_linear_quantized_f16_wmma_linked", "srcs": [ - "qwen3_moe/dense_linear_quantized_f16_wmma.loom" + "qwen3_moe/dense_linear_quantized_f16_wmma.loom", + "qwen3_moe/motifs/dense_linear_q4k_next_q8.loom", + "qwen3_moe/motifs/attention_prepare_target.loom" ], "libraries": [ + "../loom-libs/motifs/matmul/q4_q8_1_x4_dense.loom", + "../loom-libs/motifs/matmul/q4_q8_1_x4_decode.loom", "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", + "../loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom", + "../loom-libs/motifs/q6_k_f16.loom", + "../loom-libs/motifs/normalization/rmsnorm_q8_1_x4.loom", "qwen3_moe/model_config.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom", "../loom-libs/ops/matmul_id/vector_pair_binary.loom", "../loom-libs/motifs/matmul_id/vector_pair.loom", "../loom-libs/motifs/matmul_id/vector_q4k_q8.loom", @@ -2550,6 +1115,7 @@ "--config=qwen3_moe.attention.head_size=128", "--config=qwen3_moe.attention.key_value_size=512", "--config=qwen3_moe.attention.query_size=4096", + "--config=qwen3_moe.attention.rope_mscale=1", "--config=qwen3_moe.model.rms_epsilon=0.000001", "--config=qwen3_moe.workload.token_capacity=1", "--dry-run", @@ -2597,6 +1163,7 @@ "--config=qwen3_moe.attention.head_size=128", "--config=qwen3_moe.attention.key_value_size=512", "--config=qwen3_moe.attention.query_size=4096", + "--config=qwen3_moe.attention.rope_mscale=1", "--config=qwen3_moe.attention.value_uses_q6=0", "--config=qwen3_moe.dense_quantized.input_size=2048", "--config=qwen3_moe.dense_quantized.output_accumulation=0", @@ -2618,6 +1185,7 @@ "--config=qwen3_moe.attention.head_size=128", "--config=qwen3_moe.attention.key_value_size=512", "--config=qwen3_moe.attention.query_size=4096", + "--config=qwen3_moe.attention.rope_mscale=1", "--config=qwen3_moe.attention.value_uses_q6=1", "--config=qwen3_moe.dense_quantized.input_size=2048", "--config=qwen3_moe.dense_quantized.output_accumulation=0", @@ -2639,6 +1207,7 @@ "--config=qwen3_moe.attention.head_size=128", "--config=qwen3_moe.attention.key_value_size=512", "--config=qwen3_moe.attention.query_size=4096", + "--config=qwen3_moe.attention.rope_mscale=1", "--config=qwen3_moe.attention.value_uses_q6=1", "--config=qwen3_moe.dense_quantized.input_size=2048", "--config=qwen3_moe.dense_quantized.output_accumulation=0", @@ -2662,28 +1231,6 @@ ], "source": "ggml/quantize_q8_1_x4.loom" }, - { - "name": "routed_gate_up_swiglu_q4k_plan_test", - "args": [ - "$(location :routed_gate_up_swiglu_q4k_linked)", - "--benchmark=@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_prefill_512", - "--config=qwen3_moe.routed_gate_up.expert_count=128", - "--config=qwen3_moe.routed_gate_up.input_size=2048", - "--config=qwen3_moe.routed_gate_up.output_size=768", - "--config=qwen3_moe.routed_gate_up.route_count=8", - "--config=qwen3_moe.workload.token_capacity=512", - "--dry-run", - "--output-format=jsonl", - "--config=ggml.moe_routing.descriptor_expert_mask=127", - "--config=ggml.moe_routing.descriptor_partition_shift=7", - "--config=ggml.moe_routing.descriptor_row_count_shift=13", - "--config=ggml.moe_routing.expert_count=128", - "--config=ggml.moe_routing.partition_workgroup_size=128", - "--config=ggml.moe_routing.route_count=8", - "--config=ggml.workload.token_capacity=2048" - ], - "link_module": "routed_gate_up_swiglu_q4k_linked" - }, { "name": "linear_q6k_q8_1_x4_plan_test", "args": [ @@ -2717,86 +1264,6 @@ ], "link_module": "routed_down_weighted_reduce_next_rmsnorm_f32_linked" }, - { - "name": "routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_plan_test", - "args": [ - "$(location :routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_linked)", - "--benchmark=@qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_fused_prefill_14", - "--config=qwen3_moe.model.hidden_size=2048", - "--config=qwen3_moe.model.rms_epsilon=0.000001", - "--config=qwen3_moe.routed_down.output_size=2048", - "--config=qwen3_moe.routed_down.route_count=8", - "--config=qwen3_moe.workload.token_capacity=512", - "--dry-run", - "--output-format=jsonl", - "--config=ggml.moe_routing.descriptor_expert_mask=127", - "--config=ggml.moe_routing.descriptor_partition_shift=7", - "--config=ggml.moe_routing.descriptor_row_count_shift=13", - "--config=ggml.moe_routing.expert_count=128", - "--config=ggml.moe_routing.partition_workgroup_size=128", - "--config=ggml.moe_routing.route_count=8", - "--config=ggml.workload.token_capacity=2048" - ], - "link_module": "routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_linked" - }, - { - "name": "routed_down_quantized_f16_wmma_plan_test", - "args": [ - "$(location :routed_down_quantized_f16_wmma_linked)", - "--config=qwen3_moe.routed_down.expert_count=128", - "--config=qwen3_moe.routed_down.input_size=768", - "--config=qwen3_moe.routed_down.output_size=2048", - "--config=qwen3_moe.routed_down.route_count=8", - "--config=qwen3_moe.workload.token_capacity=512", - "--dry-run", - "--output-format=jsonl", - "--config=ggml.moe_routing.descriptor_expert_mask=127", - "--config=ggml.moe_routing.descriptor_partition_shift=7", - "--config=ggml.moe_routing.descriptor_row_count_shift=13", - "--config=ggml.moe_routing.expert_count=128", - "--config=ggml.moe_routing.partition_workgroup_size=128", - "--config=ggml.moe_routing.route_count=8", - "--config=ggml.workload.token_capacity=2048" - ], - "link_module": "routed_down_quantized_f16_wmma_linked" - }, - { - "name": "routed_down_q4k_plan_test", - "args": [ - "$(location :routed_down_q4k_linked)", - "--benchmark=@qwen3_moe_routed_down_q4k_q8_1_x4_prefill_512", - "--config=qwen3_moe.routed_down.expert_count=128", - "--config=qwen3_moe.routed_down.input_size=768", - "--config=qwen3_moe.routed_down.output_size=2048", - "--config=qwen3_moe.routed_down.route_count=8", - "--config=qwen3_moe.workload.token_capacity=512", - "--dry-run", - "--output-format=jsonl", - "--config=ggml.moe_routing.descriptor_expert_mask=127", - "--config=ggml.moe_routing.descriptor_partition_shift=7", - "--config=ggml.moe_routing.descriptor_row_count_shift=13", - "--config=ggml.moe_routing.expert_count=128", - "--config=ggml.moe_routing.partition_workgroup_size=128", - "--config=ggml.moe_routing.route_count=8", - "--config=ggml.workload.token_capacity=2048" - ], - "link_module": "routed_down_q4k_linked" - }, - { - "name": "routed_down_q6k_plan_test", - "args": [ - "$(location :routed_down_q6k_linked)", - "--benchmark=@qwen3_moe_routed_down_q6k_f32_wave64_decode", - "--config=qwen3_moe.routed_down.expert_count=128", - "--config=qwen3_moe.routed_down.input_size=768", - "--config=qwen3_moe.routed_down.output_size=2048", - "--config=qwen3_moe.routed_down.route_count=8", - "--config=qwen3_moe.workload.token_capacity=1", - "--dry-run", - "--output-format=jsonl" - ], - "link_module": "routed_down_q6k_linked" - }, { "name": "expert_table_partition_fused_plan_test", "args": [ @@ -2823,6 +1290,7 @@ "--config=qwen3_moe.dense_quantized.output_accumulation=0", "--config=qwen3_moe.dense_quantized.output_size=512", "--config=qwen3_moe.workload.token_capacity=1", + "--config=ggml.quantize_q8_1_x4.group_capacity=8192", "--dry-run", "--output-format=jsonl", "--config=ggml.moe_routing.descriptor_expert_mask=127", @@ -2846,6 +1314,7 @@ "--config=qwen3_moe.model.hidden_size=2048", "--config=qwen3_moe.model.rms_epsilon=0.000001", "--config=qwen3_moe.workload.token_capacity=1", + "--config=ggml.quantize_q8_1_x4.group_capacity=8192", "--dry-run", "--output-format=jsonl", "--config=ggml.moe_routing.descriptor_expert_mask=127", @@ -2858,58 +1327,13 @@ ], "link_module": "dense_linear_quantized_f16_wmma_linked" }, - { - "name": "routed_gate_up_swiglu_q4k_f16_wmma_plan_test", - "args": [ - "$(location :routed_gate_up_swiglu_q4k_f16_wmma_linked)", - "--config=qwen3_moe.routed_gate_up.expert_count=128", - "--config=qwen3_moe.routed_gate_up.input_size=2048", - "--config=qwen3_moe.routed_gate_up.output_size=768", - "--config=qwen3_moe.routed_gate_up.route_count=8", - "--config=qwen3_moe.workload.token_capacity=2048", - "--dry-run", - "--output-format=jsonl", - "--config=ggml.moe_routing.descriptor_expert_mask=127", - "--config=ggml.moe_routing.descriptor_partition_shift=7", - "--config=ggml.moe_routing.descriptor_row_count_shift=13", - "--config=ggml.moe_routing.expert_count=128", - "--config=ggml.moe_routing.partition_workgroup_size=128", - "--config=ggml.moe_routing.route_count=8", - "--config=ggml.workload.token_capacity=2048" - ], - "link_module": "routed_gate_up_swiglu_q4k_f16_wmma_linked" - }, - { - "name": "flash_attention_decode_f32_f16_wmma_plan_test", - "args": [ - "$(location qwen3_moe/flash_attention_decode_f32_f16_wmma.loom)", - "--config=qwen3_moe.attention.decode.output_partition_count=1", - "--config=qwen3_moe.attention.key_value_head_count=4", - "--config=qwen3_moe.attention.query_head_count=32", - "--dry-run", - "--output-format=jsonl" - ], - "source": "qwen3_moe/flash_attention_decode_f32_f16_wmma.loom" - }, - { - "name": "flash_attention_decode_q128_f32_f16_wmma_plan_test", - "args": [ - "$(location qwen3_moe/flash_attention_decode_q128_f32_f16_wmma.loom)", - "--config=qwen3_moe.attention.key_value_head_count=4", - "--config=qwen3_moe.attention.query_head_count=32", - "--dry-run", - "--output-format=jsonl" - ], - "source": "qwen3_moe/flash_attention_decode_q128_f32_f16_wmma.loom" - }, { "name": "owned_attention_metadata_prefill_plan_test", "args": [ "$(location ../qwen/attention_metadata.loom)", "--benchmark=@qwen_attention_metadata_prefill_512", "--dry-run", - "--output-format=jsonl", - "--sample-compilation=per_sample" + "--output-format=jsonl" ], "source": "../qwen/attention_metadata.loom" } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_postprocess_f32_f16.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_postprocess_f32_f16.loom index ed381348a7a4..6c97509cdef1 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_postprocess_f32_f16.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_postprocess_f32_f16.loom @@ -27,165 +27,6 @@ config.decl @qwen3_moe.attention.key_value_size : %value: index where [range(%va config.decl @qwen3_moe.attention.rope_mscale : f32 -// Applies one RMSNorm and two adjacent NEOX pairs. NEOX pairs corresponding -// channels from the low and high halves, while keeping each half contiguous. -// A two-pair packet therefore gives both halves naturally coalesced loads and -// stores without changing the model's pairing rule. -func.def inline @qwen3_moe_rmsnorm_neox_packet(%row_sum: f32, %head_size: f32, %epsilon: f32, %position: f32, %rope_mscale: f32, %inverse_frequencies: vector<2xf32>, %low_values: vector<2xf32>, %high_values: vector<2xf32>, %low_weights: vector<2xf32>, %high_weights: vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) { - %mean = scalar.divf %row_sum, %head_size : f32 - %biased_mean = scalar.addf %mean, %epsilon : f32 - %scale = scalar.rsqrtf %biased_mean : f32 - %scale_vector = vector.splat %scale : vector<2xf32> - %position_vector = vector.splat %position : vector<2xf32> - %inverse_two_pi = scalar.constant 0.15915494309189535 : f32 - %inverse_two_pi_vector = vector.splat %inverse_two_pi : vector<2xf32> - %normalized_low = vector.mulf %low_values, %scale_vector : vector<2xf32> - %normalized_high = vector.mulf %high_values, %scale_vector : vector<2xf32> - %scaled_low = vector.mulf %normalized_low, %low_weights : vector<2xf32> - %scaled_high = vector.mulf %normalized_high, %high_weights : vector<2xf32> - %angles = vector.mulf %position_vector, %inverse_frequencies : vector<2xf32> - %turns = vector.mulf %angles, %inverse_two_pi_vector : vector<2xf32> - %cosines = vector.costurnsf %turns : vector<2xf32> - %sines = vector.sinturnsf %turns : vector<2xf32> - %low_cosines = vector.mulf %scaled_low, %cosines : vector<2xf32> - %high_sines = vector.mulf %scaled_high, %sines : vector<2xf32> - %low_sines = vector.mulf %scaled_low, %sines : vector<2xf32> - %high_cosines = vector.mulf %scaled_high, %cosines : vector<2xf32> - %rotated_low0 = vector.subf %low_cosines, %high_sines : vector<2xf32> - %rotated_high0 = vector.addf %low_sines, %high_cosines : vector<2xf32> - %rope_mscale_vector = vector.splat %rope_mscale : vector<2xf32> - %rotated_low = vector.mulf %rotated_low0, %rope_mscale_vector : vector<2xf32> - %rotated_high = vector.mulf %rotated_high0, %rope_mscale_vector : vector<2xf32> - func.return %rotated_low, %rotated_high : vector<2xf32>, vector<2xf32> -} - -// Processes one head domain after its raw projection row is visible. Callers -// may use a larger workgroup than the packet count; inactive workitems -// contribute zero to normalization and perform no loads or stores. -template.def<@qwen3_moe.attention.postprocess.head_body> device @qwen3_moe_attention_postprocess_head_body(%publish_output: i1, %token_count: index, %cache_row_count: index, %head_domain0: index, %token0: index, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer, %query_input: buffer, %key_input: buffer, %value_input: buffer, %query_norm_weight: buffer, %key_norm_weight: buffer, %inverse_frequencies: buffer, %query_output: buffer, %key_cache: buffer, %value_cache: buffer) { - %query_size0 = config.get @qwen3_moe.attention.query_size : index - %key_value_size0 = config.get @qwen3_moe.attention.key_value_size : index - %head_size0 = config.get @qwen3_moe.attention.head_size : index - %rope_mscale = config.get @qwen3_moe.attention.rope_mscale : f32 - %epsilon = config.get @qwen3_moe.model.rms_epsilon : f32 - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index - %bounded_cache_row_count = index.assume %cache_row_count [range(%cache_row_count, 1, 1048576)] : index - %query_size, %key_value_size, %head_size = index.assume %query_size0, %key_value_size0, %head_size0 [range(%query_size0, 1, 262144), range(%key_value_size0, 1, 262144), range(%head_size0, 4, 1024), mul(%head_size0, 4), mul(%query_size0, %head_size0), mul(%key_value_size0, %head_size0)] : index, index, index - %channel0 = kernel.workitem.id : index - %c0 = index.constant 0 : index - %c2 = index.constant 2 : index - %c4 = index.constant 4 : index - %c0_f32 = scalar.constant 0.0 : f32 - %c0_offset = index.constant 0 : offset - %token, %launch_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index - %half_head_size = index.div %head_size, %c2 : index - %pair_packet_count = index.div %head_size, %c4 : index - %query_head_count = index.div %query_size, %head_size : index - %key_value_head_count = index.div %key_value_size, %head_size : index - %key_value_domain_count = index.mul %key_value_head_count, %c2 : index - %head_domain_count = index.add %query_head_count, %key_value_domain_count : index - %head_domain = index.assume %head_domain0 [lt(%head_domain0, %head_domain_count)] : index - %key_domain_end = index.add %query_head_count, %key_value_head_count : index - %is_query = index.cmp ult, %head_domain, %query_head_count : index - %is_query_or_key = index.cmp ult, %head_domain, %key_domain_end : index - %key_value_head = index.rem %head_domain, %key_value_head_count : index - %active_channel = index.cmp ult, %channel0, %pair_packet_count : index - %head_size_i32 = index.cast %head_size : index to i32 - %head_size_f32 = scalar.sitofp %head_size_i32 : i32 to f32 - %positions_noalias, %key_cache_indices_noalias, %value_cache_indices_noalias, %query_input_noalias, %key_input_noalias, %value_input_noalias, %query_norm_weight_noalias, %key_norm_weight_noalias, %inverse_frequencies_noalias, %query_output_noalias, %key_cache_noalias, %value_cache_noalias = buffer.assume.noalias %positions, %key_cache_indices, %value_cache_indices, %query_input, %key_input, %value_input, %query_norm_weight, %key_norm_weight, %inverse_frequencies, %query_output, %key_cache, %value_cache : buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer - %positions_view = buffer.view %positions_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi32> - %key_cache_indices_view = buffer.view %key_cache_indices_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi64> - %value_cache_indices_view = buffer.view %value_cache_indices_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi64> - %query_input_view = buffer.view %query_input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> - %key_input_view = buffer.view %key_input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> - %value_input_view = buffer.view %value_input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> - %query_norm_weight_view = buffer.view %query_norm_weight_noalias[%c0_offset] : buffer -> view<[%head_size]xf32> - %key_norm_weight_view = buffer.view %key_norm_weight_noalias[%c0_offset] : buffer -> view<[%head_size]xf32> - %inverse_frequencies_view = buffer.view %inverse_frequencies_noalias[%c0_offset] : buffer -> view<[%half_head_size]xf32> - %query_output_view = buffer.view %query_output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> - %key_cache_words_view = buffer.view %key_cache_noalias[%c0_offset] : buffer -> view<[%bounded_cache_row_count]x[%key_value_head_count]x[%half_head_size]xi32> - %value_cache_words_view = buffer.view %value_cache_noalias[%c0_offset] : buffer -> view<[%bounded_cache_row_count]x[%key_value_head_count]x[%half_head_size]xi32> - %row_sum = scf.if %is_query_or_key -> (f32) { - %partial_sum = scf.if %publish_output -> (f32) { - %channel_sum = scf.if %active_channel -> (f32) { - %channel = index.assume %channel0 [lt(%channel0, %pair_packet_count)] : index - %reduction_channel = index.mul %channel, %c4 : index - %reduction_values = scf.if %is_query -> (vector<4xf32>) { - %query_values = vector.load %query_input_view[%token, %head_domain, %reduction_channel] : view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> -> vector<4xf32> - scf.yield %query_values : vector<4xf32> - } else { - %key_values = vector.load %key_input_view[%token, %key_value_head, %reduction_channel] : view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> -> vector<4xf32> - scf.yield %key_values : vector<4xf32> - } - %squares = vector.mulf %reduction_values, %reduction_values : vector<4xf32> - %sum = vector.reduce %squares, %c0_f32 : vector<4xf32>, f32 - scf.yield %sum : f32 - } else { - scf.yield %c0_f32 : f32 - } - scf.yield %channel_sum : f32 - } else { - scf.yield %c0_f32 : f32 - } - %sum = kernel.workgroup.reduce %partial_sum : f32 - scf.yield %sum : f32 - } else { - scf.yield %c0_f32 : f32 - } - scf.if %publish_output { - scf.if %active_channel { - %channel = index.assume %channel0 [lt(%channel0, %pair_packet_count)] : index - %pair_channel = index.mul %channel, %c2 : index - %paired_channel = index.add %pair_channel, %half_head_size : index - %paired_word = index.add %channel, %pair_packet_count : index - %reduction_channel = index.mul %channel, %c4 : index - scf.if %is_query { - %low_values = vector.load %query_input_view[%token, %head_domain, %pair_channel] : view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> -> vector<2xf32> - %high_values = vector.load %query_input_view[%token, %head_domain, %paired_channel] : view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> -> vector<2xf32> - %low_weights = vector.load %query_norm_weight_view[%pair_channel] : view<[%head_size]xf32> -> vector<2xf32> - %high_weights = vector.load %query_norm_weight_view[%paired_channel] : view<[%head_size]xf32> -> vector<2xf32> - %inverse_frequencies_packet = vector.load %inverse_frequencies_view[%pair_channel] : view<[%half_head_size]xf32> -> vector<2xf32> - %position_i32 = view.load %positions_view[%token] : view<[%launch_token_count]xi32> -> i32 - %position = scalar.sitofp %position_i32 : i32 to f32 - %rotated_low, %rotated_high = func.call @qwen3_moe_rmsnorm_neox_packet(%row_sum, %head_size_f32, %epsilon, %position, %rope_mscale, %inverse_frequencies_packet, %low_values, %high_values, %low_weights, %high_weights) : (f32, f32, f32, f32, f32, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) - vector.store %rotated_low, %query_output_view[%token, %head_domain, %pair_channel] : vector<2xf32>, view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> - vector.store %rotated_high, %query_output_view[%token, %head_domain, %paired_channel] : vector<2xf32>, view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> - } else { - scf.if %is_query_or_key { - %low_values = vector.load %key_input_view[%token, %key_value_head, %pair_channel] : view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> -> vector<2xf32> - %high_values = vector.load %key_input_view[%token, %key_value_head, %paired_channel] : view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> -> vector<2xf32> - %low_weights = vector.load %key_norm_weight_view[%pair_channel] : view<[%head_size]xf32> -> vector<2xf32> - %high_weights = vector.load %key_norm_weight_view[%paired_channel] : view<[%head_size]xf32> -> vector<2xf32> - %inverse_frequencies_packet = vector.load %inverse_frequencies_view[%pair_channel] : view<[%half_head_size]xf32> -> vector<2xf32> - %position_i32 = view.load %positions_view[%token] : view<[%launch_token_count]xi32> -> i32 - %position = scalar.sitofp %position_i32 : i32 to f32 - %cache_index_raw = view.load %key_cache_indices_view[%token] : view<[%launch_token_count]xi64> -> i64 - %cache_index_i64 = scalar.assume %cache_index_raw [range(%cache_index_raw, 0, 1048575)] : i64 - %cache_index0 = index.cast %cache_index_i64 : i64 to index - %cache_index = index.assume %cache_index0 [lt(%cache_index0, %bounded_cache_row_count)] : index - %rotated_low, %rotated_high = func.call @qwen3_moe_rmsnorm_neox_packet(%row_sum, %head_size_f32, %epsilon, %position, %rope_mscale, %inverse_frequencies_packet, %low_values, %high_values, %low_weights, %high_weights) : (f32, f32, f32, f32, f32, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) - %half_low = vector.fptrunc %rotated_low : vector<2xf32> to vector<2xf16> - %half_high = vector.fptrunc %rotated_high : vector<2xf32> to vector<2xf16> - %packed_low = vector.bitcast %half_low : vector<2xf16> to vector<1xi32> - %packed_high = vector.bitcast %half_high : vector<2xf16> to vector<1xi32> - vector.store %packed_low, %key_cache_words_view[%cache_index, %key_value_head, %channel] : vector<1xi32>, view<[%bounded_cache_row_count]x[%key_value_head_count]x[%half_head_size]xi32> - vector.store %packed_high, %key_cache_words_view[%cache_index, %key_value_head, %paired_word] : vector<1xi32>, view<[%bounded_cache_row_count]x[%key_value_head_count]x[%half_head_size]xi32> - } else { - %cache_index_raw = view.load %value_cache_indices_view[%token] : view<[%launch_token_count]xi64> -> i64 - %cache_index_i64 = scalar.assume %cache_index_raw [range(%cache_index_raw, 0, 1048575)] : i64 - %cache_index0 = index.cast %cache_index_i64 : i64 to index - %cache_index = index.assume %cache_index0 [lt(%cache_index0, %bounded_cache_row_count)] : index - %values = vector.load %value_input_view[%token, %key_value_head, %reduction_channel] : view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> -> vector<4xf32> - %half_values = vector.fptrunc %values : vector<4xf32> to vector<4xf16> - %packed_values = vector.bitcast %half_values : vector<4xf16> to vector<2xi32> - vector.store %packed_values, %value_cache_words_view[%cache_index, %key_value_head, %pair_channel] : vector<2xi32>, view<[%bounded_cache_row_count]x[%key_value_head_count]x[%half_head_size]xi32> - } - } - } - } - template.return -} - kernel.def target(@qwen3_moe_attention_postprocess_gfx11_wave32) @qwen3_moe_attention_postprocess_f32_f16(%token_count: index, %cache_row_count: index) { %query_size = config.get @qwen3_moe.attention.query_size : index %key_value_size = config.get @qwen3_moe.attention.key_value_size : index diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_prepare_quantized.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_prepare_quantized.loom index 84a966e0f15a..c0091da9a755 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_prepare_quantized.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_prepare_quantized.loom @@ -17,7 +17,7 @@ // Scratch is fixed at 1152 bytes regardless of hidden size. template.decl @qwen3_moe.rmsnorm_quantize_q8_1_x4.body(%publish_normalized: i1, %reduction_subgroup_count0: index, %token_count: index, %token0: index, %input: buffer, %weight: buffer, %normalized_output: buffer, %q8_output: buffer) -amdgpu.target @qwen3_moe_attention_prepare_gfx11_wave32 {subgroup_size = 32} +target.decl @qwen3_moe_attention_prepare_gfx11_wave32 config.decl @qwen3_moe.model.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] @@ -103,180 +103,6 @@ kernel.def target(@qwen3_moe_attention_prepare_gfx11_wave32) @qwen3_moe_rmsnorm_ kernel.return } -// Shared RMSNorm and GGML Q8_1 x4 row producer. The caller owns launch geometry -// and supplies the logical token whose complete hidden row this workgroup owns. -// The attention export discards the ordinary F32 row, while feed-forward -// publishes it for the router without rereading and requantizing the normalized -// values. -template.def<@qwen3_moe.rmsnorm_quantize_q8_1_x4.body> device @qwen3_moe_rmsnorm_quantize_q8_1_x4_body(%publish_normalized: i1, %reduction_subgroup_count0: index, %token_count: index, %token0: index, %input: buffer, %weight: buffer, %normalized_output: buffer, %q8_output: buffer) { - %hidden_size0 = config.get @qwen3_moe.model.hidden_size : index - %epsilon = config.get @qwen3_moe.model.rms_epsilon : f32 - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index - %reduction_subgroup_count = index.assume %reduction_subgroup_count0 [range(%reduction_subgroup_count0, 1, 8)] : index - %hidden_size = index.assume %hidden_size0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128)] : index - %workitem = kernel.workitem.id : index - %subgroup0 = kernel.subgroup.id : index - %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index - %lane = kernel.subgroup.lane.id : index - %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 - %c32 = index.constant 32 : index - %c128 = index.constant 128 : index - %c256 = index.constant 256 : index - %c1024 = index.constant 1024 : index - %group_bytes = index.constant 144 : offset - %payload_byte_add = index.constant 16 : offset - %scratch_d_byte_add = index.constant 1024 : offset - %scratch_bytes = index.constant 1152 : offset - %c0_offset = index.constant 0 : offset - %c0_f32 = scalar.constant 0.0 : f32 - %c1_f32 = scalar.constant 1.0 : f32 - %c127 = scalar.constant 127.0 : f32 - %c0_f32x4 = vector.constant 0.0 : vector<4xf32> - %valid_token = index.cmp ult, %token0, %bounded_token_count : index - %safe_token0 = scf.select %valid_token, %token0, %c0 : index - %token, %launch_token_count = index.assume %safe_token0, %bounded_token_count [lt(%safe_token0, %bounded_token_count)] : index, index - %hidden_size_i32 = index.cast %hidden_size : index to i32 - %hidden_size_f32 = scalar.sitofp %hidden_size_i32 : i32 to f32 - %physical_group_count = index.div %hidden_size, %c128 : index - %row_bytes = index.scale %physical_group_count, %group_bytes : index, offset -> offset - %token_output_byte_base = index.scale %token, %row_bytes : index, offset -> offset - %input_noalias, %weight_noalias, %q8_output_noalias = buffer.assume.noalias %input, %weight, %q8_output : buffer, buffer, buffer - %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> - %weight_view = buffer.view %weight_noalias[%c0_offset] : buffer -> view<[%hidden_size]xf32> - %normalized_output_view = buffer.view %normalized_output[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> - %scratch = buffer.alloca align(16) %scratch_bytes : buffer - %scratch_values = buffer.view %scratch[%c0_offset] : buffer -> view<256xf32> - %scratch_d = buffer.view %scratch[%scratch_d_byte_add] : buffer -> view<32xf32> - // Reduce the complete row before any block-local quantization. - %thread_sum = scf.for %channel = [%workitem to %hidden_size step %c256](%running_sum = %c0_f32 : f32) -> (f32) { - %value = view.load %input_view[%token, %channel] : view<[%launch_token_count]x[%hidden_size]xf32> -> f32 - %square = scalar.mulf %value, %value : f32 - %next_sum = scalar.addf %running_sum, %square : f32 - scf.yield %next_sum : f32 - } - %subgroup_sum = kernel.subgroup.reduce %thread_sum : f32 - %is_subgroup_leader = index.cmp eq, %lane, %c0 : index - scf.if %is_subgroup_leader { - view.store %subgroup_sum, %scratch_values[%subgroup] : f32, view<256xf32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %is_reduction_subgroup = index.cmp eq, %subgroup, %c0 : index - %is_reduction_lane = index.cmp ult, %lane, %reduction_subgroup_count : index - %loads_subgroup_sum = scalar.andi %is_reduction_subgroup, %is_reduction_lane : i1 - %subgroup_partial = scf.if %loads_subgroup_sum -> (f32) { - %value = view.load %scratch_values[%lane] : view<256xf32> -> f32 - scf.yield %value : f32 - } else { - scf.yield %c0_f32 : f32 - } - %row_sum = kernel.subgroup.reduce %subgroup_partial : f32 - %writes_scale = scalar.andi %is_reduction_subgroup, %is_subgroup_leader : i1 - scf.if %writes_scale { - %mean = scalar.divf %row_sum, %hidden_size_f32 : f32 - %biased_mean = scalar.addf %mean, %epsilon : f32 - %scale = scalar.rsqrtf %biased_mean : f32 - view.store %scale, %scratch_values[%c0] : f32, view<256xf32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %row_scale = view.load %scratch_values[%c0] : view<256xf32> -> f32 - %row_scale_vector = vector.splat %row_scale : vector<4xf32> - %publishes_normalized = scalar.andi %publish_normalized, %valid_token : i1 - // Reuse one scratch frame for each 1024-element stripe. Hidden sizes need - // only be divisible by 128; inactive workitems in the final stripe carry - // zeros and never publish. - scf.for %stripe_base = [%c0 to %hidden_size step %c1024] { - %word_element_add = index.mul %workitem, %c4 : index - %channel = index.add %stripe_base, %word_element_add : index - %valid_word = index.cmp ult, %channel, %hidden_size : index - %mask = vector.mask.range [%channel to %hidden_size step %c1] : index -> vector<4xi1> - %input_values = vector.load.mask %input_view[%token, %channel], %mask, %c0_f32x4 : view<[%launch_token_count]x[%hidden_size]xf32>, vector<4xi1>, vector<4xf32> - %learned_weights = vector.load.mask %weight_view[%channel], %mask, %c0_f32x4 : view<[%hidden_size]xf32>, vector<4xi1>, vector<4xf32> - %normalized0 = vector.mulf %input_values, %row_scale_vector : vector<4xf32> - %normalized = vector.mulf %normalized0, %learned_weights : vector<4xf32> - scf.if %publishes_normalized { - vector.store.mask %normalized, %normalized_output_view[%token, %channel], %mask : vector<4xf32>, view<[%launch_token_count]x[%hidden_size]xf32>, vector<4xi1> - } - %absolute_values = vector.absf %normalized : vector<4xf32> - %thread_max = vector.reduce %absolute_values, %c0_f32 : vector<4xf32>, f32 - view.store %thread_max, %scratch_values[%workitem] : f32, view<256xf32> - kernel.barrier scope(workgroup) ordering(acq_rel) - %word_in_block = index.rem %workitem, %c8 : index - %block_in_stripe = index.div %workitem, %c8 : index - %is_block_leader = index.cmp eq, %word_in_block, %c0 : index - %writes_block_d = scalar.andi %valid_word, %is_block_leader : i1 - scf.if %writes_block_d { - %cohort_base = index.mul %block_in_stripe, %c8 : index - %cohort_maxima = vector.load %scratch_values[%cohort_base] : view<256xf32> -> vector<8xf32> - %amax = vector.reduce %cohort_maxima, %c0_f32 : vector<8xf32>, f32 - %d = scalar.divf %amax, %c127 : f32 - view.store %d, %scratch_d[%block_in_stripe] : f32, view<32xf32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %d = scf.if %valid_word -> (f32) { - %block_d = view.load %scratch_d[%block_in_stripe] : view<32xf32> -> f32 - scf.yield %block_d : f32 - } else { - scf.yield %c0_f32 : f32 - } - %d_nonzero = scalar.cmpf one, %d, %c0_f32 : f32 - %d_inverse = scf.if %d_nonzero -> (f32) { - %inverse = scalar.divf %c1_f32, %d : f32 - scf.yield %inverse : f32 - } else { - scf.yield %c0_f32 : f32 - } - %d_inverse_vector = vector.splat %d_inverse : vector<4xf32> - %scaled_values = vector.mulf %normalized, %d_inverse_vector : vector<4xf32> - %rounded_values = vector.roundf %scaled_values : vector<4xf32> - %quantized_values = vector.fptosi %rounded_values : vector<4xf32> to vector<4xi8> - %packed_word = vector.bitcast %quantized_values : vector<4xi8> to vector<1xi32> - %publishes_q8_word = scalar.andi %valid_word, %valid_token : i1 - scf.if %publishes_q8_word { - %q8_block = index.div %channel, %c32 : index - %physical_group = index.div %q8_block, %c4 : index - %block_in_group = index.rem %q8_block, %c4 : index - %group_byte_add = index.scale %physical_group, %group_bytes : index, offset -> offset - %group_byte_offset = index.add %token_output_byte_base, %group_byte_add : offset - %payload_byte_offset = index.add %group_byte_offset, %payload_byte_add : offset - %group_ds = buffer.view %q8_output_noalias[%group_byte_offset] : buffer -> view<8xf16> - %group_qs = buffer.view %q8_output_noalias[%payload_byte_offset] : buffer -> view<32xi32> - %block_word_base = index.mul %block_in_group, %c8 : index - %packed_word_index0 = index.add %block_word_base, %word_in_block : index - %packed_word_index = index.assume %packed_word_index0 [range(%packed_word_index0, 0, 31)] : index - vector.store %packed_word, %group_qs[%packed_word_index] : vector<1xi32>, view<32xi32> - } - %thread_quantized_sum = vector.reduce %rounded_values, %c0_f32 : vector<4xf32>, f32 - view.store %thread_quantized_sum, %scratch_values[%workitem] : f32, view<256xf32> - kernel.barrier scope(workgroup) ordering(acq_rel) - %publishes_block_ds = scalar.andi %writes_block_d, %valid_token : i1 - scf.if %publishes_block_ds { - %cohort_base = index.mul %block_in_stripe, %c8 : index - %cohort_sums = vector.load %scratch_values[%cohort_base] : view<256xf32> -> vector<8xf32> - %quantized_sum = vector.reduce %cohort_sums, %c0_f32 : vector<8xf32>, f32 - %s = scalar.mulf %quantized_sum, %d : f32 - %q8_block = index.div %channel, %c32 : index - %physical_group = index.div %q8_block, %c4 : index - %block_in_group = index.rem %q8_block, %c4 : index - %group_byte_add = index.scale %physical_group, %group_bytes : index, offset -> offset - %group_byte_offset = index.add %token_output_byte_base, %group_byte_add : offset - %group_ds = buffer.view %q8_output_noalias[%group_byte_offset] : buffer -> view<8xf16> - %d_f16 = scalar.fptrunc %d : f32 to f16 - %s_f16 = scalar.fptrunc %s : f32 to f16 - %ds_index = index.mul %block_in_group, %c2 : index - %s_index = index.add %ds_index, %c1 : index - view.store %d_f16, %group_ds[%ds_index] : f16, view<8xf16> - view.store %s_f16, %group_ds[%s_index] : f16, view<8xf16> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - } - template.return -} - // Fuses attention RMSNorm directly into the Q8_1 x4 row consumed by decode // projections. The output has hidden_size / 128 physical groups of 144 bytes. kernel.def target(@qwen3_moe_attention_prepare_gfx11_wave32) @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4(%token_count: index) { diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_qkv_postprocess_fused.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_qkv_postprocess_fused.loom index c6170dd87e72..62d91587f23f 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_qkv_postprocess_fused.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_qkv_postprocess_fused.loom @@ -37,65 +37,6 @@ kernel.decl @qwen3_moe_attention_qkv_quantized(%token_count$61: index) launch(%t kernel.decl @qwen3_moe_attention_postprocess_f32_f16(%token_count$70: index, %cache_row_count$71: index) launch(%token_count$72: index, %cache_row_count$73: index, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer, %query_input: buffer, %key_input: buffer, %value_input: buffer, %query_norm_weight: buffer, %key_norm_weight: buffer, %inverse_frequencies: buffer, %query_output: buffer, %key_cache: buffer, %value_cache: buffer) -// The fused export sources the value storage format from model config. -template.def<@qwen3_moe.attention.qkv_postprocess_fused_decode.body> device @qwen3_moe_attention_qkv_postprocess_fused_decode_body(%value_uses_q6_index: index, %token_count: index, %cache_row_count: index, %q8_input: buffer, %query_weight: buffer, %key_weight: buffer, %value_weight: buffer, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer, %query_output_raw: buffer, %key_output_raw: buffer, %value_output_raw: buffer, %query_norm_weight: buffer, %key_norm_weight: buffer, %inverse_frequencies: buffer, %query_output: buffer, %key_cache: buffer, %value_cache: buffer, %completion_counters: buffer) { - %publish_projection_output = scalar.constant true : i1 - %projection_token = kernel.workgroup.id : index - template.apply<@qwen3_moe.attention.qkv_quantized.body>(%value_uses_q6_index, %publish_projection_output, %token_count, %projection_token, %q8_input, %query_weight, %key_weight, %value_weight, %query_output_raw, %key_output_raw, %value_output_raw) : (index, i1, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) - %query_size0 = config.get @qwen3_moe.attention.query_size : index - %key_value_size0 = config.get @qwen3_moe.attention.key_value_size : index - %head_size0 = config.get @qwen3_moe.attention.head_size : index - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index - %query_size, %key_value_size, %head_size = index.assume %query_size0, %key_value_size0, %head_size0 [range(%query_size0, 128, 262144), range(%key_value_size0, 128, 262144), range(%head_size0, 128, 128), mul(%query_size0, %head_size0), mul(%key_value_size0, %head_size0)] : index, index, index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c2 = index.constant 2 : index - %c8 = index.constant 8 : index - %c0_i32 = scalar.constant 0 : i32 - %c1_i32 = scalar.constant 1 : i32 - %c0_offset = index.constant 0 : offset - %counter_scratch_bytes = index.constant 4 : offset - %channel_tile = kernel.workgroup.id : index - %workitem = kernel.workitem.id : index - %token = index.assume %c0 [lt(%c0, %bounded_token_count)] : index - %head_tile_count = index.div %head_size, %c8 : index - %head_domain0 = index.div %channel_tile, %head_tile_count : index - %query_head_count = index.div %query_size, %head_size : index - %key_value_head_count = index.div %key_value_size, %head_size : index - %key_value_domain_count = index.mul %key_value_head_count, %c2 : index - %head_domain_count = index.add %query_head_count, %key_value_domain_count : index - %head_domain, %completion_counter_count = index.assume %head_domain0, %head_domain_count [lt(%head_domain0, %head_domain_count)] : index, index - %is_arrival_workitem = index.cmp eq, %workitem, %c0 : index - %completion_counters_aligned = buffer.assume.alignment %completion_counters {minimum_alignment = 16} : buffer - %completion_counters_view = buffer.view %completion_counters_aligned[%c0_offset] : buffer -> view<[%completion_counter_count]xi32> - %counter_scratch = buffer.alloca align(4) %counter_scratch_bytes : buffer - %counter_scratch_view = buffer.view %counter_scratch[%c0_offset] : buffer -> view<1xi32> - // Publish every producer's projection stores before the leader advances one - // workgroup arrival. The last arrival then acquires the complete head. - kernel.barrier scope(workgroup) ordering(release) - scf.if %is_arrival_workitem { - %old_counter = view.atomic.rmw %c1_i32, %completion_counters_view[%head_domain] {ordering = acq_rel, scope = device} : i32, view<[%completion_counter_count]xi32> -> i32 - view.store %old_counter, %counter_scratch_view[%c0] : i32, view<1xi32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %old_counter = view.load %counter_scratch_view[%c0] : view<1xi32> -> i32 - %head_tile_count_i32 = index.cast %head_tile_count : index to i32 - %last_head_tile_i32 = scalar.subi %head_tile_count_i32, %c1_i32 : i32 - %negative_head_tile_count_i32 = scalar.subi %c0_i32, %head_tile_count_i32 : i32 - %is_last_head_tile = scalar.cmpi eq, %old_counter, %last_head_tile_i32 : i32 - scf.if %is_last_head_tile { - kernel.barrier scope(workgroup) ordering(acquire) - %publish_postprocess_output = scalar.constant true : i1 - template.apply<@qwen3_moe.attention.postprocess.head_body>(%publish_postprocess_output, %bounded_token_count, %cache_row_count, %head_domain, %token, %positions, %key_cache_indices, %value_cache_indices, %query_output_raw, %key_output_raw, %value_output_raw, %query_norm_weight, %key_norm_weight, %inverse_frequencies, %query_output, %key_cache, %value_cache) : (i1, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) - // The counter becomes reusable only after all semantic output stores. - kernel.barrier scope(workgroup) ordering(release) - scf.if %is_arrival_workitem { - view.atomic.reduce %negative_head_tile_count_i32, %completion_counters_view[%head_domain] {ordering = release, scope = device} : i32, view<[%completion_counter_count]xi32> - } - } - template.return -} - template.def<@qwen3_moe.attention.qkv_postprocess_fused_decode.launch> @qwen3_moe_attention_qkv_postprocess_fused_decode_launch() -> (index, index, index) { %query_size = config.get @qwen3_moe.attention.query_size : index %key_value_size = config.get @qwen3_moe.attention.key_value_size : index diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_qkv_quantized.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_qkv_quantized.loom index 3a5d0ddc5f79..4967c5942e92 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_qkv_quantized.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_qkv_quantized.loom @@ -42,105 +42,6 @@ kernel.decl @qwen3_moe_dense_linear_q4k_q8_1_x4(%token_count$41: index) launch(% kernel.decl @ggml_linear_q6k_q8_1_x4(%token_count$46: index, %input_size$47: index, %output_size$48: index) launch(%token_count$49: index, %input_size$50: index, %output_size$51: index, %q8_input: buffer, %weight: buffer, %output: buffer) -// Device body shared by the ordinary projection and completion-fused -// postprocess exports. Each caller owns its boundary after the raw row stores. -template.def<@qwen3_moe.attention.qkv_quantized.body> device @qwen3_moe_attention_qkv_quantized_body(%value_uses_q6_index: index, %publish_output: i1, %token_count: index, %token0: index, %q8_input: buffer, %query_weight: buffer, %key_weight: buffer, %value_weight: buffer, %query_output: buffer, %key_output: buffer, %value_output: buffer) { - %hidden_size0 = config.get @qwen3_moe.model.hidden_size : index - %query_size0 = config.get @qwen3_moe.attention.query_size : index - %key_value_size0 = config.get @qwen3_moe.attention.key_value_size : index - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index - %hidden_size = index.assume %hidden_size0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128)] : index - %query_size, %key_value_size = index.assume %query_size0, %key_value_size0 [range(%query_size0, 1, 262144), range(%key_value_size0, 1, 262144), mul(%query_size0, %key_value_size0)] : index, index - %channel_tile = kernel.workgroup.id : index - %subgroup0 = kernel.subgroup.id : index - %lane = kernel.subgroup.lane.id : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c2 = index.constant 2 : index - %c8 = index.constant 8 : index - %c128 = index.constant 128 : index - %c144_bytes = index.constant 144 : offset - %c210_bytes = index.constant 210 : offset - %c256 = index.constant 256 : index - %c0_i32 = scalar.constant 0 : i32 - %c0_f32 = scalar.constant 0.0 : f32 - %c0_offset = index.constant 0 : offset - %token, %launch_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index - %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index - %channel_base = index.mul %channel_tile, %c8 : index - %global_channel = index.add %channel_base, %subgroup : index - %key_value_end = index.add %query_size, %key_value_size : index - %key_value_output_size = index.mul %key_value_size, %c2 : index - %total_output_size = index.add %query_size, %key_value_output_size : index - %valid_channel = index.cmp ult, %global_channel, %total_output_size : index - %is_query = index.cmp ult, %global_channel, %query_size : index - %is_key = index.cmp ult, %global_channel, %key_value_end : index - %key_value_channel = index.rem %global_channel, %key_value_size : index - %value_uses_q6 = index.cmp eq, %value_uses_q6_index, %c1 : index - %lane_i32 = index.cast %lane : index to i32 - %is_lane_zero = scalar.cmpi eq, %lane_i32, %c0_i32 : i32 - %quant_block_count = index.div %hidden_size, %c256 : index - %q4_row_bytes = index.scale %quant_block_count, %c144_bytes : index, offset -> offset - %q6_row_bytes = index.scale %quant_block_count, %c210_bytes : index, offset -> offset - %q8_group_count = index.div %hidden_size, %c128 : index - %q8_row_bytes = index.scale %q8_group_count, %c144_bytes : index, offset -> offset - %q8_row_byte_base = index.scale %token, %q8_row_bytes : index, offset -> offset - %q8_noalias, %query_weight_noalias, %key_weight_noalias, %value_weight_noalias, %query_output_noalias, %key_output_noalias, %value_output_noalias = buffer.assume.noalias %q8_input, %query_weight, %key_weight, %value_weight, %query_output, %key_output, %value_output : buffer, buffer, buffer, buffer, buffer, buffer, buffer - %lane_sum = scf.if %publish_output -> (f32) { - %channel_sum = scf.if %valid_channel -> (f32) { - %projection_sum = scf.if %is_query -> (f32) { - %row_byte_base = index.scale %global_channel, %q4_row_bytes : index, offset -> offset - %sum = func.call @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%hidden_size, %query_weight_noalias, %row_byte_base, %q8_noalias, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) - scf.yield %sum : f32 - } else { - %key_or_value_sum = scf.if %is_key -> (f32) { - %row_byte_base = index.scale %key_value_channel, %q4_row_bytes : index, offset -> offset - %sum = func.call @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%hidden_size, %key_weight_noalias, %row_byte_base, %q8_noalias, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) - scf.yield %sum : f32 - } else { - %value_sum = scf.if %value_uses_q6 -> (f32) { - %row_byte_base = index.scale %key_value_channel, %q6_row_bytes : index, offset -> offset - %sum = func.call @ggml_q6k_q8_1_x4_row_lane(%hidden_size, %value_weight_noalias, %row_byte_base, %q8_noalias, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) - scf.yield %sum : f32 - } else { - %row_byte_base = index.scale %key_value_channel, %q4_row_bytes : index, offset -> offset - %sum = func.call @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%hidden_size, %value_weight_noalias, %row_byte_base, %q8_noalias, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) - scf.yield %sum : f32 - } - scf.yield %value_sum : f32 - } - scf.yield %key_or_value_sum : f32 - } - scf.yield %projection_sum : f32 - } else { - scf.yield %c0_f32 : f32 - } - scf.yield %channel_sum : f32 - } else { - scf.yield %c0_f32 : f32 - } - %dot = kernel.subgroup.reduce %lane_sum : f32 - %query_output_view = buffer.view %query_output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%query_size]xf32> - %key_output_view = buffer.view %key_output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%key_value_size]xf32> - %value_output_view = buffer.view %value_output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%key_value_size]xf32> - scf.if %publish_output { - scf.if %valid_channel { - scf.if %is_lane_zero { - scf.if %is_query { - view.store %dot, %query_output_view[%token, %global_channel] : f32, view<[%launch_token_count]x[%query_size]xf32> - } else { - scf.if %is_key { - view.store %dot, %key_output_view[%token, %key_value_channel] : f32, view<[%launch_token_count]x[%key_value_size]xf32> - } else { - view.store %dot, %value_output_view[%token, %key_value_channel] : f32, view<[%launch_token_count]x[%key_value_size]xf32> - } - } - } - } - } - template.return -} - kernel.def target(@qwen3_moe_attention_qkv_gfx11_wave32) @qwen3_moe_attention_qkv_quantized(%token_count: index) { %query_size = config.get @qwen3_moe.attention.query_size : index %key_value_size = config.get @qwen3_moe.attention.key_value_size : index diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/dense_linear_quantized_f16_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/dense_linear_quantized_f16_wmma.loom index 07d4d1e03bbc..6a62f60e558b 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/dense_linear_quantized_f16_wmma.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/dense_linear_quantized_f16_wmma.loom @@ -17,12 +17,6 @@ config.decl @qwen3_moe.dense_quantized.output_accumulation : %value: index where config.decl @qwen3_moe.workload.token_capacity : %value: index where [range(%value, 1, 2048)] -// These declarations are used only by the differential correctness case. The -// dense production kernel has no routing input or dependency. -kernel.decl @ggml_moe_build_expert_table(%token_count$34: index, %route_count$35: index, %route_stride$36: index, %expert_count$37: index) launch(%token_count$38: index, %route_count$39: index, %route_stride$40: index, %expert_count$41: index, %route_ids: buffer, %expert_table: buffer) - -kernel.decl @qwen3_moe_routed_linear_q4k_f16_wmma(%token_count$44: index) launch(%token_count$45: index, %input: buffer, %expert_table: buffer, %weight: buffer, %output: buffer) - // Shared raw-layout Q4_K row contraction primitive. Each lane consumes both // nibbles of one packed-code load. func.decl @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32) @@ -31,69 +25,6 @@ kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count$63: index, %input_size$64: i kernel.decl @qwen3_moe_rmsnorm_f32_quantize_q8_1_x4(%token_count$78: index) launch(%token_count$79: index, %input: buffer, %rms_weight: buffer, %normalized_f32_output: buffer, %q8_1_x4_output: buffer) -template.def<@qwen3_moe.dense.q4k_q8_1_x4.body> device @qwen3_moe_dense_linear_q4k_q8_1_x4_body(%publish_output: i1, %token_count: index, %token0: index, %q8_input: buffer, %weight: buffer, %output: buffer) { - %input_size = config.get @qwen3_moe.dense_quantized.input_size : index - %output_size = config.get @qwen3_moe.dense_quantized.output_size : index - %output_accumulation = config.get @qwen3_moe.dense_quantized.output_accumulation : index - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index - %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index - %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144)] : index - %channel_tile = kernel.workgroup.id : index - %subgroup0 = kernel.subgroup.id : index - %lane = kernel.subgroup.lane.id : index - %c1 = index.constant 1 : index - %c8 = index.constant 8 : index - %c128 = index.constant 128 : index - %c144_bytes = index.constant 144 : offset - %c256 = index.constant 256 : index - %c0_i32 = scalar.constant 0 : i32 - %c0_f32 = scalar.constant 0.0 : f32 - %c0_offset = index.constant 0 : offset - %token, %launch_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index - %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index - %channel_base = index.mul %channel_tile, %c8 : index - %channel = index.add %channel_base, %subgroup : index - %valid_channel = index.cmp ult, %channel, %bounded_output_size : index - %lane_i32 = index.cast %lane : index to i32 - %is_lane_zero = scalar.cmpi eq, %lane_i32, %c0_i32 : i32 - %q4_block_count = index.div %bounded_input_size, %c256 : index - %weight_row_bytes = index.scale %q4_block_count, %c144_bytes : index, offset -> offset - %weight_row_byte_base = index.scale %channel, %weight_row_bytes : index, offset -> offset - %q8_group_count = index.div %bounded_input_size, %c128 : index - %q8_row_bytes = index.scale %q8_group_count, %c144_bytes : index, offset -> offset - %q8_row_byte_base = index.scale %token, %q8_row_bytes : index, offset -> offset - %q8_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %q8_input, %weight, %output : buffer, buffer, buffer - %lane_sum = scf.if %publish_output -> (f32) { - %channel_sum = scf.if %valid_channel -> (f32) { - %sum = func.call @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%bounded_input_size, %weight_noalias, %weight_row_byte_base, %q8_noalias, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) - scf.yield %sum : f32 - } else { - scf.yield %c0_f32 : f32 - } - scf.yield %channel_sum : f32 - } else { - scf.yield %c0_f32 : f32 - } - %dot = kernel.subgroup.reduce %lane_sum : f32 - %accumulates_output = index.cmp eq, %output_accumulation, %c1 : index - scf.if %publish_output { - scf.if %valid_channel { - scf.if %is_lane_zero { - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%bounded_output_size]xf32> - %result = scf.if %accumulates_output -> (f32) { - %residual = view.load %output_view[%token, %channel] : view<[%launch_token_count]x[%bounded_output_size]xf32> -> f32 - %sum = scalar.addf %residual, %dot : f32 - scf.yield %sum : f32 - } else { - scf.yield %dot : f32 - } - view.store %result, %output_view[%token, %channel] : f32, view<[%launch_token_count]x[%bounded_output_size]xf32> - } - } - } - template.return -} - // Direct decode entry point. This complements the grouped WMMA schedule below: // it gives up cross-token weight reuse in exchange for doing no padded matrix // work when fewer than 32 tokens are available. @@ -118,86 +49,7 @@ kernel.def target(@qwen3_moe_dense_gfx11_wave32) @qwen3_moe_dense_linear_q4k_q8_ kernel.return } -// Decode-only route that publishes the normalized F32 and Q8_1 x4 rows -// consumed by the following feed-forward boundary. -kernel.def target(@qwen3_moe_attention_prepare_gfx11_wave32) @qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8(%token_count: index) { - %output_size = config.get @qwen3_moe.dense_quantized.output_size : index - %c1 = index.constant 1 : index - %c8 = index.constant 8 : index - %c256 = index.constant 256 : index - %output_tiles = index.div %output_size, %c8 : index - kernel.launch.config workgroups(%output_tiles, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index -} launch(%token_count: index, %q8_input: buffer, %weight: buffer, %output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counter: buffer, %next_q8_output: buffer) { - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index - %publish_output = scalar.constant true : i1 - %projection_token = kernel.workgroup.id : index - template.apply<@qwen3_moe.dense.q4k_q8_1_x4.body>(%publish_output, %bounded_token_count, %projection_token, %q8_input, %weight, %output) : (i1, index, index, buffer, buffer, buffer) - %output_size0 = config.get @qwen3_moe.dense_quantized.output_size : index - %output_size = index.assume %output_size0 [range(%output_size0, 128, 32768), mul(%output_size0, 128)] : index - %token0 = kernel.workgroup.id : index - %token = index.assume %token0 [lt(%token0, %bounded_token_count)] : index - %c0 = index.constant 0 : index - %c8 = index.constant 8 : index - %workitem = kernel.workitem.id : index - %output_tile_count = index.div %output_size, %c8 : index - %is_arrival_workitem = index.cmp eq, %workitem, %c0 : index - %c0_i32 = scalar.constant 0 : i32 - %c1_i32 = scalar.constant 1 : i32 - %c0_offset = index.constant 0 : offset - %counter_scratch_bytes = index.constant 4 : offset - %output_noalias, %norm_weight_noalias, %normalized_output_noalias, %completion_counter_noalias, %next_q8_output_noalias = buffer.assume.noalias %output, %norm_weight, %normalized_output, %completion_counter, %next_q8_output : buffer, buffer, buffer, buffer, buffer - %completion_counter_aligned = buffer.assume.alignment %completion_counter_noalias {minimum_alignment = 16} : buffer - %completion_counter_view = buffer.view %completion_counter_aligned[%c0_offset] : buffer -> view<1xi32> - %counter_scratch = buffer.alloca align(4) %counter_scratch_bytes : buffer - %counter_scratch_view = buffer.view %counter_scratch[%c0_offset] : buffer -> view<1xi32> - // Publish every producer's residual stores before the leader advances one - // workgroup arrival. The last arrival then acquires the complete row. - kernel.barrier scope(workgroup) ordering(release) - scf.if %is_arrival_workitem { - %old_counter = view.atomic.rmw %c1_i32, %completion_counter_view[%c0] {ordering = acq_rel, scope = device} : i32, view<1xi32> -> i32 - view.store %old_counter, %counter_scratch_view[%c0] : i32, view<1xi32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %old_counter = view.load %counter_scratch_view[%c0] : view<1xi32> -> i32 - %output_tile_count_i32 = index.cast %output_tile_count : index to i32 - %last_output_tile_i32 = scalar.subi %output_tile_count_i32, %c1_i32 : i32 - %negative_output_tile_count_i32 = scalar.subi %c0_i32, %output_tile_count_i32 : i32 - %is_last_output_tile = scalar.cmpi eq, %old_counter, %last_output_tile_i32 : i32 - scf.if %is_last_output_tile { - kernel.barrier scope(workgroup) ordering(acquire) - %publish_normalized = scalar.constant true : i1 - template.apply<@qwen3_moe.rmsnorm_quantize_q8_1_x4.body>(%publish_normalized, %c8, %bounded_token_count, %token, %output_noalias, %norm_weight_noalias, %normalized_output_noalias, %next_q8_output_noalias) : (i1, index, index, index, buffer, buffer, buffer, buffer) - // Reset only after every normalized F32 and Q8 store completes. - kernel.barrier scope(workgroup) ordering(release) - scf.if %is_arrival_workitem { - view.atomic.reduce %negative_output_tile_count_i32, %completion_counter_view[%c0] {ordering = release, scope = device} : i32, view<1xi32> - } - } - kernel.return -} - -// Keep a nonzero production-size differential for the retained raw Q4xQ8 -// control without depending on the retired dense Q4/F16 implementation. -check.case public @qwen3_moe_dense_linear_q4k_q8_1_x4_o_differential_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(4096) : index - %route_count = check.literal value(1) : index - %route_stride = check.literal value(1) : index - %expert_count = check.literal value(1) : index - %input = check.generate.fill value(0.00390625) : tensor<1x4096xf32> - %q8_input = check.generate.fill value(0) : tensor<1x4608xi8> - %route_ids = check.generate.fill value(0) : tensor<1x1xi32> - %expert_table = check.generate.fill value(-1) : tensor<2xi32> - %weight = check.generate.fill value(34) : tensor<1x2048x16x144xi8> - %expected = check.generate.fill value(0.0) : tensor<1x2048xf32> - %actual = check.generate.fill value(0.0) : tensor<1x2048xf32> - kernel.launch @ggml_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<1x1xi32>, tensor<2xi32>) - kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<1x4096xf32>, tensor<1x4608xi8>) - kernel.launch @qwen3_moe_routed_linear_q4k_f16_wmma[%token_count](%token_count, %input, %expert_table, %weight, %expected) : [index](index, tensor<1x4096xf32>, tensor<2xi32>, tensor<1x2048x16x144xi8>, tensor<1x2048xf32>) - kernel.launch @qwen3_moe_dense_linear_q4k_q8_1_x4[%token_count](%token_count, %q8_input, %weight, %actual) : [index](index, tensor<1x4608xi8>, tensor<1x2048x16x144xi8>, tensor<1x2048xf32>) - check.expect.close actual(%actual) expected(%expected) atol(0.25) rtol(0.01) nan(same) : tensor<1x2048xf32> - check.return -} +kernel.decl @qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8(%token_count$1: index) launch(%token_count$2: index, %q8_input: buffer, %weight: buffer, %output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counter: buffer, %next_q8_output: buffer) // A zero projection leaves the nonzero caller-provided residual unchanged. // Overwrite publication would instead produce zero, making the accumulation @@ -321,8 +173,6 @@ check.case public @qwen3_moe_dense_linear_q4k_q8_1_x4_o_projection_benchmark_cas check.return } -check.benchmark<@qwen3_moe_dense_linear_q4k_q8_1_x4_o_differential_case> @qwen3_moe_dense_linear_q4k_q8_1_x4_o_differential - check.benchmark<@qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8_benchmark_case> @qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8_decode check.benchmark<@qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8_composed_benchmark_case> @qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8_composed_decode diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/expert_table_partition_fused.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/expert_table_partition_fused.loom index ae5571e1ad33..2eaba9e6d6f3 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/expert_table_partition_fused.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/expert_table_partition_fused.loom @@ -6,169 +6,13 @@ // The last workgroup acquires all counts, compacts them into deterministic // 32-row descriptors, and resets the counter so reusable command buffers can // issue the kernel again against the same storage. -amdgpu.target @qwen3_moe_expert_table_partition_gfx11_wave32 {subgroup_size = 32} +target.decl @qwen3_moe_expert_table_partition_gfx11_wave32 kernel.decl @ggml_moe_build_expert_table(%token_count$0: index, %route_count$1: index, %route_stride$2: index, %expert_count$3: index) launch(%token_count$4: index, %route_count$5: index, %route_stride$6: index, %expert_count$7: index, %route_ids: buffer, %expert_table: buffer) kernel.decl @ggml_moe_build_expert_partition_table(%token_count$10: index, %route_count$11: index, %expert_count$12: index) launch(%token_count$13: index, %route_count$14: index, %expert_count$15: index, %expert_table: buffer, %partition_table: buffer) -kernel.def target(@qwen3_moe_expert_table_partition_gfx11_wave32) @qwen3_moe_build_expert_table_partition_prefill_512(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index) { - %c1 = index.constant 1 : index - %launch_expert_count = index.constant 128 : index - %workgroup_size = index.constant 256 : index - kernel.launch.config workgroups(%launch_expert_count, %c1, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index -} launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %route_ids: buffer, %expert_table: buffer, %partition_table: buffer, %completion_counter: buffer) { - %bounded_token_count = index.assume %token_count [range(%token_count, 512, 512)] : index - %bounded_route_count = index.assume %route_count [range(%route_count, 8, 8)] : index - %bounded_route_stride = index.assume %route_stride [range(%route_stride, 8, 8)] : index - %bounded_expert_count = index.assume %expert_count [range(%expert_count, 128, 128)] : index - %expert0 = kernel.workgroup.id : index - %expert, %launch_expert_count = index.assume %expert0, %bounded_expert_count [lt(%expert0, %bounded_expert_count)] : index, index - %lane0 = kernel.workitem.id : index - %lane = index.assume %lane0 [range(%lane0, 0, 255)] : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c31 = index.constant 31 : index - %c32 = index.constant 32 : index - %workgroup_size = index.constant 256 : index - %c0_i32 = scalar.constant 0 : i32 - %c1_i32 = scalar.constant 1 : i32 - %c7_i32 = scalar.constant 7 : i32 - %c13_i32 = scalar.constant 13 : i32 - %c0_offset = index.constant 0 : offset - %c4_bytes = index.constant 4 : offset - %assignment_count = index.mul %bounded_token_count, %bounded_route_count : index - %assignment_table_byte_base = index.scale %launch_expert_count, %c4_bytes : index, offset -> offset - %route_ids_noalias, %expert_table_noalias, %partition_table_noalias, %completion_counter_noalias = buffer.assume.noalias %route_ids, %expert_table, %partition_table, %completion_counter : buffer, buffer, buffer, buffer - %completion_counter_aligned = buffer.assume.alignment %completion_counter_noalias {minimum_alignment = 16} : buffer - %route_view = buffer.view %route_ids_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_route_stride]xi32> - %count_view = buffer.view %expert_table_noalias[%c0_offset] : buffer -> view<[%launch_expert_count]xi32> - %assignment_view = buffer.view %expert_table_noalias[%assignment_table_byte_base] : buffer -> view<[%launch_expert_count]x[%bounded_token_count]xi32> - %completion_counter_view = buffer.view %completion_counter_aligned[%c0_offset] : buffer -> view<1xi32> - - %expert_route_count = scf.for %block_base = [%c0 to %assignment_count step %workgroup_size](%matched_base = %c0_i32 : i32) -> (i32) { - %assignment = index.add %block_base, %lane : index - %in_range = index.cmp ult, %assignment, %assignment_count : index - %route_expert_i32 = scf.if %in_range -> (i32) { - %token0 = index.div %assignment, %bounded_route_count : index - %route0 = index.rem %assignment, %bounded_route_count : index - %token, %route_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index - %route, %route_row_stride = index.assume %route0, %bounded_route_stride [lt(%route0, %bounded_route_stride)] : index, index - %loaded = view.load %route_view[%token, %route] : view<[%bounded_token_count]x[%bounded_route_stride]xi32> -> i32 - scf.yield %loaded : i32 - } else { - %cn1_i32 = scalar.constant -1 : i32 - scf.yield %cn1_i32 : i32 - } - %route_expert0 = index.cast %route_expert_i32 : i32 to index - %route_expert = index.assume %route_expert0 [range(%route_expert0, -1, 127)] : index - %matches = index.cmp eq, %route_expert, %expert : index - %match_i32 = scf.if %matches -> (i32) { - scf.yield %c1_i32 : i32 - } else { - scf.yield %c0_i32 : i32 - } - %block_prefix = kernel.workgroup.scan %match_i32 {direction = forward, mode = exclusive} : i32 - %block_match_count_reduced = kernel.workgroup.reduce %match_i32 : i32 - %block_match_count = kernel.subgroup.broadcast.first %block_match_count_reduced : i32 - scf.if %matches { - %match_ordinal_i32 = scalar.addi %matched_base, %block_prefix : i32 - %match_ordinal0 = index.cast %match_ordinal_i32 : i32 to index - %match_ordinal = index.assume %match_ordinal0 [range(%match_ordinal0, 0, 511)] : index - %bounded_match_ordinal, %table_token_count = index.assume %match_ordinal, %bounded_token_count [lt(%match_ordinal, %bounded_token_count)] : index, index - %assignment_i32 = index.cast %assignment : index to i32 - view.store %assignment_i32, %assignment_view[%expert, %bounded_match_ordinal] : i32, view<[%launch_expert_count]x[%bounded_token_count]xi32> - } - %next_matched_base = scalar.addi %matched_base, %block_match_count : i32 - scf.yield %next_matched_base : i32 - } - - %is_lane_zero = index.cmp eq, %lane, %c0 : index - scf.if %is_lane_zero { - view.store %expert_route_count, %count_view[%expert] : i32, view<[%launch_expert_count]xi32> - } - - // Every count store precedes its workgroup's release. The last arrival - // acquires all preceding releases before any lane loads the complete table. - %local_old_counter = scf.if %is_lane_zero -> (i32) { - %old_counter = view.atomic.rmw %c1_i32, %completion_counter_view[%c0] {ordering = acq_rel, scope = device} : i32, view<1xi32> -> i32 - scf.yield %old_counter : i32 - } else { - scf.yield %c0_i32 : i32 - } - %old_counter = kernel.workgroup.reduce %local_old_counter : i32 - %expert_count_i32 = index.cast %launch_expert_count : index to i32 - %last_expert_i32 = scalar.subi %expert_count_i32, %c1_i32 : i32 - %negative_expert_count_i32 = scalar.subi %c0_i32, %expert_count_i32 : i32 - %is_last_expert = scalar.cmpi eq, %old_counter, %last_expert_i32 : i32 - scf.if %is_last_expert { - %rounded_assignment_count = index.add %assignment_count, %c31 : index - %assignment_partition_count = index.div %rounded_assignment_count, %c32 : index - %maximum_partition_count = index.add %assignment_partition_count, %launch_expert_count : index - %partition_count_view = buffer.view %partition_table_noalias[%c0_offset] : buffer -> view<1xi32> - %partition_descriptor_view = buffer.view %partition_table_noalias[%c4_bytes] : buffer -> view<[%maximum_partition_count]xi32> - %has_expert = index.cmp ult, %lane, %launch_expert_count : index - %expert_assignment_count_i32 = scf.if %has_expert -> (i32) { - %partition_expert, %table_expert_count = index.assume %lane, %launch_expert_count [lt(%lane, %launch_expert_count)] : index, index - %loaded = view.load %count_view[%partition_expert] : view<[%launch_expert_count]xi32> -> i32 - scf.yield %loaded : i32 - } else { - scf.yield %c0_i32 : i32 - } - %expert_assignment_count0 = index.cast %expert_assignment_count_i32 : i32 to index - %expert_assignment_count = index.assume %expert_assignment_count0 [range(%expert_assignment_count0, 0, 512)] : index - %rounded_expert_assignment_count = index.add %expert_assignment_count, %c31 : index - %expert_partition_count = index.div %rounded_expert_assignment_count, %c32 : index - %expert_partition_count_i32 = index.cast %expert_partition_count : index to i32 - %expert_partition_base_i32 = kernel.workgroup.scan %expert_partition_count_i32 {direction = forward, mode = exclusive} : i32 - %partition_count_i32 = kernel.workgroup.reduce %expert_partition_count_i32 : i32 - %partition_count = index.cast %partition_count_i32 : i32 to index - %bounded_partition_count, %table_partition_capacity = index.assume %partition_count, %maximum_partition_count [lt(%partition_count, %maximum_partition_count)] : index, index - %expert_partition_base0 = index.cast %expert_partition_base_i32 : i32 to index - %expert_partition_base = index.assume %expert_partition_base0 [range(%expert_partition_base0, 0, 255)] : index - scf.if %has_expert { - %partition_expert, %table_expert_count = index.assume %lane, %launch_expert_count [lt(%lane, %launch_expert_count)] : index, index - %expert_i32 = index.cast %partition_expert : index to i32 - scf.for %partition = [%c0 to %expert_partition_count step %c1] { - %descriptor_ordinal0 = index.add %expert_partition_base, %partition : index - %descriptor_ordinal, %descriptor_count = index.assume %descriptor_ordinal0, %bounded_partition_count [lt(%descriptor_ordinal0, %bounded_partition_count)] : index, index - %table_descriptor_ordinal, %table_descriptor_capacity = index.assume %descriptor_ordinal, %maximum_partition_count [lt(%descriptor_ordinal, %maximum_partition_count)] : index, index - %partition_remainder = index.rem %expert_assignment_count, %c32 : index - %has_partial_tail = index.cmp ne, %partition_remainder, %c0 : index - %partition_row_count = scf.if %has_partial_tail -> (index) { - %next_partition = index.add %partition, %c1 : index - %is_tail_partition = index.cmp eq, %next_partition, %expert_partition_count : index - %tail_row_count = scf.if %is_tail_partition -> (index) { - scf.yield %partition_remainder : index - } else { - scf.yield %c32 : index - } - scf.yield %tail_row_count : index - } else { - scf.yield %c32 : index - } - %partition_i32 = index.cast %partition : index to i32 - %partition_row_count_i32 = index.cast %partition_row_count : index to i32 - %packed_partition = scalar.shli %partition_i32, %c7_i32 : i32 - %partition_row_count_minus_one = scalar.subi %partition_row_count_i32, %c1_i32 : i32 - %packed_row_count = scalar.shli %partition_row_count_minus_one, %c13_i32 : i32 - %packed_expert_partition = scalar.ori %expert_i32, %packed_partition : i32 - %packed_descriptor = scalar.ori %packed_expert_partition, %packed_row_count : i32 - view.store %packed_descriptor, %partition_descriptor_view[%table_descriptor_ordinal] : i32, view<[%maximum_partition_count]xi32> - } - } - scf.if %is_lane_zero { - view.store %partition_count_i32, %partition_count_view[%c0] : i32, view<1xi32> - } - // The counter cannot become reusable until every descriptor store is - // complete and visible to the following dispatch. - kernel.barrier scope(workgroup) ordering(acq_rel) - scf.if %is_lane_zero { - view.atomic.reduce %negative_expert_count_i32, %completion_counter_view[%c0] {ordering = release, scope = device} : i32, view<1xi32> - } - } - kernel.return -} +kernel.decl @qwen3_moe_build_expert_table_partition_prefill_512(%token_count$20: index, %route_count$21: index, %route_stride$22: index, %expert_count$23: index) launch(%token_count$24: index, %route_count$25: index, %route_stride$26: index, %expert_count$27: index, %route_ids: buffer, %expert_table: buffer, %partition_table: buffer, %completion_counter: buffer) // Compare both fused outputs against the production composition, then invoke // the fused route twice against one counter to make reset correctness visible. @@ -224,6 +68,8 @@ check.case public @qwen3_moe_expert_table_partition_composed_benchmark_case { check.return } +check.benchmark<@qwen3_moe_expert_table_partition_fused_differential_case> @qwen3_moe_expert_table_partition_fused_differential + check.benchmark<@qwen3_moe_expert_table_partition_fused_benchmark_case> @qwen3_moe_expert_table_partition_fused_prefill_512 check.benchmark<@qwen3_moe_expert_table_partition_composed_benchmark_case> @qwen3_moe_expert_table_partition_composed_prefill_512 diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_decode_f32_f16_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_decode_f32_f16_wmma.loom deleted file mode 100644 index fc523520717a..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_decode_f32_f16_wmma.loom +++ /dev/null @@ -1,690 +0,0 @@ -// Qwen3 MoE grouped-query decode FlashAttention. -// -// One or two four-wave workgroups compute up to 16 query heads that share one -// KV head. Each workgroup walks the exact KV length in 64-row blocks while -// carrying online-softmax max, sum, and output state in registers. With two -// output partitions, each workgroup owns one disjoint 64-channel output half; -// this duplicates QK work but exposes more parallelism without publishing -// partial softmax state. The ownership changes mirror the cooperative-matrix -// schedule used by llama.cpp's Vulkan CM1 kernel: -// -// 1. All workitems stage a scaled 16x128 F16 GQA-head tile. -// 2. Each wave computes one 16x16 QK score slice. -// 3. Scores cross LDS so each wave can normalize four complete query heads. -// 4. F16 probabilities cross LDS for four P*V WMMA steps. -// 5. Each active lane retains one four-channel F16 packet for each of its -// four query heads across subsequent KV blocks. -// -// K and V remain in the row-major llama.cpp cache layout -// [KV token][KV head][128]. Their aligned F16 fragments load directly from -// global memory; there is no expanded or repacked persistent allocation. QK -// and the online-softmax statistics remain F32, while the P*V accumulation and -// carried output match the Vulkan oracle's F16 policy. Unlike the split-K -// fallback, this kernel never publishes partial tensors or completion atomics. -amdgpu.target @qwen3_moe_attention_decode_gfx11_wave64 {subgroup_size = 64} - -config.decl @qwen3_moe.attention.query_head_count : %value: index where [range(%value, 1, 64)] - -config.decl @qwen3_moe.attention.key_value_head_count : %value: index where [range(%value, 1, 64)] - -// Number of independent 64-channel output partitions per KV head. One avoids -// redundant QK work; two exposes more workgroups for short decode contexts. -config.decl @qwen3_moe.attention.decode.output_partition_count : %value: index where [range(%value, 1, 2)] - -kernel.def target(@qwen3_moe_attention_decode_gfx11_wave64) @qwen3_moe_flash_attention_decode_f32_f16_wmma(%key_value_token_count: index) { - %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index - %output_partition_count = config.get @qwen3_moe.attention.decode.output_partition_count : index - %c1 = index.constant 1 : index - %c256 = index.constant 256 : index - %workgroup_count = index.mul %key_value_head_count, %output_partition_count : index - kernel.launch.config workgroups(%workgroup_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index -} launch(%key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %output: buffer) { - %bounded_key_value_token_count = index.assume %key_value_token_count [range(%key_value_token_count, 1, 32768)] : index - %query_head_count = config.get @qwen3_moe.attention.query_head_count : index - %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index - %output_partition_count = config.get @qwen3_moe.attention.decode.output_partition_count : index - %workgroup_x0 = kernel.workgroup.id : index - %workgroup_count0 = index.mul %key_value_head_count, %output_partition_count : index - %workgroup_x, %workgroup_count = index.assume %workgroup_x0, %workgroup_count0 [lt(%workgroup_x0, %workgroup_count0)] : index, index - %key_value_head0 = index.div %workgroup_x, %output_partition_count : index - %key_value_head = index.assume %key_value_head0 [range(%key_value_head0, 0, 63)] : index - %output_partition = index.rem %workgroup_x, %output_partition_count : index - %workitem = kernel.workitem.id : index - %subgroup0 = kernel.subgroup.id : index - %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 3)] : index - %lane = kernel.subgroup.lane.id : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c2 = index.constant 2 : index - %c3 = index.constant 3 : index - %c4 = index.constant 4 : index - %c8 = index.constant 8 : index - %c15 = index.constant 15 : index - %c16 = index.constant 16 : index - %c32 = index.constant 32 : index - %c64 = index.constant 64 : index - %c128 = index.constant 128 : index - %c256 = index.constant 256 : index - %c0_offset = index.constant 0 : offset - %query_stage_bytes = index.constant 4352 : offset - %score_stage_bytes = index.constant 6144 : offset - %probability_stage_bytes = index.constant 3072 : offset - %product_stage_bytes = index.constant 2048 : offset - %tail_key_value_stage_capacity = index.constant 8192 : offset - %c0_f32 = scalar.constant 0.0 : f32 - %negative_large = scalar.constant -1e+30 : f32 - %head_size_f32 = scalar.constant 128.0 : f32 - %attention_scale = scalar.rsqrtf %head_size_f32 : f32 - %c0_f16 = scalar.constant 0.0 : f16 - %c1_f32 = scalar.constant 1.0 : f32 - %output_zero0 = vector.constant 0.0 : vector<4xf16> - %output_zero1 = vector.constant 0.0 : vector<4xf16> - %output_zero2 = vector.constant 0.0 : vector<4xf16> - %output_zero3 = vector.constant 0.0 : vector<4xf16> - %c0_f16x8 = vector.constant 0.0 : vector<8xf16> - %c0_f32x4 = vector.constant 0.0 : vector<4xf32> - %negative_f32x4 = vector.constant -1e+30 : vector<4xf32> - %m = index.constant 16 : index - %n = index.constant 16 : index - %k = index.constant 16 : index - %query_heads_per_key_value_head0 = index.div %query_head_count, %key_value_head_count : index - %query_heads_per_key_value_head = index.assume %query_heads_per_key_value_head0 [range(%query_heads_per_key_value_head0, 1, 16)] : index - %query_head_base = index.mul %key_value_head, %query_heads_per_key_value_head : index - %key_value_width = index.mul %key_value_head_count, %c128 : index - %key_value_head_base = index.mul %key_value_head, %c128 : index - %full_key_value_block_count = index.div %bounded_key_value_token_count, %c64 : index - %full_key_value_token_count0 = index.mul %full_key_value_block_count, %c64 : index - %full_key_value_token_count = index.assume %full_key_value_token_count0 [range(%full_key_value_token_count0, 0, 32768), mul(%full_key_value_token_count0, 64)] : index - %tail_key_value_token_count = index.sub %bounded_key_value_token_count, %full_key_value_token_count : index - %has_key_value_tail = index.cmp ne, %tail_key_value_token_count, %c0 : index - %has_single_key_value_tail = index.cmp eq, %tail_key_value_token_count, %c1 : index - %tail_key_value_stage_bytes = scf.select %has_key_value_tail, %tail_key_value_stage_capacity, %c0_offset : offset - %subgroup_score_column = index.mul %subgroup, %c16 : index - %subgroup_query_row = index.mul %subgroup, %c4 : index - %query_row0 = index.add %subgroup_query_row, %c0 : index - %query_row1 = index.add %subgroup_query_row, %c1 : index - %query_row2 = index.add %subgroup_query_row, %c2 : index - %query_row3 = index.add %subgroup_query_row, %c3 : index - %query_head0 = index.add %query_head_base, %query_row0 : index - %query_head1 = index.add %query_head_base, %query_row1 : index - %query_head2 = index.add %query_head_base, %query_row2 : index - %query_head3 = index.add %query_head_base, %query_row3 : index - %query_valid0 = index.cmp ult, %query_head0, %query_head_count : index - %query_valid1 = index.cmp ult, %query_head1, %query_head_count : index - %query_valid2 = index.cmp ult, %query_head2, %query_head_count : index - %query_valid3 = index.cmp ult, %query_head3, %query_head_count : index - %query_valid = vector.from_elements %query_valid0, %query_valid1, %query_valid2, %query_valid3 : vector<4xi1> - %subgroup_product_channel = index.mul %subgroup, %c16 : index - %lane_output_tile = index.div %lane, %c16 : index - %output_tile_count = index.div %c2, %output_partition_count : index - %output_tile_end = index.add %output_partition, %output_tile_count : index - %lane_product_channel0 = index.rem %lane, %c16 : index - %lane_product_channel = index.mul %lane_product_channel0, %c4 : index - %lane_output_channel = index.mul %lane, %c4 : index - %lane_output_at_or_after_partition = index.cmp uge, %lane_output_tile, %output_partition : index - %lane_output_before_end = index.cmp ult, %lane_output_tile, %output_tile_end : index - %lane_has_output = scalar.andi %lane_output_at_or_after_partition, %lane_output_before_end : i1 - // The padded LDS rows mirror the Vulkan oracle's ownership changes. Q uses - // eight spare F16 columns after its 128 channels. Score and probability - // transpose to key-major rows with eight spare columns after 16 queries. - // These strides avoid the bank pattern produced by dense transposed rows. - %query_transposed_layout = encoding.layout.strided [1, 136] : encoding - %probability_transposed_layout = encoding.layout.strided [1, 24] : encoding - %query_noalias, %key_noalias, %value_noalias, %mask_noalias, %output_noalias = buffer.assume.noalias %query, %key, %value, %mask, %output : buffer, buffer, buffer, buffer, buffer - %query_aligned = buffer.assume.alignment %query_noalias {minimum_alignment = 16} : buffer - %key_aligned = buffer.assume.alignment %key_noalias {minimum_alignment = 16} : buffer - %value_aligned = buffer.assume.alignment %value_noalias {minimum_alignment = 16} : buffer - %mask_aligned = buffer.assume.alignment %mask_noalias {minimum_alignment = 16} : buffer - %output_aligned = buffer.assume.alignment %output_noalias {minimum_alignment = 16} : buffer - %query_view = buffer.view %query_aligned[%c0_offset] : buffer -> view<[%query_head_count]x128xf32> - %key_view = buffer.view %key_aligned[%c0_offset] : buffer -> view<[%bounded_key_value_token_count]x[%key_value_width]xf16> - %value_view = buffer.view %value_aligned[%c0_offset] : buffer -> view<[%bounded_key_value_token_count]x[%key_value_width]xf16> - %mask_view = buffer.view %mask_aligned[%c0_offset] : buffer -> view<[%bounded_key_value_token_count]xf16> - %output_view = buffer.view %output_aligned[%c0_offset] : buffer -> view<[%query_head_count]x128xf32> - %query_stage = buffer.alloca align(16) %query_stage_bytes : buffer - %score_stage = buffer.alloca align(16) %score_stage_bytes : buffer - %probability_stage = buffer.alloca align(16) %probability_stage_bytes : buffer - %product_stage = buffer.alloca align(16) %product_stage_bytes : buffer - %tail_key_value_stage = buffer.alloca align(16) %tail_key_value_stage_bytes : buffer - %query_stage_view = buffer.view %query_stage[%c0_offset] : buffer -> view<16x136xf16> - %query_transposed_view = buffer.view %query_stage[%c0_offset] : buffer -> view<128x16xf16, %query_transposed_layout> - %score_stage_view = buffer.view %score_stage[%c0_offset] : buffer -> view<64x24xf32> - %probability_stage_view = buffer.view %probability_stage[%c0_offset] : buffer -> view<16x64xf16, %probability_transposed_layout> - %product_stage_view = buffer.view %product_stage[%c0_offset] : buffer -> view<16x64xf16> - %tail_key_value_stage_view = buffer.view %tail_key_value_stage[%c0_offset] : buffer -> view<32x128xf16> - // Scale and truncate Q exactly once. The Vulkan reference does this before - // entering its KV loop, making QK a native F16 WMMA while retaining F32 - // accumulation. - scf.for %load_iteration = [%c0 to %c8 step %c1] unroll { - %linear = index.madd %load_iteration, %c256, %workitem : index - %local_query_row = index.div %linear, %c128 : index - %query_channel = index.rem %linear, %c128 : index - %local_query_head = index.add %query_head_base, %local_query_row : index - %query_valid_load = index.cmp ult, %local_query_head, %query_head_count : index - %query_value = scf.if %query_valid_load -> (f16) { - %loaded = view.load %query_view[%local_query_head, %query_channel] : view<[%query_head_count]x128xf32> -> f32 - %scaled = scalar.mulf %loaded, %attention_scale : f32 - %truncated = scalar.fptrunc %scaled : f32 to f16 - scf.yield %truncated : f16 - } else { - scf.yield %c0_f16 : f16 - } - view.store %query_value, %query_stage_view[%local_query_row, %query_channel] : f16, view<16x136xf16> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %full_max, %full_sum, %full_output0, %full_output1, %full_output2, %full_output3 = scf.for %key_origin = [%c0 to %full_key_value_token_count step %c64](%current_max = %negative_f32x4 : vector<4xf32>, %current_sum = %c0_f32x4 : vector<4xf32>, %current_output0 = %output_zero0 : vector<4xf16>, %current_output1 = %output_zero1 : vector<4xf16>, %current_output2 = %output_zero2 : vector<4xf16>, %current_output3 = %output_zero3 : vector<4xf16>) -> (vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { - // Four independent wave-level WMMAs produce a 16x64 score tile. - %score_key_origin0 = index.add %key_origin, %subgroup_score_column : index - %last_full_key_tile_start = index.sub %bounded_key_value_token_count, %c15 : index - %score_key_origin = index.assume %score_key_origin0 [lt(%score_key_origin0, %last_full_key_tile_start)] : index - %score_init_values = vector.constant 0.0 : vector<4xf32> - %score_init = vector.fragment %score_init_values shape [%m, %n] : vector<4xf32> - %score_fragment = scf.for %head_tile = [%c0 to %c128 step %c16](%score_accumulator = %score_init : vector<4xf32>) -> (vector<4xf32>) unroll { - %key_channel = index.add %key_value_head_base, %head_tile : index - %key_fragment = vector.fragment.load %key_view[%score_key_origin, %key_channel] shape [%m, %k] : view<[%bounded_key_value_token_count]x[%key_value_width]xf16> -> vector<16xf16> - %query_fragment = vector.fragment.load %query_transposed_view[%head_tile, %c0] shape [%k, %n] : view<128x16xf16, %query_transposed_layout> -> vector<16xf16> - %next_score_accumulator = vector.mma %key_fragment, %query_fragment, %score_accumulator : vector<16xf16>, vector<16xf16>, vector<4xf32> - scf.yield %next_score_accumulator : vector<4xf32> - } - vector.fragment.store %score_fragment, %score_stage_view[%subgroup_score_column, %c0] shape [%m, %n] : vector<4xf32>, view<64x24xf32> - kernel.barrier scope(workgroup) ordering(acq_rel) - // LDS transposes ownership from one 16-column score slice per wave to - // four complete query rows per wave. Every lane contributes one key - // column to each of those rows. - %key_token0 = index.add %key_origin, %lane : index - %key_token = index.assume %key_token0 [lt(%key_token0, %bounded_key_value_token_count)] : index - %raw_score0 = view.load %score_stage_view[%lane, %query_row0] : view<64x24xf32> -> f32 - %raw_score1 = view.load %score_stage_view[%lane, %query_row1] : view<64x24xf32> -> f32 - %raw_score2 = view.load %score_stage_view[%lane, %query_row2] : view<64x24xf32> -> f32 - %raw_score3 = view.load %score_stage_view[%lane, %query_row3] : view<64x24xf32> -> f32 - %mask_f16 = view.load %mask_view[%key_token] : view<[%bounded_key_value_token_count]xf16> -> f16 - %mask_f32 = scalar.extf %mask_f16 : f16 to f32 - %masked_score0 = scf.if %query_valid0 -> (f32) { - %score = scalar.addf %raw_score0, %mask_f32 : f32 - scf.yield %score : f32 - } else { - scf.yield %negative_large : f32 - } - %masked_score1 = scf.if %query_valid1 -> (f32) { - %score = scalar.addf %raw_score1, %mask_f32 : f32 - scf.yield %score : f32 - } else { - scf.yield %negative_large : f32 - } - %masked_score2 = scf.if %query_valid2 -> (f32) { - %score = scalar.addf %raw_score2, %mask_f32 : f32 - scf.yield %score : f32 - } else { - scf.yield %negative_large : f32 - } - %masked_score3 = scf.if %query_valid3 -> (f32) { - %score = scalar.addf %raw_score3, %mask_f32 : f32 - scf.yield %score : f32 - } else { - scf.yield %negative_large : f32 - } - %masked_scores = vector.from_elements %masked_score0, %masked_score1, %masked_score2, %masked_score3 : vector<4xf32> - %block_max = kernel.subgroup.reduce %masked_scores : vector<4xf32> - %next_max = vector.maxnumf %current_max, %block_max : vector<4xf32> - %score_delta = vector.subf %masked_scores, %next_max : vector<4xf32> - %raw_probability = vector.expf %score_delta : vector<4xf32> - %probability = vector.select %query_valid, %raw_probability, %c0_f32x4 : vector<4xf32> - %block_sum = kernel.subgroup.reduce %probability : vector<4xf32> - %old_delta = vector.subf %current_max, %next_max : vector<4xf32> - %old_scale = vector.expf %old_delta : vector<4xf32> - %scaled_current_sum = vector.mulf %current_sum, %old_scale : vector<4xf32> - %next_sum = vector.addf %scaled_current_sum, %block_sum : vector<4xf32> - %probability_f16 = vector.fptrunc %probability : vector<4xf32> to vector<4xf16> - %probability0 = vector.extract %probability_f16[0] : vector<4xf16> -> f16 - %probability1 = vector.extract %probability_f16[1] : vector<4xf16> -> f16 - %probability2 = vector.extract %probability_f16[2] : vector<4xf16> -> f16 - %probability3 = vector.extract %probability_f16[3] : vector<4xf16> -> f16 - view.store %probability0, %probability_stage_view[%query_row0, %lane] : f16, view<16x64xf16, %probability_transposed_layout> - view.store %probability1, %probability_stage_view[%query_row1, %lane] : f16, view<16x64xf16, %probability_transposed_layout> - view.store %probability2, %probability_stage_view[%query_row2, %lane] : f16, view<16x64xf16, %probability_transposed_layout> - view.store %probability3, %probability_stage_view[%query_row3, %lane] : f16, view<16x64xf16, %probability_transposed_layout> - kernel.barrier scope(workgroup) ordering(acq_rel) - // Match the Vulkan CM1 ownership schedule by computing two sequential - // 64-channel output tiles. This keeps one P*V accumulator live per wave - // and reuses a 2 KiB exchange tile instead of retaining both halves. - %old_scale_f16 = vector.fptrunc %old_scale : vector<4xf32> to vector<4xf16> - %old_scale0_scalar = vector.extract %old_scale_f16[0] : vector<4xf16> -> f16 - %old_scale1_scalar = vector.extract %old_scale_f16[1] : vector<4xf16> -> f16 - %old_scale2_scalar = vector.extract %old_scale_f16[2] : vector<4xf16> -> f16 - %old_scale3_scalar = vector.extract %old_scale_f16[3] : vector<4xf16> -> f16 - %old_scale0 = vector.splat %old_scale0_scalar : vector<4xf16> - %old_scale1 = vector.splat %old_scale1_scalar : vector<4xf16> - %old_scale2 = vector.splat %old_scale2_scalar : vector<4xf16> - %old_scale3 = vector.splat %old_scale3_scalar : vector<4xf16> - %scaled_current_output0 = vector.mulf %current_output0, %old_scale0 : vector<4xf16> - %scaled_current_output1 = vector.mulf %current_output1, %old_scale1 : vector<4xf16> - %scaled_current_output2 = vector.mulf %current_output2, %old_scale2 : vector<4xf16> - %scaled_current_output3 = vector.mulf %current_output3, %old_scale3 : vector<4xf16> - %next_output0, %next_output1, %next_output2, %next_output3 = scf.for %output_tile = [%c0 to %c2 step %c1](%tile_output0 = %scaled_current_output0 : vector<4xf16>, %tile_output1 = %scaled_current_output1 : vector<4xf16>, %tile_output2 = %scaled_current_output2 : vector<4xf16>, %tile_output3 = %scaled_current_output3 : vector<4xf16>) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) unroll { - %tile_at_or_after_partition = index.cmp uge, %output_tile, %output_partition : index - %tile_before_partition_end = index.cmp ult, %output_tile, %output_tile_end : index - %tile_selected = scalar.andi %tile_at_or_after_partition, %tile_before_partition_end : i1 - %updated_output0, %updated_output1, %updated_output2, %updated_output3 = scf.if %tile_selected -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { - %output_tile_channel = index.mul %output_tile, %c64 : index - %value_channel0 = index.add %key_value_head_base, %output_tile_channel : index - %value_channel = index.add %value_channel0, %subgroup_product_channel : index - %product_init = vector.fragment %c0_f16x8 shape [%m, %n] : vector<8xf16> - %product_fragment = scf.for %key_tile = [%c0 to %c64 step %c16](%product_accumulator = %product_init : vector<8xf16>) -> (vector<8xf16>) unroll { - %value_token0 = index.add %key_origin, %key_tile : index - %value_token = index.assume %value_token0 [lt(%value_token0, %last_full_key_tile_start)] : index - %probability_fragment = vector.fragment.load %probability_stage_view[%c0, %key_tile] shape [%m, %k] : view<16x64xf16, %probability_transposed_layout> -> vector<16xf16> - %value_fragment = vector.fragment.load %value_view[%value_token, %value_channel] shape [%k, %n] : view<[%bounded_key_value_token_count]x[%key_value_width]xf16> -> vector<16xf16> - %next_product_accumulator = vector.mma %probability_fragment, %value_fragment, %product_accumulator : vector<16xf16>, vector<16xf16>, vector<8xf16> - scf.yield %next_product_accumulator : vector<8xf16> - } - vector.fragment.store %product_fragment, %product_stage_view[%c0, %subgroup_product_channel] shape [%m, %n] : vector<8xf16>, view<16x64xf16> - kernel.barrier scope(workgroup) ordering(acq_rel) - %owns_output_tile = index.cmp eq, %lane_output_tile, %output_tile : index - %next_tile_output0, %next_tile_output1, %next_tile_output2, %next_tile_output3 = scf.if %owns_output_tile -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { - %block_output0 = vector.load %product_stage_view[%query_row0, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> - %block_output1 = vector.load %product_stage_view[%query_row1, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> - %block_output2 = vector.load %product_stage_view[%query_row2, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> - %block_output3 = vector.load %product_stage_view[%query_row3, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> - %next_tile_output0 = vector.addf %tile_output0, %block_output0 : vector<4xf16> - %next_tile_output1 = vector.addf %tile_output1, %block_output1 : vector<4xf16> - %next_tile_output2 = vector.addf %tile_output2, %block_output2 : vector<4xf16> - %next_tile_output3 = vector.addf %tile_output3, %block_output3 : vector<4xf16> - scf.yield %next_tile_output0, %next_tile_output1, %next_tile_output2, %next_tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } else { - scf.yield %tile_output0, %tile_output1, %tile_output2, %tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } - // Complete every read before another output tile overwrites LDS. - kernel.barrier scope(workgroup) ordering(acq_rel) - scf.yield %next_tile_output0, %next_tile_output1, %next_tile_output2, %next_tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } else { - scf.yield %tile_output0, %tile_output1, %tile_output2, %tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } - scf.yield %updated_output0, %updated_output1, %updated_output2, %updated_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } - scf.yield %next_max, %next_sum, %next_output0, %next_output1, %next_output2, %next_output3 : vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } - // A single trailing KV row is cheaper as a native wave reduction and direct - // V update. Skip the general WMMA tail loop for that exact JIT-specialized - // shape; larger tails continue through the masked 32-row schedule below. - %wmma_tail_start = scf.select %has_single_key_value_tail, %bounded_key_value_token_count, %full_key_value_token_count : index - %tail_score_wave = index.cmp ult, %subgroup, %c2 : index - %wmma_tail_max, %wmma_tail_sum, %wmma_tail_output0, %wmma_tail_output1, %wmma_tail_output2, %wmma_tail_output3 = scf.for %tail_key_origin = [%wmma_tail_start to %bounded_key_value_token_count step %c32](%current_max = %full_max : vector<4xf32>, %current_sum = %full_sum : vector<4xf32>, %current_output0 = %full_output0 : vector<4xf16>, %current_output1 = %full_output1 : vector<4xf16>, %current_output2 = %full_output2 : vector<4xf16>, %current_output3 = %full_output3 : vector<4xf16>) -> (vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { - %tail_remaining = index.sub %bounded_key_value_token_count, %tail_key_origin : index - %tail_key_count = index.min %tail_remaining, %c32 : index - // Cooperatively stage one K tile, explicitly zeroing the padded rows. - scf.for %load_iteration = [%c0 to %c16 step %c1] unroll { - %linear = index.madd %load_iteration, %c256, %workitem : index - %tail_key_row = index.div %linear, %c128 : index - %tail_key_channel = index.rem %linear, %c128 : index - %tail_key_valid = index.cmp ult, %tail_key_row, %tail_key_count : index - %tail_key_value = scf.if %tail_key_valid -> (f16) { - %tail_key_token0 = index.add %tail_key_origin, %tail_key_row : index - %tail_key_token = index.assume %tail_key_token0 [lt(%tail_key_token0, %bounded_key_value_token_count)] : index - %global_key_channel = index.add %key_value_head_base, %tail_key_channel : index - %loaded = view.load %key_view[%tail_key_token, %global_key_channel] : view<[%bounded_key_value_token_count]x[%key_value_width]xf16> -> f16 - scf.yield %loaded : f16 - } else { - scf.yield %c0_f16 : f16 - } - view.store %tail_key_value, %tail_key_value_stage_view[%tail_key_row, %tail_key_channel] : f16, view<32x128xf16> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - // Waves zero and one compute the two 16-column QK fragments in this tile. - scf.if %tail_score_wave { - %tail_score_subgroup = index.assume %subgroup [range(%subgroup, 0, 1)] : index - %tail_score_column = index.mul %tail_score_subgroup, %c16 : index - %tail_score_init_values = vector.constant 0.0 : vector<4xf32> - %tail_score_init = vector.fragment %tail_score_init_values shape [%m, %n] : vector<4xf32> - %tail_score_fragment = scf.for %head_tile = [%c0 to %c128 step %c16](%score_accumulator = %tail_score_init : vector<4xf32>) -> (vector<4xf32>) unroll { - %key_fragment = vector.fragment.load %tail_key_value_stage_view[%tail_score_column, %head_tile] shape [%m, %k] : view<32x128xf16> -> vector<16xf16> - %query_fragment = vector.fragment.load %query_transposed_view[%head_tile, %c0] shape [%k, %n] : view<128x16xf16, %query_transposed_layout> -> vector<16xf16> - %next_score_accumulator = vector.mma %key_fragment, %query_fragment, %score_accumulator : vector<16xf16>, vector<16xf16>, vector<4xf32> - scf.yield %next_score_accumulator : vector<4xf32> - } - vector.fragment.store %tail_score_fragment, %score_stage_view[%tail_score_column, %c0] shape [%m, %n] : vector<4xf32>, view<64x24xf32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - // LDS changes ownership from the QK wave to four query-row waves. Lanes - // beyond the logical tail never read the score or mask buffers. - %tail_lane_valid = index.cmp ult, %lane, %tail_key_count : index - %tail_valid0 = scalar.andi %tail_lane_valid, %query_valid0 : i1 - %tail_valid1 = scalar.andi %tail_lane_valid, %query_valid1 : i1 - %tail_valid2 = scalar.andi %tail_lane_valid, %query_valid2 : i1 - %tail_valid3 = scalar.andi %tail_lane_valid, %query_valid3 : i1 - %tail_valid = vector.from_elements %tail_valid0, %tail_valid1, %tail_valid2, %tail_valid3 : vector<4xi1> - %tail_key_token0 = index.add %tail_key_origin, %lane : index - %tail_mask_f32 = scf.if %tail_lane_valid -> (f32) { - %tail_key_token = index.assume %tail_key_token0 [lt(%tail_key_token0, %bounded_key_value_token_count)] : index - %mask_f16 = view.load %mask_view[%tail_key_token] : view<[%bounded_key_value_token_count]xf16> -> f16 - %mask_f32 = scalar.extf %mask_f16 : f16 to f32 - scf.yield %mask_f32 : f32 - } else { - scf.yield %c0_f32 : f32 - } - %masked_score0 = scf.if %tail_valid0 -> (f32) { - %raw_score = view.load %score_stage_view[%lane, %query_row0] : view<64x24xf32> -> f32 - %score = scalar.addf %raw_score, %tail_mask_f32 : f32 - scf.yield %score : f32 - } else { - scf.yield %negative_large : f32 - } - %masked_score1 = scf.if %tail_valid1 -> (f32) { - %raw_score = view.load %score_stage_view[%lane, %query_row1] : view<64x24xf32> -> f32 - %score = scalar.addf %raw_score, %tail_mask_f32 : f32 - scf.yield %score : f32 - } else { - scf.yield %negative_large : f32 - } - %masked_score2 = scf.if %tail_valid2 -> (f32) { - %raw_score = view.load %score_stage_view[%lane, %query_row2] : view<64x24xf32> -> f32 - %score = scalar.addf %raw_score, %tail_mask_f32 : f32 - scf.yield %score : f32 - } else { - scf.yield %negative_large : f32 - } - %masked_score3 = scf.if %tail_valid3 -> (f32) { - %raw_score = view.load %score_stage_view[%lane, %query_row3] : view<64x24xf32> -> f32 - %score = scalar.addf %raw_score, %tail_mask_f32 : f32 - scf.yield %score : f32 - } else { - scf.yield %negative_large : f32 - } - %masked_scores = vector.from_elements %masked_score0, %masked_score1, %masked_score2, %masked_score3 : vector<4xf32> - %block_max = kernel.subgroup.reduce %masked_scores : vector<4xf32> - %next_max = vector.maxnumf %current_max, %block_max : vector<4xf32> - %score_delta = vector.subf %masked_scores, %next_max : vector<4xf32> - %raw_probability = vector.expf %score_delta : vector<4xf32> - %probability = vector.select %tail_valid, %raw_probability, %c0_f32x4 : vector<4xf32> - %block_sum = kernel.subgroup.reduce %probability : vector<4xf32> - %old_delta = vector.subf %current_max, %next_max : vector<4xf32> - %old_scale = vector.expf %old_delta : vector<4xf32> - %scaled_current_sum = vector.mulf %current_sum, %old_scale : vector<4xf32> - %next_sum = vector.addf %scaled_current_sum, %block_sum : vector<4xf32> - %probability_f16 = vector.fptrunc %probability : vector<4xf32> to vector<4xf16> - %probability0 = vector.extract %probability_f16[0] : vector<4xf16> -> f16 - %probability1 = vector.extract %probability_f16[1] : vector<4xf16> -> f16 - %probability2 = vector.extract %probability_f16[2] : vector<4xf16> -> f16 - %probability3 = vector.extract %probability_f16[3] : vector<4xf16> -> f16 - view.store %probability0, %probability_stage_view[%query_row0, %lane] : f16, view<16x64xf16, %probability_transposed_layout> - view.store %probability1, %probability_stage_view[%query_row1, %lane] : f16, view<16x64xf16, %probability_transposed_layout> - view.store %probability2, %probability_stage_view[%query_row2, %lane] : f16, view<16x64xf16, %probability_transposed_layout> - view.store %probability3, %probability_stage_view[%query_row3, %lane] : f16, view<16x64xf16, %probability_transposed_layout> - kernel.barrier scope(workgroup) ordering(acq_rel) - // Reuse product scratch for V after every probability is resident in its - // disjoint LDS tile. - scf.for %load_iteration = [%c0 to %c16 step %c1] unroll { - %linear = index.madd %load_iteration, %c256, %workitem : index - %tail_value_row = index.div %linear, %c128 : index - %tail_value_channel = index.rem %linear, %c128 : index - %tail_value_valid = index.cmp ult, %tail_value_row, %tail_key_count : index - %tail_value = scf.if %tail_value_valid -> (f16) { - %tail_value_token0 = index.add %tail_key_origin, %tail_value_row : index - %tail_value_token = index.assume %tail_value_token0 [lt(%tail_value_token0, %bounded_key_value_token_count)] : index - %global_value_channel = index.add %key_value_head_base, %tail_value_channel : index - %loaded = view.load %value_view[%tail_value_token, %global_value_channel] : view<[%bounded_key_value_token_count]x[%key_value_width]xf16> -> f16 - scf.yield %loaded : f16 - } else { - scf.yield %c0_f16 : f16 - } - view.store %tail_value, %tail_key_value_stage_view[%tail_value_row, %tail_value_channel] : f16, view<32x128xf16> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - // Keep tail V staging separate from the 2 KiB product exchange, then use - // the same two 64-channel phases as the aligned path. - %old_scale_f16 = vector.fptrunc %old_scale : vector<4xf32> to vector<4xf16> - %old_scale0_scalar = vector.extract %old_scale_f16[0] : vector<4xf16> -> f16 - %old_scale1_scalar = vector.extract %old_scale_f16[1] : vector<4xf16> -> f16 - %old_scale2_scalar = vector.extract %old_scale_f16[2] : vector<4xf16> -> f16 - %old_scale3_scalar = vector.extract %old_scale_f16[3] : vector<4xf16> -> f16 - %old_scale0 = vector.splat %old_scale0_scalar : vector<4xf16> - %old_scale1 = vector.splat %old_scale1_scalar : vector<4xf16> - %old_scale2 = vector.splat %old_scale2_scalar : vector<4xf16> - %old_scale3 = vector.splat %old_scale3_scalar : vector<4xf16> - %scaled_current_output0 = vector.mulf %current_output0, %old_scale0 : vector<4xf16> - %scaled_current_output1 = vector.mulf %current_output1, %old_scale1 : vector<4xf16> - %scaled_current_output2 = vector.mulf %current_output2, %old_scale2 : vector<4xf16> - %scaled_current_output3 = vector.mulf %current_output3, %old_scale3 : vector<4xf16> - %next_output0, %next_output1, %next_output2, %next_output3 = scf.for %output_tile = [%c0 to %c2 step %c1](%tile_output0 = %scaled_current_output0 : vector<4xf16>, %tile_output1 = %scaled_current_output1 : vector<4xf16>, %tile_output2 = %scaled_current_output2 : vector<4xf16>, %tile_output3 = %scaled_current_output3 : vector<4xf16>) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) unroll { - %tile_at_or_after_partition = index.cmp uge, %output_tile, %output_partition : index - %tile_before_partition_end = index.cmp ult, %output_tile, %output_tile_end : index - %tile_selected = scalar.andi %tile_at_or_after_partition, %tile_before_partition_end : i1 - %updated_output0, %updated_output1, %updated_output2, %updated_output3 = scf.if %tile_selected -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { - %output_tile_channel = index.mul %output_tile, %c64 : index - %value_channel = index.add %output_tile_channel, %subgroup_product_channel : index - %tail_product_init = vector.fragment %c0_f16x8 shape [%m, %n] : vector<8xf16> - %tail_product_fragment = scf.for %key_tile = [%c0 to %c32 step %c16](%product_accumulator = %tail_product_init : vector<8xf16>) -> (vector<8xf16>) unroll { - %probability_fragment = vector.fragment.load %probability_stage_view[%c0, %key_tile] shape [%m, %k] : view<16x64xf16, %probability_transposed_layout> -> vector<16xf16> - %value_fragment = vector.fragment.load %tail_key_value_stage_view[%key_tile, %value_channel] shape [%k, %n] : view<32x128xf16> -> vector<16xf16> - %next_product_accumulator = vector.mma %probability_fragment, %value_fragment, %product_accumulator : vector<16xf16>, vector<16xf16>, vector<8xf16> - scf.yield %next_product_accumulator : vector<8xf16> - } - vector.fragment.store %tail_product_fragment, %product_stage_view[%c0, %subgroup_product_channel] shape [%m, %n] : vector<8xf16>, view<16x64xf16> - kernel.barrier scope(workgroup) ordering(acq_rel) - %owns_output_tile = index.cmp eq, %lane_output_tile, %output_tile : index - %next_tile_output0, %next_tile_output1, %next_tile_output2, %next_tile_output3 = scf.if %owns_output_tile -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { - %block_output0 = vector.load %product_stage_view[%query_row0, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> - %block_output1 = vector.load %product_stage_view[%query_row1, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> - %block_output2 = vector.load %product_stage_view[%query_row2, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> - %block_output3 = vector.load %product_stage_view[%query_row3, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> - %next_tile_output0 = vector.addf %tile_output0, %block_output0 : vector<4xf16> - %next_tile_output1 = vector.addf %tile_output1, %block_output1 : vector<4xf16> - %next_tile_output2 = vector.addf %tile_output2, %block_output2 : vector<4xf16> - %next_tile_output3 = vector.addf %tile_output3, %block_output3 : vector<4xf16> - scf.yield %next_tile_output0, %next_tile_output1, %next_tile_output2, %next_tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } else { - scf.yield %tile_output0, %tile_output1, %tile_output2, %tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - scf.yield %next_tile_output0, %next_tile_output1, %next_tile_output2, %next_tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } else { - scf.yield %tile_output0, %tile_output1, %tile_output2, %tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } - scf.yield %updated_output0, %updated_output1, %updated_output2, %updated_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } - scf.yield %next_max, %next_sum, %next_output0, %next_output1, %next_output2, %next_output3 : vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } - %final_max, %final_sum, %final_output0, %final_output1, %final_output2, %final_output3 = scf.if %has_single_key_value_tail -> (vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { - %single_tail_token = index.assume %full_key_value_token_count [lt(%full_key_value_token_count, %bounded_key_value_token_count)] : index - %single_tail_channel = index.mul %lane, %c2 : index - %single_tail_key_channel = index.add %key_value_head_base, %single_tail_channel : index - %key_pair_f16 = vector.load %key_view[%single_tail_token, %single_tail_key_channel] : view<[%bounded_key_value_token_count]x[%key_value_width]xf16> -> vector<2xf16> - %query_pair0_f16 = vector.load %query_stage_view[%query_row0, %single_tail_channel] : view<16x136xf16> -> vector<2xf16> - %query_pair1_f16 = vector.load %query_stage_view[%query_row1, %single_tail_channel] : view<16x136xf16> -> vector<2xf16> - %query_pair2_f16 = vector.load %query_stage_view[%query_row2, %single_tail_channel] : view<16x136xf16> -> vector<2xf16> - %query_pair3_f16 = vector.load %query_stage_view[%query_row3, %single_tail_channel] : view<16x136xf16> -> vector<2xf16> - %key_pair = vector.extf %key_pair_f16 : vector<2xf16> to vector<2xf32> - %query_pair0 = vector.extf %query_pair0_f16 : vector<2xf16> to vector<2xf32> - %query_pair1 = vector.extf %query_pair1_f16 : vector<2xf16> to vector<2xf32> - %query_pair2 = vector.extf %query_pair2_f16 : vector<2xf16> to vector<2xf32> - %query_pair3 = vector.extf %query_pair3_f16 : vector<2xf16> to vector<2xf32> - %product_pair0 = vector.mulf %query_pair0, %key_pair : vector<2xf32> - %product_pair1 = vector.mulf %query_pair1, %key_pair : vector<2xf32> - %product_pair2 = vector.mulf %query_pair2, %key_pair : vector<2xf32> - %product_pair3 = vector.mulf %query_pair3, %key_pair : vector<2xf32> - %partial_score0 = vector.reduce %product_pair0, %c0_f32 : vector<2xf32>, f32 - %partial_score1 = vector.reduce %product_pair1, %c0_f32 : vector<2xf32>, f32 - %partial_score2 = vector.reduce %product_pair2, %c0_f32 : vector<2xf32>, f32 - %partial_score3 = vector.reduce %product_pair3, %c0_f32 : vector<2xf32>, f32 - %partial_scores = vector.from_elements %partial_score0, %partial_score1, %partial_score2, %partial_score3 : vector<4xf32> - %reduced_scores = kernel.subgroup.reduce %partial_scores : vector<4xf32> - %mask_f16 = view.load %mask_view[%single_tail_token] : view<[%bounded_key_value_token_count]xf16> -> f16 - %mask_f32 = scalar.extf %mask_f16 : f16 to f32 - %mask_vector = vector.splat %mask_f32 : vector<4xf32> - %raw_scores = vector.addf %reduced_scores, %mask_vector : vector<4xf32> - %masked_scores = vector.select %query_valid, %raw_scores, %negative_f32x4 : vector<4xf32> - %next_max = vector.maxnumf %wmma_tail_max, %masked_scores : vector<4xf32> - %score_delta = vector.subf %masked_scores, %next_max : vector<4xf32> - %raw_probability = vector.expf %score_delta : vector<4xf32> - %probability = vector.select %query_valid, %raw_probability, %c0_f32x4 : vector<4xf32> - %old_delta = vector.subf %wmma_tail_max, %next_max : vector<4xf32> - %old_scale = vector.expf %old_delta : vector<4xf32> - %scaled_current_sum = vector.mulf %wmma_tail_sum, %old_scale : vector<4xf32> - %next_sum = vector.addf %scaled_current_sum, %probability : vector<4xf32> - %next_output0, %next_output1, %next_output2, %next_output3 = scf.if %lane_has_output -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { - %output_packet_channel = index.assume %lane_output_channel [range(%lane_output_channel, 0, 124), mul(%lane_output_channel, 4)] : index - %value_channel = index.add %key_value_head_base, %output_packet_channel : index - %value_packet = vector.load %value_view[%single_tail_token, %value_channel] : view<[%bounded_key_value_token_count]x[%key_value_width]xf16> -> vector<4xf16> - %old_scale_f16 = vector.fptrunc %old_scale : vector<4xf32> to vector<4xf16> - %probability_f16 = vector.fptrunc %probability : vector<4xf32> to vector<4xf16> - %old_scale0_scalar = vector.extract %old_scale_f16[0] : vector<4xf16> -> f16 - %old_scale1_scalar = vector.extract %old_scale_f16[1] : vector<4xf16> -> f16 - %old_scale2_scalar = vector.extract %old_scale_f16[2] : vector<4xf16> -> f16 - %old_scale3_scalar = vector.extract %old_scale_f16[3] : vector<4xf16> -> f16 - %probability0_scalar = vector.extract %probability_f16[0] : vector<4xf16> -> f16 - %probability1_scalar = vector.extract %probability_f16[1] : vector<4xf16> -> f16 - %probability2_scalar = vector.extract %probability_f16[2] : vector<4xf16> -> f16 - %probability3_scalar = vector.extract %probability_f16[3] : vector<4xf16> -> f16 - %old_scale0 = vector.splat %old_scale0_scalar : vector<4xf16> - %old_scale1 = vector.splat %old_scale1_scalar : vector<4xf16> - %old_scale2 = vector.splat %old_scale2_scalar : vector<4xf16> - %old_scale3 = vector.splat %old_scale3_scalar : vector<4xf16> - %probability0 = vector.splat %probability0_scalar : vector<4xf16> - %probability1 = vector.splat %probability1_scalar : vector<4xf16> - %probability2 = vector.splat %probability2_scalar : vector<4xf16> - %probability3 = vector.splat %probability3_scalar : vector<4xf16> - %scaled_current_output0 = vector.mulf %wmma_tail_output0, %old_scale0 : vector<4xf16> - %scaled_current_output1 = vector.mulf %wmma_tail_output1, %old_scale1 : vector<4xf16> - %scaled_current_output2 = vector.mulf %wmma_tail_output2, %old_scale2 : vector<4xf16> - %scaled_current_output3 = vector.mulf %wmma_tail_output3, %old_scale3 : vector<4xf16> - %tail_output0 = vector.mulf %value_packet, %probability0 : vector<4xf16> - %tail_output1 = vector.mulf %value_packet, %probability1 : vector<4xf16> - %tail_output2 = vector.mulf %value_packet, %probability2 : vector<4xf16> - %tail_output3 = vector.mulf %value_packet, %probability3 : vector<4xf16> - %updated_output0 = vector.addf %scaled_current_output0, %tail_output0 : vector<4xf16> - %updated_output1 = vector.addf %scaled_current_output1, %tail_output1 : vector<4xf16> - %updated_output2 = vector.addf %scaled_current_output2, %tail_output2 : vector<4xf16> - %updated_output3 = vector.addf %scaled_current_output3, %tail_output3 : vector<4xf16> - scf.yield %updated_output0, %updated_output1, %updated_output2, %updated_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } else { - scf.yield %wmma_tail_output0, %wmma_tail_output1, %wmma_tail_output2, %wmma_tail_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } - scf.yield %next_max, %next_sum, %next_output0, %next_output1, %next_output2, %next_output3 : vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } else { - scf.yield %wmma_tail_max, %wmma_tail_sum, %wmma_tail_output0, %wmma_tail_output1, %wmma_tail_output2, %wmma_tail_output3 : vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } - // Normalize and publish the lane-owned packets. Query-tail rows never - // participate in the output store. - scf.if %lane_has_output { - %output_packet_channel = index.assume %lane_output_channel [range(%lane_output_channel, 0, 124), mul(%lane_output_channel, 4)] : index - %sum0_scalar = vector.extract %final_sum[0] : vector<4xf32> -> f32 - %sum1_scalar = vector.extract %final_sum[1] : vector<4xf32> -> f32 - %sum2_scalar = vector.extract %final_sum[2] : vector<4xf32> -> f32 - %sum3_scalar = vector.extract %final_sum[3] : vector<4xf32> -> f32 - %inverse_sum0_f32 = scalar.divf %c1_f32, %sum0_scalar : f32 - %inverse_sum1_f32 = scalar.divf %c1_f32, %sum1_scalar : f32 - %inverse_sum2_f32 = scalar.divf %c1_f32, %sum2_scalar : f32 - %inverse_sum3_f32 = scalar.divf %c1_f32, %sum3_scalar : f32 - %inverse_sum0_f16 = scalar.fptrunc %inverse_sum0_f32 : f32 to f16 - %inverse_sum1_f16 = scalar.fptrunc %inverse_sum1_f32 : f32 to f16 - %inverse_sum2_f16 = scalar.fptrunc %inverse_sum2_f32 : f32 to f16 - %inverse_sum3_f16 = scalar.fptrunc %inverse_sum3_f32 : f32 to f16 - %inverse_sum0 = vector.splat %inverse_sum0_f16 : vector<4xf16> - %inverse_sum1 = vector.splat %inverse_sum1_f16 : vector<4xf16> - %inverse_sum2 = vector.splat %inverse_sum2_f16 : vector<4xf16> - %inverse_sum3 = vector.splat %inverse_sum3_f16 : vector<4xf16> - %normalized0_f16 = vector.mulf %final_output0, %inverse_sum0 : vector<4xf16> - %normalized1_f16 = vector.mulf %final_output1, %inverse_sum1 : vector<4xf16> - %normalized2_f16 = vector.mulf %final_output2, %inverse_sum2 : vector<4xf16> - %normalized3_f16 = vector.mulf %final_output3, %inverse_sum3 : vector<4xf16> - %normalized0 = vector.extf %normalized0_f16 : vector<4xf16> to vector<4xf32> - %normalized1 = vector.extf %normalized1_f16 : vector<4xf16> to vector<4xf32> - %normalized2 = vector.extf %normalized2_f16 : vector<4xf16> to vector<4xf32> - %normalized3 = vector.extf %normalized3_f16 : vector<4xf16> to vector<4xf32> - scf.if %query_valid0 { - vector.store %normalized0, %output_view[%query_head0, %output_packet_channel] : vector<4xf32>, view<[%query_head_count]x128xf32> - } - scf.if %query_valid1 { - vector.store %normalized1, %output_view[%query_head1, %output_packet_channel] : vector<4xf32>, view<[%query_head_count]x128xf32> - } - scf.if %query_valid2 { - vector.store %normalized2, %output_view[%query_head2, %output_packet_channel] : vector<4xf32>, view<[%query_head_count]x128xf32> - } - scf.if %query_valid3 { - vector.store %normalized3, %output_view[%query_head3, %output_packet_channel] : vector<4xf32>, view<[%query_head_count]x128xf32> - } - } - kernel.return -} - -// The mask selects the first KV row exactly. QK, F16 probability conversion, -// and P*V all execute, while the expected result remains an auditable iota. -check.case public @qwen3_moe_flash_attention_decode_f32_f16_wmma_selected_row_case { - %key_value_token_count = check.literal value(64) : index - %query = check.generate.fill value(1.0) : tensor<1x128xf32> - %key = check.generate.fill value(1.0) : tensor<64x1x128xf16> - %value = check.generate.iota offset(0.0) step(0.125) : tensor<64x1x128xf16> - %mask = check.generate.iota offset(0.0) step(-10000.0) : tensor<64xf16> - %output = check.generate.fill value(-1.0) : tensor<1x128xf32> - %expected = check.generate.iota offset(0.0) step(0.125) : tensor<1x128xf32> - kernel.launch @qwen3_moe_flash_attention_decode_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %output) : [index](index, tensor<1x128xf32>, tensor<64x1x128xf16>, tensor<64x1x128xf16>, tensor<64xf16>, tensor<1x128xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<1x128xf32> - check.return -} - -// Thirty-two query heads share four KV heads in the production GQA ratio. -// Equal scores and constant values make every output exactly two while all -// head-group addressing and output ownership execute. -check.case public @qwen3_moe_flash_attention_decode_f32_f16_wmma_gqa_case { - %key_value_token_count = check.literal value(128) : index - %query = check.generate.fill value(1.0) : tensor<32x128xf32> - %key = check.generate.fill value(1.0) : tensor<128x4x128xf16> - %value = check.generate.fill value(2.0) : tensor<128x4x128xf16> - %mask = check.generate.fill value(0.0) : tensor<128xf16> - %output = check.generate.fill value(-1.0) : tensor<32x128xf32> - %expected = check.generate.fill value(2.0) : tensor<32x128xf32> - kernel.launch @qwen3_moe_flash_attention_decode_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %output) : [index](index, tensor<32x128xf32>, tensor<128x4x128xf16>, tensor<128x4x128xf16>, tensor<128xf16>, tensor<32x128xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<32x128xf32> - check.return -} - -// Sixty-five KV rows force one masked cleanup tile after a full WMMA block. -// The mask selects only that final row, whose iota values begin at 1024. -check.case public @qwen3_moe_flash_attention_decode_f32_f16_wmma_tail_case { - %key_value_token_count = check.literal value(65) : index - %query = check.generate.fill value(1.0) : tensor<1x128xf32> - %key = check.generate.fill value(1.0) : tensor<65x1x128xf16> - %value = check.generate.iota offset(0.0) step(0.125) : tensor<65x1x128xf16> - %mask = check.generate.iota offset(-64000.0) step(1000.0) : tensor<65xf16> - %output = check.generate.fill value(-1.0) : tensor<1x128xf32> - %expected = check.generate.iota offset(1024.0) step(0.125) : tensor<1x128xf32> - kernel.launch @qwen3_moe_flash_attention_decode_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %output) : [index](index, tensor<1x128xf32>, tensor<65x1x128xf16>, tensor<65x1x128xf16>, tensor<65xf16>, tensor<1x128xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<1x128xf32> - check.return -} - -check.case public @qwen3_moe_flash_attention_decode_f32_f16_wmma_benchmark_case { - %key_value_token_count = check.param.choice values([64, 65, 128, 256, 512, 768, 1024, 1280, 2048]) name("key_value_token_count") : index - %query = check.generate.fill value(0.0) : tensor<32x128xf32> - %key = check.generate.fill value(0.0) : tensor<[%key_value_token_count]x4x128xf16> - %value = check.generate.fill value(0.0) : tensor<[%key_value_token_count]x4x128xf16> - %mask = check.generate.fill value(0.0) : tensor<[%key_value_token_count]xf16> - %output = check.generate.fill value(1.0) : tensor<32x128xf32> - %expected = check.generate.fill value(0.0) : tensor<32x128xf32> - kernel.launch @qwen3_moe_flash_attention_decode_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %output) : [index](index, tensor<32x128xf32>, tensor<[%key_value_token_count]x4x128xf16>, tensor<[%key_value_token_count]x4x128xf16>, tensor<[%key_value_token_count]xf16>, tensor<32x128xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<32x128xf32> - check.return -} - -check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_selected_row_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_selected_row - -check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_gqa_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_gqa - -check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_tail_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_tail - -check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_decode_64 {key_value_token_count = 64} - -check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_decode_65 {key_value_token_count = 65} - -check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_decode_128 {key_value_token_count = 128} - -check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_decode_256 {key_value_token_count = 256} - -check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_decode_512 {key_value_token_count = 512} - -check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_decode_768 {key_value_token_count = 768} - -check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_decode_1024 {key_value_token_count = 1024} - -check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_decode_1280 {key_value_token_count = 1280} - -check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_decode_2048 {key_value_token_count = 2048} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_decode_q128_f32_f16_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_decode_q128_f32_f16_wmma.loom deleted file mode 100644 index 5a0c8037f856..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_decode_q128_f32_f16_wmma.loom +++ /dev/null @@ -1,266 +0,0 @@ -// Exact-q128 grouped-query decode FlashAttention. -// -// One eight-wave workgroup owns each KV head. The waves cooperatively cover -// all 128 KV rows for QK, normalize the complete score matrix in LDS, and then -// each own one 16-channel P*V tile. This preserves 32 active waves for the -// production 32Q/4KV shape while removing duplicate QK work, global partial -// tensors, completion atomics, and a separate resolve dispatch. -amdgpu.target @qwen3_moe_attention_decode_q128_gfx11_wave64 {subgroup_size = 64} - -config.decl @qwen3_moe.attention.query_head_count : %value: index where [range(%value, 1, 64)] - -config.decl @qwen3_moe.attention.key_value_head_count : %value: index where [range(%value, 1, 64)] - -kernel.def target(@qwen3_moe_attention_decode_q128_gfx11_wave64) @qwen3_moe_flash_attention_decode_q128_fused_f32_f16_wmma(%key_value_token_count: index) { - %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index - %c1 = index.constant 1 : index - %c512 = index.constant 512 : index - kernel.launch.config workgroups(%key_value_head_count, %c1, %c1) workgroup_size(%c512, %c1, %c1) : index -} launch(%key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %output: buffer) { - %query_head_count = config.get @qwen3_moe.attention.query_head_count : index - %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index - %workgroup_x0 = kernel.workgroup.id : index - %key_value_head = index.assume %workgroup_x0 [range(%workgroup_x0, 0, 63)] : index - %workitem0 = kernel.workitem.id : index - %workitem = index.assume %workitem0 [range(%workitem0, 0, 511)] : index - %subgroup0 = kernel.subgroup.id : index - %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index - %lane0 = kernel.subgroup.lane.id : index - %lane = index.assume %lane0 [range(%lane0, 0, 63)] : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c2 = index.constant 2 : index - %c4 = index.constant 4 : index - %c16 = index.constant 16 : index - %c32 = index.constant 32 : index - %c64 = index.constant 64 : index - %c128 = index.constant 128 : index - %c512 = index.constant 512 : index - %c0_offset = index.constant 0 : offset - %query_stage_bytes = index.constant 4352 : offset - %score_stage_bytes = index.constant 12288 : offset - %probability_stage_bytes = index.constant 6144 : offset - %product_stage_bytes = index.constant 8192 : offset - %c0_f16 = scalar.constant 0.0 : f16 - %c0_f32 = scalar.constant 0.0 : f32 - %c1_f32 = scalar.constant 1.0 : f32 - %negative_large = scalar.constant -1e+30 : f32 - %c0_f32x2 = vector.constant 0.0 : vector<2xf32> - %m = index.constant 16 : index - %n = index.constant 16 : index - %k = index.constant 16 : index - %head_size_f32 = scalar.constant 128.0 : f32 - %attention_scale = scalar.rsqrtf %head_size_f32 : f32 - %query_heads_per_key_value_head0 = index.div %query_head_count, %key_value_head_count : index - %query_heads_per_key_value_head = index.assume %query_heads_per_key_value_head0 [range(%query_heads_per_key_value_head0, 1, 16)] : index - %query_head_base = index.mul %key_value_head, %query_heads_per_key_value_head : index - %key_value_width = index.mul %key_value_head_count, %c128 : index - %key_value_head_base = index.mul %key_value_head, %c128 : index - %score_key_origin = index.mul %subgroup, %c16 : index - %query_row_base = index.mul %subgroup, %c2 : index - %query_row0 = index.add %query_row_base, %c0 : index - %query_row1 = index.add %query_row_base, %c1 : index - %query_head0 = index.add %query_head_base, %query_row0 : index - %query_head1 = index.add %query_head_base, %query_row1 : index - %query_row_valid0 = index.cmp ult, %query_row0, %query_heads_per_key_value_head : index - %query_row_valid1 = index.cmp ult, %query_row1, %query_heads_per_key_value_head : index - %query_head_in_range0 = index.cmp ult, %query_head0, %query_head_count : index - %query_head_in_range1 = index.cmp ult, %query_head1, %query_head_count : index - %query_valid0 = scalar.andi %query_row_valid0, %query_head_in_range0 : i1 - %query_valid1 = scalar.andi %query_row_valid1, %query_head_in_range1 : i1 - %query_valid0x2 = vector.from_elements %query_valid0, %query_valid0 : vector<2xi1> - %query_valid1x2 = vector.from_elements %query_valid1, %query_valid1 : vector<2xi1> - %key_token0 = index.add %lane, %c0 : index - %key_token1 = index.add %lane, %c64 : index - %output_tile_channel = index.mul %subgroup, %c16 : index - %query_transposed_layout = encoding.layout.strided [1, 136] : encoding - %probability_transposed_layout = encoding.layout.strided [1, 24] : encoding - %query_noalias, %key_noalias, %value_noalias, %mask_noalias, %output_noalias = buffer.assume.noalias %query, %key, %value, %mask, %output : buffer, buffer, buffer, buffer, buffer - %query_aligned = buffer.assume.alignment %query_noalias {minimum_alignment = 16} : buffer - %key_aligned = buffer.assume.alignment %key_noalias {minimum_alignment = 16} : buffer - %value_aligned = buffer.assume.alignment %value_noalias {minimum_alignment = 16} : buffer - %mask_aligned = buffer.assume.alignment %mask_noalias {minimum_alignment = 16} : buffer - %output_aligned = buffer.assume.alignment %output_noalias {minimum_alignment = 16} : buffer - %query_view = buffer.view %query_aligned[%c0_offset] : buffer -> view<[%query_head_count]x128xf32> - %key_view = buffer.view %key_aligned[%c0_offset] : buffer -> view<128x[%key_value_width]xf16> - %value_view = buffer.view %value_aligned[%c0_offset] : buffer -> view<128x[%key_value_width]xf16> - %mask_view = buffer.view %mask_aligned[%c0_offset] : buffer -> view<128xf16> - %output_view = buffer.view %output_aligned[%c0_offset] : buffer -> view<[%query_head_count]x128xf32> - %query_stage = buffer.alloca align(16) %query_stage_bytes : buffer - %score_stage = buffer.alloca align(16) %score_stage_bytes : buffer - %probability_stage = buffer.alloca align(16) %probability_stage_bytes : buffer - %product_stage = buffer.alloca align(16) %product_stage_bytes : buffer - %query_stage_view = buffer.view %query_stage[%c0_offset] : buffer -> view<16x136xf16> - %query_transposed_view = buffer.view %query_stage[%c0_offset] : buffer -> view<128x16xf16, %query_transposed_layout> - %score_stage_view = buffer.view %score_stage[%c0_offset] : buffer -> view<128x24xf32> - %probability_stage_view = buffer.view %probability_stage[%c0_offset] : buffer -> view<16x128xf16, %probability_transposed_layout> - %product_stage_view = buffer.view %product_stage[%c0_offset] : buffer -> view<16x128xf32> - // Scale and transpose the 16-row Q tile once for all eight score waves. - scf.for %load_iteration = [%c0 to %c4 step %c1] unroll { - %linear = index.madd %load_iteration, %c512, %workitem : index - %local_query_row = index.div %linear, %c128 : index - %query_channel = index.rem %linear, %c128 : index - %local_query_head = index.add %query_head_base, %local_query_row : index - %local_query_row_valid = index.cmp ult, %local_query_row, %query_heads_per_key_value_head : index - %local_query_head_in_range = index.cmp ult, %local_query_head, %query_head_count : index - %local_query_valid = scalar.andi %local_query_row_valid, %local_query_head_in_range : i1 - %query_value = scf.if %local_query_valid -> (f16) { - %loaded = view.load %query_view[%local_query_head, %query_channel] : view<[%query_head_count]x128xf32> -> f32 - %scaled = scalar.mulf %loaded, %attention_scale : f32 - %truncated = scalar.fptrunc %scaled : f32 to f16 - scf.yield %truncated : f16 - } else { - scf.yield %c0_f16 : f16 - } - view.store %query_value, %query_stage_view[%local_query_row, %query_channel] : f16, view<16x136xf16> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - // Each wave produces one 16x16 score slice, covering all 128 KV rows. - %score_init_values = vector.constant 0.0 : vector<4xf32> - %score_init = vector.fragment %score_init_values shape [%m, %n] : vector<4xf32> - %score_fragment = scf.for %head_tile = [%c0 to %c128 step %c16](%score_accumulator = %score_init : vector<4xf32>) -> (vector<4xf32>) unroll { - %key_channel = index.add %key_value_head_base, %head_tile : index - %key_fragment = vector.fragment.load %key_view[%score_key_origin, %key_channel] shape [%m, %k] : view<128x[%key_value_width]xf16> -> vector<16xf16> - %query_fragment = vector.fragment.load %query_transposed_view[%head_tile, %c0] shape [%k, %n] : view<128x16xf16, %query_transposed_layout> -> vector<16xf16> - %next_score_accumulator = vector.mma %key_fragment, %query_fragment, %score_accumulator : vector<16xf16>, vector<16xf16>, vector<4xf32> - scf.yield %next_score_accumulator : vector<4xf32> - } - vector.fragment.store %score_fragment, %score_stage_view[%score_key_origin, %c0] shape [%m, %n] : vector<4xf32>, view<128x24xf32> - kernel.barrier scope(workgroup) ordering(acq_rel) - // Each wave normalizes two query rows. Every lane owns two KV columns, so a - // vector reduction first combines its pair and the subgroup reduction then - // spans the complete 128-token row. - %mask0_f16 = view.load %mask_view[%key_token0] : view<128xf16> -> f16 - %mask1_f16 = view.load %mask_view[%key_token1] : view<128xf16> -> f16 - %mask0 = scalar.extf %mask0_f16 : f16 to f32 - %mask1 = scalar.extf %mask1_f16 : f16 to f32 - %mask_pair = vector.from_elements %mask0, %mask1 : vector<2xf32> - %raw_score00 = view.load %score_stage_view[%key_token0, %query_row0] : view<128x24xf32> -> f32 - %raw_score01 = view.load %score_stage_view[%key_token1, %query_row0] : view<128x24xf32> -> f32 - %raw_score10 = view.load %score_stage_view[%key_token0, %query_row1] : view<128x24xf32> -> f32 - %raw_score11 = view.load %score_stage_view[%key_token1, %query_row1] : view<128x24xf32> -> f32 - %raw_scores0 = vector.from_elements %raw_score00, %raw_score01 : vector<2xf32> - %raw_scores1 = vector.from_elements %raw_score10, %raw_score11 : vector<2xf32> - %added_scores0 = vector.addf %raw_scores0, %mask_pair : vector<2xf32> - %added_scores1 = vector.addf %raw_scores1, %mask_pair : vector<2xf32> - %negative_pair = vector.constant -1e+30 : vector<2xf32> - %masked_scores0 = vector.select %query_valid0x2, %added_scores0, %negative_pair : vector<2xf32> - %masked_scores1 = vector.select %query_valid1x2, %added_scores1, %negative_pair : vector<2xf32> - %lane_max0 = vector.reduce %masked_scores0, %negative_large : vector<2xf32>, f32 - %lane_max1 = vector.reduce %masked_scores1, %negative_large : vector<2xf32>, f32 - %lane_maxima = vector.from_elements %lane_max0, %lane_max1 : vector<2xf32> - %row_maxima = kernel.subgroup.reduce %lane_maxima : vector<2xf32> - %row_maximum0 = vector.extract %row_maxima[0] : vector<2xf32> -> f32 - %row_maximum1 = vector.extract %row_maxima[1] : vector<2xf32> -> f32 - %row_maximum0x2 = vector.splat %row_maximum0 : vector<2xf32> - %row_maximum1x2 = vector.splat %row_maximum1 : vector<2xf32> - %score_delta0 = vector.subf %masked_scores0, %row_maximum0x2 : vector<2xf32> - %score_delta1 = vector.subf %masked_scores1, %row_maximum1x2 : vector<2xf32> - %raw_probability0 = vector.expf %score_delta0 : vector<2xf32> - %raw_probability1 = vector.expf %score_delta1 : vector<2xf32> - %probability0 = vector.select %query_valid0x2, %raw_probability0, %c0_f32x2 : vector<2xf32> - %probability1 = vector.select %query_valid1x2, %raw_probability1, %c0_f32x2 : vector<2xf32> - %lane_sum0 = vector.reduce %probability0, %c0_f32 : vector<2xf32>, f32 - %lane_sum1 = vector.reduce %probability1, %c0_f32 : vector<2xf32>, f32 - %lane_sums = vector.from_elements %lane_sum0, %lane_sum1 : vector<2xf32> - %row_sums = kernel.subgroup.reduce %lane_sums : vector<2xf32> - %row_sum0 = vector.extract %row_sums[0] : vector<2xf32> -> f32 - %row_sum1 = vector.extract %row_sums[1] : vector<2xf32> -> f32 - %reciprocal_sum0 = scf.if %query_valid0 -> (f32) { - %reciprocal = scalar.divf %c1_f32, %row_sum0 : f32 - scf.yield %reciprocal : f32 - } else { - scf.yield %c0_f32 : f32 - } - %reciprocal_sum1 = scf.if %query_valid1 -> (f32) { - %reciprocal = scalar.divf %c1_f32, %row_sum1 : f32 - scf.yield %reciprocal : f32 - } else { - scf.yield %c0_f32 : f32 - } - %reciprocal_sum0x2 = vector.splat %reciprocal_sum0 : vector<2xf32> - %reciprocal_sum1x2 = vector.splat %reciprocal_sum1 : vector<2xf32> - %normalized_probability0 = vector.mulf %probability0, %reciprocal_sum0x2 : vector<2xf32> - %normalized_probability1 = vector.mulf %probability1, %reciprocal_sum1x2 : vector<2xf32> - %probability0_f16 = vector.fptrunc %normalized_probability0 : vector<2xf32> to vector<2xf16> - %probability1_f16 = vector.fptrunc %normalized_probability1 : vector<2xf32> to vector<2xf16> - %probability00 = vector.extract %probability0_f16[0] : vector<2xf16> -> f16 - %probability01 = vector.extract %probability0_f16[1] : vector<2xf16> -> f16 - %probability10 = vector.extract %probability1_f16[0] : vector<2xf16> -> f16 - %probability11 = vector.extract %probability1_f16[1] : vector<2xf16> -> f16 - view.store %probability00, %probability_stage_view[%query_row0, %key_token0] : f16, view<16x128xf16, %probability_transposed_layout> - view.store %probability01, %probability_stage_view[%query_row0, %key_token1] : f16, view<16x128xf16, %probability_transposed_layout> - view.store %probability10, %probability_stage_view[%query_row1, %key_token0] : f16, view<16x128xf16, %probability_transposed_layout> - view.store %probability11, %probability_stage_view[%query_row1, %key_token1] : f16, view<16x128xf16, %probability_transposed_layout> - kernel.barrier scope(workgroup) ordering(acq_rel) - // Each wave owns one disjoint 16-channel output tile. - %product_init_values = vector.constant 0.0 : vector<4xf32> - %product_init = vector.fragment %product_init_values shape [%m, %n] : vector<4xf32> - %value_channel = index.add %key_value_head_base, %output_tile_channel : index - %product_fragment = scf.for %key_tile = [%c0 to %c128 step %c16](%product_accumulator = %product_init : vector<4xf32>) -> (vector<4xf32>) unroll { - %probability_fragment = vector.fragment.load %probability_stage_view[%c0, %key_tile] shape [%m, %k] : view<16x128xf16, %probability_transposed_layout> -> vector<16xf16> - %value_fragment = vector.fragment.load %value_view[%key_tile, %value_channel] shape [%k, %n] : view<128x[%key_value_width]xf16> -> vector<16xf16> - %next_product_accumulator = vector.mma %probability_fragment, %value_fragment, %product_accumulator : vector<16xf16>, vector<16xf16>, vector<4xf32> - scf.yield %next_product_accumulator : vector<4xf32> - } - vector.fragment.store %product_fragment, %product_stage_view[%c0, %output_tile_channel] shape [%m, %n] : vector<4xf32>, view<16x128xf32> - kernel.barrier scope(workgroup) ordering(acq_rel) - // One vector packet per workitem drains the 16x128 result tile while - // suppressing padded GQA rows. - %output_row = index.div %workitem, %c32 : index - %output_packet = index.rem %workitem, %c32 : index - %output_channel = index.mul %output_packet, %c4 : index - %output_head = index.add %query_head_base, %output_row : index - %output_row_valid = index.cmp ult, %output_row, %query_heads_per_key_value_head : index - %output_head_in_range = index.cmp ult, %output_head, %query_head_count : index - %output_valid = scalar.andi %output_row_valid, %output_head_in_range : i1 - scf.if %output_valid { - %output_packet_value = vector.load %product_stage_view[%output_row, %output_channel] : view<16x128xf32> -> vector<4xf32> - vector.store %output_packet_value, %output_view[%output_head, %output_channel] : vector<4xf32>, view<[%query_head_count]x128xf32> - } - kernel.return -} - -// A sharp mask selects the first KV row and makes the expected result an -// auditable iota while QK, global softmax, F16 probabilities, and P*V execute. -check.case public @qwen3_moe_flash_attention_decode_q128_fused_selected_row_case { - %key_value_token_count = check.literal value(128) : index - %query = check.generate.fill value(1.0) : tensor<1x128xf32> - %key = check.generate.fill value(1.0) : tensor<128x1x128xf16> - %value = check.generate.iota offset(0.0) step(0.125) : tensor<128x1x128xf16> - %mask = check.generate.iota offset(0.0) step(-10000.0) : tensor<128xf16> - %output = check.generate.fill value(-1.0) : tensor<1x128xf32> - %expected = check.generate.iota offset(0.0) step(0.125) : tensor<1x128xf32> - kernel.launch @qwen3_moe_flash_attention_decode_q128_fused_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %output) : [index](index, tensor<1x128xf32>, tensor<128x1x128xf16>, tensor<128x1x128xf16>, tensor<128xf16>, tensor<1x128xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<1x128xf32> - check.return -} - -// The production 32Q/4KV GQA shape verifies all head groups and output tiles. -check.case public @qwen3_moe_flash_attention_decode_q128_fused_gqa_case { - %key_value_token_count = check.literal value(128) : index - %query = check.generate.fill value(1.0) : tensor<32x128xf32> - %key = check.generate.fill value(1.0) : tensor<128x4x128xf16> - %value = check.generate.fill value(2.0) : tensor<128x4x128xf16> - %mask = check.generate.fill value(0.0) : tensor<128xf16> - %output = check.generate.fill value(-1.0) : tensor<32x128xf32> - %expected = check.generate.fill value(2.0) : tensor<32x128xf32> - kernel.launch @qwen3_moe_flash_attention_decode_q128_fused_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %output) : [index](index, tensor<32x128xf32>, tensor<128x4x128xf16>, tensor<128x4x128xf16>, tensor<128xf16>, tensor<32x128xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<32x128xf32> - check.return -} - -check.case public @qwen3_moe_flash_attention_decode_q128_fused_benchmark_case { - %key_value_token_count = check.literal value(128) : index - %query = check.generate.fill value(0.0) : tensor<32x128xf32> - %key = check.generate.fill value(0.0) : tensor<128x4x128xf16> - %value = check.generate.fill value(0.0) : tensor<128x4x128xf16> - %mask = check.generate.fill value(0.0) : tensor<128xf16> - %output = check.generate.fill value(1.0) : tensor<32x128xf32> - %expected = check.generate.fill value(0.0) : tensor<32x128xf32> - kernel.launch @qwen3_moe_flash_attention_decode_q128_fused_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %output) : [index](index, tensor<32x128xf32>, tensor<128x4x128xf16>, tensor<128x4x128xf16>, tensor<128xf16>, tensor<32x128xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<32x128xf32> - check.return -} - -check.benchmark<@qwen3_moe_flash_attention_decode_q128_fused_benchmark_case> @qwen3_moe_flash_attention_decode_q128_fused diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/motifs/attention_prepare_target.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/motifs/attention_prepare_target.loom new file mode 100644 index 000000000000..feb0e4fb8a6f --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/motifs/attention_prepare_target.loom @@ -0,0 +1,4 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +amdgpu.target @qwen3_moe_attention_prepare_gfx11_wave32 {subgroup_size = 32} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/motifs/attention_qkv_postprocess_fused_decode.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/motifs/attention_qkv_postprocess_fused_decode.loom new file mode 100644 index 000000000000..d7f4d82b1a8e --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/motifs/attention_qkv_postprocess_fused_decode.loom @@ -0,0 +1,73 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @qwen3_moe.attention.qkv_postprocess_fused_decode.body(%value_uses_q6_index: index, %token_count: index, %cache_row_count: index, %q8_input: buffer, %query_weight: buffer, %key_weight: buffer, %value_weight: buffer, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer, %query_output_raw: buffer, %key_output_raw: buffer, %value_output_raw: buffer, %query_norm_weight: buffer, %key_norm_weight: buffer, %inverse_frequencies: buffer, %query_output: buffer, %key_cache: buffer, %value_cache: buffer, %completion_counters: buffer) + +template.decl @qwen3_moe.attention.qkv_quantized.body(%value_uses_q6_index: index, %publish_output: i1, %token_count: index, %token0: index, %q8_input: buffer, %query_weight: buffer, %key_weight: buffer, %value_weight: buffer, %query_output: buffer, %key_output: buffer, %value_output: buffer) + +template.decl @qwen3_moe.attention.postprocess.head_body(%arg0: i1, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: buffer, %arg6: buffer, %arg7: buffer, %arg8: buffer, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer) + +config.decl @qwen3_moe.attention.query_size : %value: index where [range(%value, 1, 262144)] + +config.decl @qwen3_moe.attention.key_value_size : %value: index where [range(%value, 1, 262144)] + +config.decl @qwen3_moe.attention.head_size : %value: index where [range(%value, 128, 128)] + +// The fused export sources the value storage format from model config. +template.def<@qwen3_moe.attention.qkv_postprocess_fused_decode.body> device @qwen3_moe_attention_qkv_postprocess_fused_decode_body(%value_uses_q6_index: index, %token_count: index, %cache_row_count: index, %q8_input: buffer, %query_weight: buffer, %key_weight: buffer, %value_weight: buffer, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer, %query_output_raw: buffer, %key_output_raw: buffer, %value_output_raw: buffer, %query_norm_weight: buffer, %key_norm_weight: buffer, %inverse_frequencies: buffer, %query_output: buffer, %key_cache: buffer, %value_cache: buffer, %completion_counters: buffer) { + %publish_projection_output = scalar.constant true : i1 + %projection_token = kernel.workgroup.id : index + template.apply<@qwen3_moe.attention.qkv_quantized.body>(%value_uses_q6_index, %publish_projection_output, %token_count, %projection_token, %q8_input, %query_weight, %key_weight, %value_weight, %query_output_raw, %key_output_raw, %value_output_raw) : (index, i1, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + %query_size0 = config.get @qwen3_moe.attention.query_size : index + %key_value_size0 = config.get @qwen3_moe.attention.key_value_size : index + %head_size0 = config.get @qwen3_moe.attention.head_size : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index + %query_size, %key_value_size, %head_size = index.assume %query_size0, %key_value_size0, %head_size0 [range(%query_size0, 128, 262144), range(%key_value_size0, 128, 262144), range(%head_size0, 128, 128), mul(%query_size0, %head_size0), mul(%key_value_size0, %head_size0)] : index, index, index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c8 = index.constant 8 : index + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %c0_offset = index.constant 0 : offset + %counter_scratch_bytes = index.constant 4 : offset + %channel_tile = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %token = index.assume %c0 [lt(%c0, %bounded_token_count)] : index + %head_tile_count = index.div %head_size, %c8 : index + %head_domain0 = index.div %channel_tile, %head_tile_count : index + %query_head_count = index.div %query_size, %head_size : index + %key_value_head_count = index.div %key_value_size, %head_size : index + %key_value_domain_count = index.mul %key_value_head_count, %c2 : index + %head_domain_count = index.add %query_head_count, %key_value_domain_count : index + %head_domain, %completion_counter_count = index.assume %head_domain0, %head_domain_count [lt(%head_domain0, %head_domain_count)] : index, index + %is_arrival_workitem = index.cmp eq, %workitem, %c0 : index + %completion_counters_aligned = buffer.assume.alignment %completion_counters {minimum_alignment = 16} : buffer + %completion_counters_view = buffer.view %completion_counters_aligned[%c0_offset] : buffer -> view<[%completion_counter_count]xi32> + %counter_scratch = buffer.alloca align(4) %counter_scratch_bytes : buffer + %counter_scratch_view = buffer.view %counter_scratch[%c0_offset] : buffer -> view<1xi32> + // Publish every producer's projection stores before the leader advances one + // workgroup arrival. The last arrival then acquires the complete head. + kernel.barrier scope(workgroup) ordering(release) + scf.if %is_arrival_workitem { + %old_counter = view.atomic.rmw %c1_i32, %completion_counters_view[%head_domain] {ordering = acq_rel, scope = device} : i32, view<[%completion_counter_count]xi32> -> i32 + view.store %old_counter, %counter_scratch_view[%c0] : i32, view<1xi32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %old_counter = view.load %counter_scratch_view[%c0] : view<1xi32> -> i32 + %head_tile_count_i32 = index.cast %head_tile_count : index to i32 + %last_head_tile_i32 = scalar.subi %head_tile_count_i32, %c1_i32 : i32 + %negative_head_tile_count_i32 = scalar.subi %c0_i32, %head_tile_count_i32 : i32 + %is_last_head_tile = scalar.cmpi eq, %old_counter, %last_head_tile_i32 : i32 + scf.if %is_last_head_tile { + kernel.barrier scope(workgroup) ordering(acquire) + %publish_postprocess_output = scalar.constant true : i1 + template.apply<@qwen3_moe.attention.postprocess.head_body>(%publish_postprocess_output, %bounded_token_count, %cache_row_count, %head_domain, %token, %positions, %key_cache_indices, %value_cache_indices, %query_output_raw, %key_output_raw, %value_output_raw, %query_norm_weight, %key_norm_weight, %inverse_frequencies, %query_output, %key_cache, %value_cache) : (i1, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + // The counter becomes reusable only after all semantic output stores. + kernel.barrier scope(workgroup) ordering(release) + scf.if %is_arrival_workitem { + view.atomic.reduce %negative_head_tile_count_i32, %completion_counters_view[%head_domain] {ordering = release, scope = device} : i32, view<[%completion_counter_count]xi32> + } + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_next_q8.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/motifs/dense_linear_q4k_next_q8.loom similarity index 52% rename from ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_next_q8.loom rename to ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/motifs/dense_linear_q4k_next_q8.loom index d6e4dea43918..30368a61bf05 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_next_q8.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/motifs/dense_linear_q4k_next_q8.loom @@ -1,34 +1,42 @@ -// Shared decode completion protocol for routed-down projections that publish -// the normalized Q8_1 x4 row consumed by the next projection boundary. Each -// x workgroup must publish one output tile before applying this body and must -// provide both its tile width and the number of target subgroups participating -// in the row reduction. The last arrival acquires the complete residual row, -// normalizes and packs it, then resets the reusable completion word after every -// Q8 store. -// -// This outer device template deliberately applies the RMSNorm/Q8 device -// template. Keeping that composition authored here ensures the compiler -// preserves the eventual kernel ancestor through nested template selection. +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @qwen3_moe.dense.q4k_q8_1_x4.body(%publish_output: i1, %token_count: index, %token0: index, %q8_input: buffer, %weight: buffer, %output: buffer) + template.decl @qwen3_moe.rmsnorm_quantize_q8_1_x4.body(%arg0: i1, %arg1: index, %arg2: index, %arg3: index, %arg4: buffer, %arg5: buffer, %arg6: buffer, %arg7: buffer) -template.decl @qwen3_moe.routed_down.next_q8_completion(%output_channels_per_workgroup0: index, %reduction_subgroup_count0: index, %token_count: index, %output_size: index, %output: buffer, %norm_weight: buffer, %completion_counter: buffer, %next_q8_output: buffer) +target.decl @qwen3_moe_attention_prepare_gfx11_wave32 + +config.decl @qwen3_moe.dense_quantized.output_size : %value: index where [range(%value, 1, 262144)] -template.def<@qwen3_moe.routed_down.next_q8_completion> device @qwen3_moe_routed_down_next_q8_completion(%output_channels_per_workgroup0: index, %reduction_subgroup_count0: index, %token_count: index, %output_size: index, %output: buffer, %norm_weight: buffer, %completion_counter: buffer, %next_q8_output: buffer) { - %output_channels_per_workgroup = index.assume %output_channels_per_workgroup0 [range(%output_channels_per_workgroup0, 1, 8)] : index - %reduction_subgroup_count = index.assume %reduction_subgroup_count0 [range(%reduction_subgroup_count0, 1, 8)] : index +// Decode-only route that publishes the normalized F32 and Q8_1 x4 rows +// consumed by the following feed-forward boundary. +kernel.def target(@qwen3_moe_attention_prepare_gfx11_wave32) @qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8(%token_count: index) { + %output_size = config.get @qwen3_moe.dense_quantized.output_size : index + %c1 = index.constant 1 : index + %c8 = index.constant 8 : index + %c256 = index.constant 256 : index + %output_tiles = index.div %output_size, %c8 : index + kernel.launch.config workgroups(%output_tiles, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %q8_input: buffer, %weight: buffer, %output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counter: buffer, %next_q8_output: buffer) { %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index - %bounded_output_size = index.assume %output_size [range(%output_size, 128, 32768), mul(%output_size, 128)] : index + %publish_output = scalar.constant true : i1 + %projection_token = kernel.workgroup.id : index + template.apply<@qwen3_moe.dense.q4k_q8_1_x4.body>(%publish_output, %bounded_token_count, %projection_token, %q8_input, %weight, %output) : (i1, index, index, buffer, buffer, buffer) + %output_size0 = config.get @qwen3_moe.dense_quantized.output_size : index + %output_size = index.assume %output_size0 [range(%output_size0, 128, 32768), mul(%output_size0, 128)] : index %token0 = kernel.workgroup.id : index %token = index.assume %token0 [lt(%token0, %bounded_token_count)] : index %c0 = index.constant 0 : index + %c8 = index.constant 8 : index %workitem = kernel.workitem.id : index - %output_tile_count = index.div %bounded_output_size, %output_channels_per_workgroup : index + %output_tile_count = index.div %output_size, %c8 : index %is_arrival_workitem = index.cmp eq, %workitem, %c0 : index %c0_i32 = scalar.constant 0 : i32 %c1_i32 = scalar.constant 1 : i32 %c0_offset = index.constant 0 : offset %counter_scratch_bytes = index.constant 4 : offset - %output_noalias, %norm_weight_noalias, %completion_counter_noalias, %next_q8_output_noalias = buffer.assume.noalias %output, %norm_weight, %completion_counter, %next_q8_output : buffer, buffer, buffer, buffer + %output_noalias, %norm_weight_noalias, %normalized_output_noalias, %completion_counter_noalias, %next_q8_output_noalias = buffer.assume.noalias %output, %norm_weight, %normalized_output, %completion_counter, %next_q8_output : buffer, buffer, buffer, buffer, buffer %completion_counter_aligned = buffer.assume.alignment %completion_counter_noalias {minimum_alignment = 16} : buffer %completion_counter_view = buffer.view %completion_counter_aligned[%c0_offset] : buffer -> view<1xi32> %counter_scratch = buffer.alloca align(4) %counter_scratch_bytes : buffer @@ -48,12 +56,13 @@ template.def<@qwen3_moe.routed_down.next_q8_completion> device @qwen3_moe_routed %is_last_output_tile = scalar.cmpi eq, %old_counter, %last_output_tile_i32 : i32 scf.if %is_last_output_tile { kernel.barrier scope(workgroup) ordering(acquire) - %publish_normalized = scalar.constant false : i1 - template.apply<@qwen3_moe.rmsnorm_quantize_q8_1_x4.body>(%publish_normalized, %reduction_subgroup_count, %bounded_token_count, %token, %output_noalias, %norm_weight_noalias, %next_q8_output_noalias, %next_q8_output_noalias) : (i1, index, index, index, buffer, buffer, buffer, buffer) + %publish_normalized = scalar.constant true : i1 + template.apply<@qwen3_moe.rmsnorm_quantize_q8_1_x4.body>(%publish_normalized, %c8, %bounded_token_count, %token, %output_noalias, %norm_weight_noalias, %normalized_output_noalias, %next_q8_output_noalias) : (i1, index, index, index, buffer, buffer, buffer, buffer) + // Reset only after every normalized F32 and Q8 store completes. kernel.barrier scope(workgroup) ordering(release) scf.if %is_arrival_workitem { view.atomic.reduce %negative_output_tile_count_i32, %completion_counter_view[%c0] {ordering = release, scope = device} : i32, view<1xi32> } } - template.return + kernel.return } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/motifs/expert_table_partition_prefill_512.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/motifs/expert_table_partition_prefill_512.loom new file mode 100644 index 000000000000..80434f33a4d0 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/motifs/expert_table_partition_prefill_512.loom @@ -0,0 +1,162 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +amdgpu.target @qwen3_moe_expert_table_partition_gfx11_wave32 {subgroup_size = 32} + +kernel.def target(@qwen3_moe_expert_table_partition_gfx11_wave32) @qwen3_moe_build_expert_table_partition_prefill_512(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index) { + %c1 = index.constant 1 : index + %launch_expert_count = index.constant 128 : index + %workgroup_size = index.constant 256 : index + kernel.launch.config workgroups(%launch_expert_count, %c1, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %route_ids: buffer, %expert_table: buffer, %partition_table: buffer, %completion_counter: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 512, 512)] : index + %bounded_route_count = index.assume %route_count [range(%route_count, 8, 8)] : index + %bounded_route_stride = index.assume %route_stride [range(%route_stride, 8, 8)] : index + %bounded_expert_count = index.assume %expert_count [range(%expert_count, 128, 128)] : index + %expert0 = kernel.workgroup.id : index + %expert, %launch_expert_count = index.assume %expert0, %bounded_expert_count [lt(%expert0, %bounded_expert_count)] : index, index + %lane0 = kernel.workitem.id : index + %lane = index.assume %lane0 [range(%lane0, 0, 255)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %workgroup_size = index.constant 256 : index + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %c7_i32 = scalar.constant 7 : i32 + %c13_i32 = scalar.constant 13 : i32 + %c0_offset = index.constant 0 : offset + %c4_bytes = index.constant 4 : offset + %assignment_count = index.mul %bounded_token_count, %bounded_route_count : index + %assignment_table_byte_base = index.scale %launch_expert_count, %c4_bytes : index, offset -> offset + %route_ids_noalias, %expert_table_noalias, %partition_table_noalias, %completion_counter_noalias = buffer.assume.noalias %route_ids, %expert_table, %partition_table, %completion_counter : buffer, buffer, buffer, buffer + %completion_counter_aligned = buffer.assume.alignment %completion_counter_noalias {minimum_alignment = 16} : buffer + %route_view = buffer.view %route_ids_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_route_stride]xi32> + %count_view = buffer.view %expert_table_noalias[%c0_offset] : buffer -> view<[%launch_expert_count]xi32> + %assignment_view = buffer.view %expert_table_noalias[%assignment_table_byte_base] : buffer -> view<[%launch_expert_count]x[%bounded_token_count]xi32> + %completion_counter_view = buffer.view %completion_counter_aligned[%c0_offset] : buffer -> view<1xi32> + + %expert_route_count = scf.for %block_base = [%c0 to %assignment_count step %workgroup_size](%matched_base = %c0_i32 : i32) -> (i32) { + %assignment = index.add %block_base, %lane : index + %in_range = index.cmp ult, %assignment, %assignment_count : index + %route_expert_i32 = scf.if %in_range -> (i32) { + %token0 = index.div %assignment, %bounded_route_count : index + %route0 = index.rem %assignment, %bounded_route_count : index + %token, %route_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %route, %route_row_stride = index.assume %route0, %bounded_route_stride [lt(%route0, %bounded_route_stride)] : index, index + %loaded = view.load %route_view[%token, %route] : view<[%bounded_token_count]x[%bounded_route_stride]xi32> -> i32 + scf.yield %loaded : i32 + } else { + %cn1_i32 = scalar.constant -1 : i32 + scf.yield %cn1_i32 : i32 + } + %route_expert0 = index.cast %route_expert_i32 : i32 to index + %route_expert = index.assume %route_expert0 [range(%route_expert0, -1, 127)] : index + %matches = index.cmp eq, %route_expert, %expert : index + %match_i32 = scf.if %matches -> (i32) { + scf.yield %c1_i32 : i32 + } else { + scf.yield %c0_i32 : i32 + } + %block_prefix = kernel.workgroup.scan %match_i32 {direction = forward, mode = exclusive} : i32 + %block_match_count_reduced = kernel.workgroup.reduce %match_i32 : i32 + %block_match_count = kernel.subgroup.broadcast.first %block_match_count_reduced : i32 + scf.if %matches { + %match_ordinal_i32 = scalar.addi %matched_base, %block_prefix : i32 + %match_ordinal0 = index.cast %match_ordinal_i32 : i32 to index + %match_ordinal = index.assume %match_ordinal0 [range(%match_ordinal0, 0, 511)] : index + %bounded_match_ordinal, %table_token_count = index.assume %match_ordinal, %bounded_token_count [lt(%match_ordinal, %bounded_token_count)] : index, index + %assignment_i32 = index.cast %assignment : index to i32 + view.store %assignment_i32, %assignment_view[%expert, %bounded_match_ordinal] : i32, view<[%launch_expert_count]x[%bounded_token_count]xi32> + } + %next_matched_base = scalar.addi %matched_base, %block_match_count : i32 + scf.yield %next_matched_base : i32 + } + + %is_lane_zero = index.cmp eq, %lane, %c0 : index + scf.if %is_lane_zero { + view.store %expert_route_count, %count_view[%expert] : i32, view<[%launch_expert_count]xi32> + } + + // Every count store precedes its workgroup's release. The last arrival + // acquires all preceding releases before any lane loads the complete table. + %local_old_counter = scf.if %is_lane_zero -> (i32) { + %old_counter = view.atomic.rmw %c1_i32, %completion_counter_view[%c0] {ordering = acq_rel, scope = device} : i32, view<1xi32> -> i32 + scf.yield %old_counter : i32 + } else { + scf.yield %c0_i32 : i32 + } + %old_counter = kernel.workgroup.reduce %local_old_counter : i32 + %expert_count_i32 = index.cast %launch_expert_count : index to i32 + %last_expert_i32 = scalar.subi %expert_count_i32, %c1_i32 : i32 + %negative_expert_count_i32 = scalar.subi %c0_i32, %expert_count_i32 : i32 + %is_last_expert = scalar.cmpi eq, %old_counter, %last_expert_i32 : i32 + scf.if %is_last_expert { + %rounded_assignment_count = index.add %assignment_count, %c31 : index + %assignment_partition_count = index.div %rounded_assignment_count, %c32 : index + %maximum_partition_count = index.add %assignment_partition_count, %launch_expert_count : index + %partition_count_view = buffer.view %partition_table_noalias[%c0_offset] : buffer -> view<1xi32> + %partition_descriptor_view = buffer.view %partition_table_noalias[%c4_bytes] : buffer -> view<[%maximum_partition_count]xi32> + %has_expert = index.cmp ult, %lane, %launch_expert_count : index + %expert_assignment_count_i32 = scf.if %has_expert -> (i32) { + %partition_expert, %table_expert_count = index.assume %lane, %launch_expert_count [lt(%lane, %launch_expert_count)] : index, index + %loaded = view.load %count_view[%partition_expert] : view<[%launch_expert_count]xi32> -> i32 + scf.yield %loaded : i32 + } else { + scf.yield %c0_i32 : i32 + } + %expert_assignment_count0 = index.cast %expert_assignment_count_i32 : i32 to index + %expert_assignment_count = index.assume %expert_assignment_count0 [range(%expert_assignment_count0, 0, 512)] : index + %rounded_expert_assignment_count = index.add %expert_assignment_count, %c31 : index + %expert_partition_count = index.div %rounded_expert_assignment_count, %c32 : index + %expert_partition_count_i32 = index.cast %expert_partition_count : index to i32 + %expert_partition_base_i32 = kernel.workgroup.scan %expert_partition_count_i32 {direction = forward, mode = exclusive} : i32 + %partition_count_i32 = kernel.workgroup.reduce %expert_partition_count_i32 : i32 + %partition_count = index.cast %partition_count_i32 : i32 to index + %bounded_partition_count, %table_partition_capacity = index.assume %partition_count, %maximum_partition_count [lt(%partition_count, %maximum_partition_count)] : index, index + %expert_partition_base0 = index.cast %expert_partition_base_i32 : i32 to index + %expert_partition_base = index.assume %expert_partition_base0 [range(%expert_partition_base0, 0, 255)] : index + scf.if %has_expert { + %partition_expert, %table_expert_count = index.assume %lane, %launch_expert_count [lt(%lane, %launch_expert_count)] : index, index + %expert_i32 = index.cast %partition_expert : index to i32 + scf.for %partition = [%c0 to %expert_partition_count step %c1] { + %descriptor_ordinal0 = index.add %expert_partition_base, %partition : index + %descriptor_ordinal, %descriptor_count = index.assume %descriptor_ordinal0, %bounded_partition_count [lt(%descriptor_ordinal0, %bounded_partition_count)] : index, index + %table_descriptor_ordinal, %table_descriptor_capacity = index.assume %descriptor_ordinal, %maximum_partition_count [lt(%descriptor_ordinal, %maximum_partition_count)] : index, index + %partition_remainder = index.rem %expert_assignment_count, %c32 : index + %has_partial_tail = index.cmp ne, %partition_remainder, %c0 : index + %partition_row_count = scf.if %has_partial_tail -> (index) { + %next_partition = index.add %partition, %c1 : index + %is_tail_partition = index.cmp eq, %next_partition, %expert_partition_count : index + %tail_row_count = scf.if %is_tail_partition -> (index) { + scf.yield %partition_remainder : index + } else { + scf.yield %c32 : index + } + scf.yield %tail_row_count : index + } else { + scf.yield %c32 : index + } + %partition_i32 = index.cast %partition : index to i32 + %partition_row_count_i32 = index.cast %partition_row_count : index to i32 + %packed_partition = scalar.shli %partition_i32, %c7_i32 : i32 + %partition_row_count_minus_one = scalar.subi %partition_row_count_i32, %c1_i32 : i32 + %packed_row_count = scalar.shli %partition_row_count_minus_one, %c13_i32 : i32 + %packed_expert_partition = scalar.ori %expert_i32, %packed_partition : i32 + %packed_descriptor = scalar.ori %packed_expert_partition, %packed_row_count : i32 + view.store %packed_descriptor, %partition_descriptor_view[%table_descriptor_ordinal] : i32, view<[%maximum_partition_count]xi32> + } + } + scf.if %is_lane_zero { + view.store %partition_count_i32, %partition_count_view[%c0] : i32, view<1xi32> + } + // The counter cannot become reusable until every descriptor store is + // complete and visible to the following dispatch. + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.if %is_lane_zero { + view.atomic.reduce %negative_expert_count_i32, %completion_counter_view[%c0] {ordering = release, scope = device} : i32, view<1xi32> + } + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_q4k.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_q4k.loom deleted file mode 100644 index 9a1f150a8e12..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_q4k.loom +++ /dev/null @@ -1,375 +0,0 @@ -// Fuses the Qwen3 MoE Q4_K down projection with route weighting, top-8 -// reduction, and residual publication. The input contains one Q8_1 x4 row per -// [token, route], while weights remain in GGUF's raw -// [expert, output, K / 256, 144 bytes] layout. -// -// Route IDs retain an independent physical stride because llama.cpp selects a -// top-8 view from its 128-entry argsort storage. Normalized route weights are -// compact [token, route]. The output enters containing the residual and is -// updated in place, so the fused boundary never materializes an unweighted or -// route-indexed [token, route, hidden] down-projection tensor. -// -// One wave owns one output channel and contracts all selected routes in -// registers. Four eight-lane cohorts contract four independent routes while -// walking every Q4_K block. Qwen's eight-route, three-block decode shape uses -// all 32 lanes across two route batches without addressing synthetic blocks. -template.decl @qwen3_moe.routed_down.next_q8_completion(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: buffer, %arg5: buffer, %arg6: buffer, %arg7: buffer) - -template.decl @qwen3_moe.routed_down.q4k_q8_1_x4.body(%publish_output: i1, %token_count: index, %token: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) - -template.decl @ggml.matmul_id.vector_weighted.q4k_q8.wave32_body(%publish_output: i1, %token_count: index, %token: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) - -template.decl @ggml.matmul_id.vector_weighted.next_q8_completion(%output_channels_per_workgroup0: index, %reduction_subgroup_count0: index, %token_count: index, %output_size: index, %output: buffer, %norm_weight: buffer, %completion_counter: buffer, %next_q8_output: buffer) - -target.decl @qwen3_moe_attention_prepare_gfx11_wave32 - -config.decl @qwen3_moe.routed_down.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] - -config.decl @qwen3_moe.routed_down.route_count : %value: index where [range(%value, 1, 8)] - -config.decl @qwen3_moe.routed_down.expert_count : %value: index where [range(%value, 1, 512)] - -config.decl @qwen3_moe.routed_down.output_size : %value: index where [range(%value, 1, 4096)] - -config.decl @qwen3_moe.workload.token_capacity : %value: index where [range(%value, 1, 2048)] - -kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count$26: index, %input_size$27: index) launch(%token_count$28: index, %input_size$29: index, %input: buffer, %output: buffer) - -func.decl @qwen3_moe_q4k_q8_1_x4_cohort_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %cohort_lane: index) -> (f32) - -kernel.decl @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4(%token_count$39: index) launch(%token_count$40: index, %input: buffer, %weight: buffer, %q8_output: buffer) - -kernel.decl @ggml_mul_mat_id_vector_weighted_wave32_publish_q8(%token_count$41: index, %input_size$42: index, %route_count$43: index, %route_id_stride$44: index, %expert_count$45: index, %output_size$46: index) launch(%token_count$47: index, %input_size$48: index, %route_count$49: index, %route_id_stride$50: index, %expert_count$51: index, %output_size$52: index, %q8_input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer, %norm_weight: buffer, %completion_counter: buffer, %next_q8_output: buffer) - -// Shared Q4_K contraction and residual publication. Each lane consumes both -// nibbles of a packed code load before the subgroup reduces the row. -template.def<@qwen3_moe.routed_down.q4k_q8_1_x4.body> device @qwen3_moe_routed_down_q4k_q8_1_x4_body(%publish_output: i1, %token_count: index, %token: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) { - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index - %bounded_token, %body_token_count = index.assume %token, %bounded_token_count [lt(%token, %bounded_token_count)] : index, index - %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index - %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index - %bounded_route_id_stride = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index - %bounded_expert_count = index.assume %expert_count [range(%expert_count, 1, 512)] : index - %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144)] : index - %channel_tile = kernel.workgroup.id : index - %subgroup0 = kernel.subgroup.id : index - %lane = kernel.subgroup.lane.id : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c3 = index.constant 3 : index - %c4 = index.constant 4 : index - %c8 = index.constant 8 : index - %c128 = index.constant 128 : index - %c144 = index.constant 144 : index - %c256 = index.constant 256 : index - %c1_byte = index.constant 1 : offset - %c0_i32 = scalar.constant 0 : i32 - %c0_f32 = scalar.constant 0.0 : f32 - %c0_offset = index.constant 0 : offset - %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index - %cohort0 = index.div %lane, %c8 : index - %cohort = index.assume %cohort0 [range(%cohort0, 0, 3)] : index - %cohort_lane0 = index.rem %lane, %c8 : index - %cohort_lane = index.assume %cohort_lane0 [range(%cohort_lane0, 0, 7)] : index - %channel_base = index.mul %channel_tile, %c8 : index - %channel = index.add %channel_base, %subgroup : index - %valid_channel = index.cmp ult, %channel, %bounded_output_size : index - %lane_i32 = index.cast %lane : index to i32 - %is_lane_zero = scalar.cmpi eq, %lane_i32, %c0_i32 : i32 - %q4_block_count = index.div %bounded_input_size, %c256 : index - %weight_row_byte_count = index.mul %q4_block_count, %c144 : index - %q8_group_count = index.div %bounded_input_size, %c128 : index - %q8_row_byte_count = index.mul %q8_group_count, %c144 : index - %padded_route_count = index.add %bounded_route_count, %c3 : index - %route_batch_count = index.div %padded_route_count, %c4 : index - %q8_noalias, %route_id_noalias, %route_weight_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %q8_input, %route_ids, %route_weights, %weight, %output : buffer, buffer, buffer, buffer, buffer - %route_id_view = buffer.view %route_id_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_route_id_stride]xi32> - %route_weight_view = buffer.view %route_weight_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_route_count]xf32> - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_output_size]xf32> - %lane_has_route = index.cmp ult, %lane, %bounded_route_count : index - %lane_within_route_stride = index.cmp ult, %lane, %bounded_route_id_stride : index - %loads_active_route = scalar.andi %publish_output, %lane_has_route : i1 - %loads_route_metadata = scalar.andi %loads_active_route, %lane_within_route_stride : i1 - %lane_expert_i32, %lane_route_weight = scf.if %loads_route_metadata -> (i32, f32) { - %route_lane0, %metadata_route_count = index.assume %lane, %bounded_route_count [lt(%lane, %bounded_route_count)] : index, index - %route_lane, %metadata_route_id_stride = index.assume %route_lane0, %bounded_route_id_stride [lt(%route_lane0, %bounded_route_id_stride)] : index, index - %expert_i32 = view.load %route_id_view[%bounded_token, %route_lane] : view<[%body_token_count]x[%bounded_route_id_stride]xi32> -> i32 - %route_weight = view.load %route_weight_view[%bounded_token, %route_lane0] : view<[%body_token_count]x[%bounded_route_count]xf32> -> f32 - scf.yield %expert_i32, %route_weight : i32, f32 - } else { - scf.yield %c0_i32, %c0_f32 : i32, f32 - } - %computes_channel = scalar.andi %publish_output, %valid_channel : i1 - %routed_lane_sum = scf.if %computes_channel -> (f32) { - %sum = scf.for %route_batch = [%c0 to %route_batch_count step %c1](%route_acc = %c0_f32 : f32) -> (f32) unroll { - %route_batch_base = index.mul %route_batch, %c4 : index - %route0 = index.add %route_batch_base, %cohort : index - %route = index.assume %route0 [range(%route0, 0, 7)] : index - %active_route = index.cmp ult, %route, %bounded_route_count : index - %weighted_lane = scf.if %active_route -> (f32) { - %bounded_route, %body_route_count = index.assume %route, %bounded_route_count [lt(%route, %bounded_route_count)] : index, index - %route_i32 = index.cast %bounded_route : index to i32 - %expert_i32 = kernel.subgroup.broadcast %lane_expert_i32 from %route_i32 : i32, i32 - %expert0 = index.cast %expert_i32 : i32 to index - %expert1 = index.assume %expert0 [range(%expert0, 0, 511)] : index - %expert, %weight_expert_count = index.assume %expert1, %bounded_expert_count [lt(%expert1, %bounded_expert_count)] : index, index - %expert_output_base = index.mul %expert, %bounded_output_size : index - %expert_channel = index.add %expert_output_base, %channel : index - %weight_row_byte_index = index.mul %expert_channel, %weight_row_byte_count : index - %weight_row_byte_base = index.scale %weight_row_byte_index, %c1_byte : index, offset -> offset - %q8_row_base0 = index.mul %bounded_token, %body_route_count : index - %q8_row = index.add %q8_row_base0, %bounded_route : index - %q8_row_byte_index = index.mul %q8_row, %q8_row_byte_count : index - %q8_row_byte_base = index.scale %q8_row_byte_index, %c1_byte : index, offset -> offset - %route_lane_sum = func.call @qwen3_moe_q4k_q8_1_x4_cohort_row_lane(%bounded_input_size, %weight_noalias, %weight_row_byte_base, %q8_noalias, %q8_row_byte_base, %cohort_lane) : (index, buffer, offset, buffer, offset, index) -> (f32) - %route_weight = kernel.subgroup.broadcast %lane_route_weight from %route_i32 : f32, i32 - %weighted = scalar.mulf %route_lane_sum, %route_weight : f32 - scf.yield %weighted : f32 - } else { - scf.yield %c0_f32 : f32 - } - %next = scalar.addf %route_acc, %weighted_lane : f32 - scf.yield %next : f32 - } - scf.yield %sum : f32 - } else { - scf.yield %c0_f32 : f32 - } - %routed_sum = kernel.subgroup.reduce %routed_lane_sum : f32 - %writes_active_channel = scalar.andi %publish_output, %valid_channel : i1 - %writes_output = scalar.andi %writes_active_channel, %is_lane_zero : i1 - scf.if %writes_output { - %bounded_channel, %body_output_size = index.assume %channel, %bounded_output_size [lt(%channel, %bounded_output_size)] : index, index - %residual = view.load %output_view[%bounded_token, %bounded_channel] : view<[%body_token_count]x[%bounded_output_size]xf32> -> f32 - %result = scalar.addf %residual, %routed_sum : f32 - view.store %result, %output_view[%bounded_token, %bounded_channel] : f32, view<[%body_token_count]x[%bounded_output_size]xf32> - } - template.return -} - -kernel.def @qwen3_moe_routed_down_q4k_q8_1_x4(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index) { - %token_capacity = config.get @qwen3_moe.workload.token_capacity : index - %configured_output_size = config.get @qwen3_moe.routed_down.output_size : index - %c8 = index.constant 8 : index - %c7 = index.constant 7 : index - %c1 = index.constant 1 : index - %workgroup_size = index.constant 256 : index - %padded_output_size = index.add %configured_output_size, %c7 : index - %output_tiles = index.div %padded_output_size, %c8 : index - kernel.launch.config workgroups(%output_tiles, %token_capacity, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index -} launch(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) { - %token_capacity = config.get @qwen3_moe.workload.token_capacity : index - %configured_input_size0 = config.get @qwen3_moe.routed_down.input_size : index - %configured_route_count0 = config.get @qwen3_moe.routed_down.route_count : index - %configured_expert_count0 = config.get @qwen3_moe.routed_down.expert_count : index - %configured_output_size0 = config.get @qwen3_moe.routed_down.output_size : index - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index - %bounded_input_size, %configured_input_size = index.assume %input_size, %configured_input_size0 [range(%input_size, 256, 32768), mul(%input_size, 256), eq(%input_size, %configured_input_size0)] : index, index - %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 8), eq(%route_count, %configured_route_count0)] : index, index - %bounded_route_id_stride = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index - %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 1, 512), eq(%expert_count, %configured_expert_count0)] : index, index - %bounded_output_size, %configured_output_size = index.assume %output_size, %configured_output_size0 [range(%output_size, 1, 4096), eq(%output_size, %configured_output_size0)] : index, index - %token0 = kernel.workgroup.id : index - %c0 = index.constant 0 : index - %valid_token = index.cmp ult, %token0, %bounded_token_count : index - %safe_token0 = scf.select %valid_token, %token0, %c0 : index - %safe_token, %body_token_count = index.assume %safe_token0, %bounded_token_count [lt(%safe_token0, %bounded_token_count)] : index, index - template.apply<@ggml.matmul_id.vector_weighted.q4k_q8.wave32_body>(%valid_token, %body_token_count, %safe_token, %configured_input_size, %configured_route_count, %bounded_route_id_stride, %configured_expert_count, %configured_output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : (i1, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer) - kernel.return -} - -// The production hidden width requires all 256 output tiles to arrive before -// normalization. Compare both residual and packed Q8 publication with the -// ordinary two-dispatch composition, then reuse the same completion word. -check.case public @qwen3_moe_routed_down_q4k_q8_1_x4_next_q8_differential_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(2) : index - %route_id_stride = check.literal value(4) : index - %expert_count = check.literal value(2) : index - %output_size = check.literal value(2048) : index - %routed_row_count = check.literal value(2) : index - %routed_input = check.generate.fill value(0.00390625) : tensor<2x768xf32> - %q8_input = check.generate.fill value(0) : tensor<2x864xi8> - %route_ids = check.generate.iota offset(0) step(1) period(2) : tensor<1x4xi32> - %route_weights = check.generate.fill value(0.5) : tensor<1x2xf32> - %weight = check.generate.iota offset(-72) step(1) period(144) : tensor<2x2048x3x144xi8> - %norm_weight = check.generate.iota offset(-1.0) step(0.0009765625) : tensor<2048xf32> - %expected_output = check.generate.iota offset(-0.5) step(0.00048828125) period(2048) : tensor<1x2048xf32> - %actual_output0 = check.generate.iota offset(-0.5) step(0.00048828125) period(2048) : tensor<1x2048xf32> - %actual_output1 = check.generate.iota offset(-0.5) step(0.00048828125) period(2048) : tensor<1x2048xf32> - %expected_q8 = check.generate.fill value(0) : tensor<2304xi8> - %actual_q8_0 = check.generate.fill value(1) : tensor<2304xi8> - %actual_q8_1 = check.generate.fill value(1) : tensor<2304xi8> - %completion_counter = check.generate.fill value(0) : tensor<1xi32> - %expected_counter = check.generate.fill value(0) : tensor<1xi32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %input_size](%routed_row_count, %input_size, %routed_input, %q8_input) : [index, index](index, index, tensor<2x768xf32>, tensor<2x864xi8>) - kernel.launch @qwen3_moe_routed_down_q4k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %expected_output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x864xi8>, tensor<1x4xi32>, tensor<1x2xf32>, tensor<2x2048x3x144xi8>, tensor<1x2048xf32>) - kernel.launch @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4[%token_count](%token_count, %expected_output, %norm_weight, %expected_q8) : [index](index, tensor<1x2048xf32>, tensor<2048xf32>, tensor<2304xi8>) - kernel.launch @ggml_mul_mat_id_vector_weighted_wave32_publish_q8[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %actual_output0, %norm_weight, %completion_counter, %actual_q8_0) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x864xi8>, tensor<1x4xi32>, tensor<1x2xf32>, tensor<2x2048x3x144xi8>, tensor<1x2048xf32>, tensor<2048xf32>, tensor<1xi32>, tensor<2304xi8>) - kernel.launch @ggml_mul_mat_id_vector_weighted_wave32_publish_q8[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %actual_output1, %norm_weight, %completion_counter, %actual_q8_1) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x864xi8>, tensor<1x4xi32>, tensor<1x2xf32>, tensor<2x2048x3x144xi8>, tensor<1x2048xf32>, tensor<2048xf32>, tensor<1xi32>, tensor<2304xi8>) - check.expect.close actual(%actual_output0) expected(%expected_output) atol(9.9999999999999995e-07) rtol(9.9999999999999995e-07) nan(same) : tensor<1x2048xf32> - check.expect.close actual(%actual_output1) expected(%expected_output) atol(9.9999999999999995e-07) rtol(9.9999999999999995e-07) nan(same) : tensor<1x2048xf32> - check.expect.equal actual(%actual_q8_0) expected(%expected_q8) : tensor<2304xi8> - check.expect.equal actual(%actual_q8_1) expected(%expected_q8) : tensor<2304xi8> - check.expect.equal actual(%completion_counter) expected(%expected_counter) : tensor<1xi32> - check.return -} - -check.case public @qwen3_moe_routed_down_q4k_q8_1_x4_next_q8_benchmark_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(8) : index - %route_id_stride = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(2048) : index - %q8_input = check.generate.fill value(0) : tensor<1x8x864xi8> - %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<1x8xi32> - %route_weights = check.generate.fill value(0.125) : tensor<1x8xf32> - %weight = check.generate.fill value(0) : tensor<128x2048x3x144xi8> - %output = check.generate.fill value(1.0) : tensor<1x2048xf32> - %norm_weight = check.generate.fill value(1.0) : tensor<2048xf32> - %completion_counter = check.generate.fill value(0) : tensor<1xi32> - %next_q8_output = check.generate.fill value(0) : tensor<2304xi8> - kernel.launch @ggml_mul_mat_id_vector_weighted_wave32_publish_q8[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output, %norm_weight, %completion_counter, %next_q8_output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<1x8x864xi8>, tensor<1x8xi32>, tensor<1x8xf32>, tensor<128x2048x3x144xi8>, tensor<1x2048xf32>, tensor<2048xf32>, tensor<1xi32>, tensor<2304xi8>) - check.return -} - -check.case public @qwen3_moe_routed_down_q4k_q8_1_x4_next_q8_composed_benchmark_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(8) : index - %route_id_stride = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(2048) : index - %q8_input = check.generate.fill value(0) : tensor<1x8x864xi8> - %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<1x8xi32> - %route_weights = check.generate.fill value(0.125) : tensor<1x8xf32> - %weight = check.generate.fill value(0) : tensor<128x2048x3x144xi8> - %output = check.generate.fill value(1.0) : tensor<1x2048xf32> - %norm_weight = check.generate.fill value(1.0) : tensor<2048xf32> - %next_q8_output = check.generate.fill value(0) : tensor<2304xi8> - kernel.launch @qwen3_moe_routed_down_q4k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<1x8x864xi8>, tensor<1x8xi32>, tensor<1x8xf32>, tensor<128x2048x3x144xi8>, tensor<1x2048xf32>) - kernel.launch @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4[%token_count](%token_count, %output, %norm_weight, %next_q8_output) : [index](index, tensor<1x2048xf32>, tensor<2048xf32>, tensor<2304xi8>) - check.return -} - -// Two selected experts exercise the noncompact route-ID stride, normalized -// weighted reduction, in-place residual update, the odd Q4_K block tail at -// K=768, and the output tile tail. Uniform 0xaa bytes decode to a deterministic -// nonzero row. -check.case public @qwen3_moe_routed_down_q4k_q8_1_x4_nonzero_residual_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(2) : index - %route_id_stride = check.literal value(4) : index - %expert_count = check.literal value(2) : index - %output_size = check.literal value(9) : index - %routed_input = check.generate.fill value(0.00390625) : tensor<2x768xf32> - %q8_input = check.generate.fill value(0) : tensor<2x864xi8> - %route_ids = check.generate.iota offset(0) step(1) period(2) : tensor<1x4xi32> - %route_weights = check.generate.fill value(0.5) : tensor<1x2xf32> - %weight = check.generate.fill value(-86) : tensor<2x9x3x144xi8> - %output = check.generate.fill value(1.0) : tensor<1x9xf32> - %expected = check.generate.fill value(-58.0394287109375) : tensor<1x9xf32> - %routed_row_count = check.literal value(2) : index - kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %input_size](%routed_row_count, %input_size, %routed_input, %q8_input) : [index, index](index, index, tensor<2x768xf32>, tensor<2x864xi8>) - kernel.launch @qwen3_moe_routed_down_q4k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x864xi8>, tensor<1x4xi32>, tensor<1x2xf32>, tensor<2x9x3x144xi8>, tensor<1x9xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.25) rtol(9.9999999999999995e-07) nan(same) : tensor<1x9xf32> - check.return -} - -// Two tokens use different route-weight sums so token, route, packed-row, and -// output addressing cannot collapse to the single-token case. -check.case public @qwen3_moe_routed_down_q4k_q8_1_x4_nonzero_prefill_case { - %token_count = check.literal value(2) : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(2) : index - %route_id_stride = check.literal value(4) : index - %expert_count = check.literal value(2) : index - %output_size = check.literal value(1) : index - %routed_row_count = check.literal value(4) : index - %routed_input = check.generate.fill value(0.00390625) : tensor<2x2x768xf32> - %q8_input = check.generate.fill value(0) : tensor<2x2x864xi8> - %route_ids = check.generate.iota offset(0) step(1) period(2) : tensor<2x4xi32> - %route_weights = check.generate.iota offset(0.25) step(0.25) period(4) : tensor<2x2xf32> - %weight = check.generate.fill value(-86) : tensor<2x1x3x144xi8> - %output = check.generate.fill value(1.0) : tensor<2x1xf32> - %expected = check.generate.iota offset(-43.279571533203125) step(-59.0394287109375) : tensor<2x1xf32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %input_size](%routed_row_count, %input_size, %routed_input, %q8_input) : [index, index](index, index, tensor<2x2x768xf32>, tensor<2x2x864xi8>) - kernel.launch @qwen3_moe_routed_down_q4k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x2x864xi8>, tensor<2x4xi32>, tensor<2x2xf32>, tensor<2x1x3x144xi8>, tensor<2x1xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.25) rtol(9.9999999999999995e-07) nan(same) : tensor<2x1xf32> - check.return -} - -check.case public @qwen3_moe_routed_down_q4k_q8_1_x4_benchmark_case { - %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512]) name("token_count") : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(8) : index - %route_id_stride = check.literal value(128) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(2048) : index - %q8_input = check.generate.fill value(0) : tensor<[%token_count]x8x864xi8> - %route_ids = check.generate.iota offset(0) step(1) period(127) : tensor<[%token_count]x128xi32> - %route_weights = check.generate.fill value(0.125) : tensor<[%token_count]x8xf32> - %weight = check.generate.fill value(0) : tensor<128x2048x3x144xi8> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - %expected = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - kernel.launch @qwen3_moe_routed_down_q4k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<[%token_count]x8x864xi8>, tensor<[%token_count]x128xi32>, tensor<[%token_count]x8xf32>, tensor<128x2048x3x144xi8>, tensor<[%token_count]x2048xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> - check.return -} - -check.case public @qwen3_moe_routed_down_q4k_q8_1_x4_pipeline_benchmark_case { - %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512]) name("token_count") : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(8) : index - %route_id_stride = check.literal value(128) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(2048) : index - // Each 768-element route contains six complete Q8_1 x4 groups, so packing - // the eight contiguous routes as one 6,144-element token row preserves the - // exact per-route physical layout without a derived testbench scalar. - %routed_input_size = check.literal value(6144) : index - %routed_input = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf32> - %q8_input = check.generate.fill value(1) : tensor<[%token_count]x8x864xi8> - %route_ids = check.generate.iota offset(0) step(1) period(127) : tensor<[%token_count]x128xi32> - %route_weights = check.generate.fill value(0.125) : tensor<[%token_count]x8xf32> - %weight = check.generate.fill value(0) : tensor<128x2048x3x144xi8> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - %expected = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %routed_input_size](%token_count, %routed_input_size, %routed_input, %q8_input) : [index, index](index, index, tensor<[%token_count]x8x768xf32>, tensor<[%token_count]x8x864xi8>) - kernel.launch @qwen3_moe_routed_down_q4k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<[%token_count]x8x864xi8>, tensor<[%token_count]x128xi32>, tensor<[%token_count]x8xf32>, tensor<128x2048x3x144xi8>, tensor<[%token_count]x2048xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> - check.return -} - -check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_nonzero_residual_case> @qwen3_moe_routed_down_q4k_q8_1_x4_small - -check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_decode {token_count = 1} - -check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_small_batch_2 {token_count = 2} - -check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_small_batch_4 {token_count = 4} - -check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_small_batch_8 {token_count = 8} - -check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_small_batch_16 {token_count = 16} - -check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_prefill_32 {token_count = 32} - -check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_prefill_128 {token_count = 128} - -check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_prefill_512 {token_count = 512} - -check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_pipeline_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_pipeline_decode {token_count = 1} - -check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_pipeline_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_pipeline_prefill_32 {token_count = 32} - -check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_pipeline_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_pipeline_prefill_128 {token_count = 128} - -check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_pipeline_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_pipeline_prefill_512 {token_count = 512} - -check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_next_q8_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_next_q8_decode - -check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_next_q8_composed_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_next_q8_composed_decode diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_q6k.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_q6k.loom deleted file mode 100644 index 62399190fe25..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_q6k.loom +++ /dev/null @@ -1,721 +0,0 @@ -// Fuses the Qwen3 MoE Q6_K down projection with route weighting, top-8 -// reduction, and residual publication. The input contains one Q8_1 x4 row per -// [token, route], while weights remain in GGUF's raw -// [expert, output, K / 256, 210 bytes] layout. -// -// Route IDs retain an independent physical stride because llama.cpp selects a -// top-8 view from its 128-entry argsort storage. Normalized route weights are -// compact [token, route]. The output enters containing the residual and is -// updated in place, so the fused boundary never materializes an unweighted or -// route-indexed [token, route, hidden] down-projection tensor. -// -// The Q8_1 provider assigns one output channel to a wave and reduces all routes -// in registers. The direct-F32 provider mirrors llama.cpp's gfx1151 Vulkan -// decode algorithm: four wave64 subgroups compute four channels, each subgroup -// contracts four raw Q6_K blocks at a time, and route weighting and residual -// publication remain fused. A grouped prefill provider can reuse the same raw -// Q6_K row primitive while staging each expert row across multiple tokens. -template.decl @qwen3_moe.rmsnorm_quantize_q8_1_x4.body(%arg0: i1, %arg1: index, %arg2: index, %arg3: index, %arg4: buffer, %arg5: buffer, %arg6: buffer, %arg7: buffer) - -template.decl @qwen3_moe.routed_down.next_q8_completion(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: buffer, %arg5: buffer, %arg6: buffer, %arg7: buffer) - -template.decl @qwen3_moe.routed_down.q6k_f32.body(%publish_output: i1, %subgroup_count: index, %block_step: index, %scale_stage_bytes: offset, %token_count: index, %token: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) - -template.decl @qwen3_moe.routed_down.q6k_f32.pair_body(%token_count: index, %token: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) - -template.decl @ggml.matmul_id.vector_weighted.q6k_f32.wave64_pair_body(%token_count: index, %token: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) - -template.decl @ggml.matmul_id.vector_weighted.next_q8_completion(%output_channels_per_workgroup0: index, %reduction_subgroup_count0: index, %token_count: index, %output_size: index, %output: buffer, %norm_weight: buffer, %completion_counter: buffer, %next_q8_output: buffer) - -template.decl @qwen3_moe.routed_down.q6k_q8_1_x4.body(%publish_output: i1, %token_count: index, %token: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) - -amdgpu.target @qwen3_moe_routed_down_q6k_gfx11_wave64 {subgroup_size = 64} - -target.decl @qwen3_moe_attention_prepare_gfx11_wave32 - -config.decl @qwen3_moe.routed_down.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] - -config.decl @qwen3_moe.routed_down.route_count : %value: index where [range(%value, 1, 8)] - -config.decl @qwen3_moe.routed_down.expert_count : %value: index where [range(%value, 1, 512)] - -config.decl @qwen3_moe.routed_down.output_size : %value: index where [range(%value, 1, 4096)] - -config.decl @qwen3_moe.workload.token_capacity : %value: index where [range(%value, 1, 2048)] - -kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count$62: index, %input_size$63: index) launch(%token_count$64: index, %input_size$65: index, %input: buffer, %output: buffer) - -func.decl @ggml_q6k_q8_1_x4_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32) - -func.decl @ggml_q6k_stage_f32_scales(%input_size: index, %row: index, %block: index, %frame_count: index, %frame: index, %lane: index, %weight: buffer, %scale_stage: buffer) - -func.decl @ggml_q6k_load_f32_block(%token_count: index, %input_size: index, %token: index, %block: index, %lane: index, %input: buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) - -func.decl @ggml_q6k_f32_block_row(%input_size: index, %row: index, %block: index, %frame_count: index, %frame: index, %lane: index, %weight: buffer, %scale_stage: buffer, %input0: vector<4xf32>, %input1: vector<4xf32>, %input2: vector<4xf32>, %input3: vector<4xf32>) -> (f32) - -kernel.decl @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4(%token_count$106: index) launch(%token_count$107: index, %input: buffer, %weight: buffer, %q8_output: buffer) - -kernel.decl @ggml_mul_mat_id_vector_weighted_wave64_publish_q8(%token_count$108: index, %input_size$109: index, %route_count$110: index, %route_id_stride$111: index, %expert_count$112: index, %output_size$113: index) launch(%token_count$114: index, %input_size$115: index, %route_count$116: index, %route_id_stride$117: index, %expert_count$118: index, %output_size$119: index, %input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer, %norm_weight: buffer, %completion_counter: buffer, %next_q8_output: buffer) - -// Shared Q6_K contraction and residual publication used by ordinary and -// completion-fused exports with identical launch geometry. -template.def<@qwen3_moe.routed_down.q6k_q8_1_x4.body> device @qwen3_moe_routed_down_q6k_q8_1_x4_body(%publish_output: i1, %token_count: index, %token: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) { - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index - %bounded_token, %body_token_count = index.assume %token, %bounded_token_count [lt(%token, %bounded_token_count)] : index, index - %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index - %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index - %bounded_route_id_stride = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index - %bounded_expert_count = index.assume %expert_count [range(%expert_count, 1, 512)] : index - %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144)] : index - %channel_tile = kernel.workgroup.id : index - %subgroup0 = kernel.subgroup.id : index - %lane = kernel.subgroup.lane.id : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c8 = index.constant 8 : index - %c128 = index.constant 128 : index - %c144_bytes = index.constant 144 : offset - %c210_bytes = index.constant 210 : offset - %c256 = index.constant 256 : index - %c0_i32 = scalar.constant 0 : i32 - %c0_f32 = scalar.constant 0.0 : f32 - %c0_offset = index.constant 0 : offset - %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index - %channel_base = index.mul %channel_tile, %c8 : index - %channel = index.add %channel_base, %subgroup : index - %valid_channel = index.cmp ult, %channel, %bounded_output_size : index - %lane_i32 = index.cast %lane : index to i32 - %is_lane_zero = scalar.cmpi eq, %lane_i32, %c0_i32 : i32 - %q6_block_count = index.div %bounded_input_size, %c256 : index - %weight_row_bytes = index.scale %q6_block_count, %c210_bytes : index, offset -> offset - %q8_group_count = index.div %bounded_input_size, %c128 : index - %q8_row_bytes = index.scale %q8_group_count, %c144_bytes : index, offset -> offset - %q8_noalias, %route_id_noalias, %route_weight_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %q8_input, %route_ids, %route_weights, %weight, %output : buffer, buffer, buffer, buffer, buffer - %route_id_view = buffer.view %route_id_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_route_id_stride]xi32> - %route_weight_view = buffer.view %route_weight_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_route_count]xf32> - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_output_size]xf32> - %lane_has_route = index.cmp ult, %lane, %bounded_route_count : index - %lane_within_route_stride = index.cmp ult, %lane, %bounded_route_id_stride : index - %loads_active_route = scalar.andi %publish_output, %lane_has_route : i1 - %loads_route_metadata = scalar.andi %loads_active_route, %lane_within_route_stride : i1 - %lane_expert_i32, %lane_route_weight = scf.if %loads_route_metadata -> (i32, f32) { - %route_lane0, %metadata_route_count = index.assume %lane, %bounded_route_count [lt(%lane, %bounded_route_count)] : index, index - %route_lane, %metadata_route_id_stride = index.assume %route_lane0, %bounded_route_id_stride [lt(%route_lane0, %bounded_route_id_stride)] : index, index - %expert_i32 = view.load %route_id_view[%bounded_token, %route_lane] : view<[%body_token_count]x[%bounded_route_id_stride]xi32> -> i32 - %route_weight = view.load %route_weight_view[%bounded_token, %route_lane0] : view<[%body_token_count]x[%bounded_route_count]xf32> -> f32 - scf.yield %expert_i32, %route_weight : i32, f32 - } else { - scf.yield %c0_i32, %c0_f32 : i32, f32 - } - %computes_channel = scalar.andi %publish_output, %valid_channel : i1 - %routed_lane_sum = scf.if %computes_channel -> (f32) { - %sum = scf.for %route = [%c0 to %bounded_route_count step %c1](%route_acc = %c0_f32 : f32) -> (f32) unroll { - %route_i32 = index.cast %route : index to i32 - %expert_i32 = kernel.subgroup.broadcast %lane_expert_i32 from %route_i32 : i32, i32 - %expert0 = index.cast %expert_i32 : i32 to index - %expert1 = index.assume %expert0 [range(%expert0, 0, 511)] : index - %expert, %weight_expert_count = index.assume %expert1, %bounded_expert_count [lt(%expert1, %bounded_expert_count)] : index, index - %expert_output_base = index.mul %expert, %bounded_output_size : index - %expert_channel = index.add %expert_output_base, %channel : index - %weight_row_byte_base = index.scale %expert_channel, %weight_row_bytes : index, offset -> offset - %q8_row_base0 = index.mul %bounded_token, %bounded_route_count : index - %q8_row = index.add %q8_row_base0, %route : index - %q8_row_byte_base = index.scale %q8_row, %q8_row_bytes : index, offset -> offset - %lane_acc = func.call @ggml_q6k_q8_1_x4_row_lane(%bounded_input_size, %weight_noalias, %weight_row_byte_base, %q8_noalias, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) - %route_weight = kernel.subgroup.broadcast %lane_route_weight from %route_i32 : f32, i32 - %weighted_lane = scalar.mulf %lane_acc, %route_weight : f32 - %next = scalar.addf %route_acc, %weighted_lane : f32 - scf.yield %next : f32 - } - scf.yield %sum : f32 - } else { - scf.yield %c0_f32 : f32 - } - %routed_sum = kernel.subgroup.reduce %routed_lane_sum : f32 - %writes_active_channel = scalar.andi %publish_output, %valid_channel : i1 - %writes_output = scalar.andi %writes_active_channel, %is_lane_zero : i1 - scf.if %writes_output { - %bounded_channel, %body_output_size = index.assume %channel, %bounded_output_size [lt(%channel, %bounded_output_size)] : index, index - %residual = view.load %output_view[%bounded_token, %bounded_channel] : view<[%body_token_count]x[%bounded_output_size]xf32> -> f32 - %result = scalar.addf %residual, %routed_sum : f32 - view.store %result, %output_view[%bounded_token, %bounded_channel] : f32, view<[%body_token_count]x[%bounded_output_size]xf32> - } - template.return -} - -kernel.def @qwen3_moe_routed_down_q6k_q8_1_x4(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index) { - %token_capacity = config.get @qwen3_moe.workload.token_capacity : index - %configured_output_size = config.get @qwen3_moe.routed_down.output_size : index - %c8 = index.constant 8 : index - %c7 = index.constant 7 : index - %c1 = index.constant 1 : index - %workgroup_size = index.constant 256 : index - %padded_output_size = index.add %configured_output_size, %c7 : index - %output_tiles = index.div %padded_output_size, %c8 : index - kernel.launch.config workgroups(%output_tiles, %token_capacity, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index -} launch(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) { - %token_capacity = config.get @qwen3_moe.workload.token_capacity : index - %configured_input_size0 = config.get @qwen3_moe.routed_down.input_size : index - %configured_route_count0 = config.get @qwen3_moe.routed_down.route_count : index - %configured_expert_count0 = config.get @qwen3_moe.routed_down.expert_count : index - %configured_output_size0 = config.get @qwen3_moe.routed_down.output_size : index - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index - %bounded_input_size, %configured_input_size = index.assume %input_size, %configured_input_size0 [range(%input_size, 256, 32768), mul(%input_size, 256), eq(%input_size, %configured_input_size0)] : index, index - %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 8), eq(%route_count, %configured_route_count0)] : index, index - %bounded_route_id_stride = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index - %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 1, 512), eq(%expert_count, %configured_expert_count0)] : index, index - %bounded_output_size, %configured_output_size = index.assume %output_size, %configured_output_size0 [range(%output_size, 1, 4096), eq(%output_size, %configured_output_size0)] : index, index - %token0 = kernel.workgroup.id : index - %c0 = index.constant 0 : index - %valid_token = index.cmp ult, %token0, %bounded_token_count : index - %safe_token0 = scf.select %valid_token, %token0, %c0 : index - %safe_token, %body_token_count = index.assume %safe_token0, %bounded_token_count [lt(%safe_token0, %bounded_token_count)] : index, index - template.apply<@qwen3_moe.routed_down.q6k_q8_1_x4.body>(%valid_token, %body_token_count, %safe_token, %configured_input_size, %configured_route_count, %bounded_route_id_stride, %configured_expert_count, %configured_output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : (i1, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer) - kernel.return -} - -// Decode-only route that also publishes the normalized Q8_1 x4 row consumed by -// the next projection boundary. -kernel.def target(@qwen3_moe_attention_prepare_gfx11_wave32) @qwen3_moe_routed_down_q6k_q8_1_x4_next_q8(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index) { - %configured_output_size = config.get @qwen3_moe.routed_down.output_size : index - %c8 = index.constant 8 : index - %c7 = index.constant 7 : index - %c1 = index.constant 1 : index - %workgroup_size = index.constant 256 : index - %padded_output_size = index.add %configured_output_size, %c7 : index - %output_tiles = index.div %padded_output_size, %c8 : index - kernel.launch.config workgroups(%output_tiles, %c1, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index -} launch(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer, %norm_weight: buffer, %completion_counter: buffer, %next_q8_output: buffer) { - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index - %configured_input_size0 = config.get @qwen3_moe.routed_down.input_size : index - %configured_route_count0 = config.get @qwen3_moe.routed_down.route_count : index - %configured_expert_count0 = config.get @qwen3_moe.routed_down.expert_count : index - %configured_output_size0 = config.get @qwen3_moe.routed_down.output_size : index - %bounded_input_size, %configured_input_size = index.assume %input_size, %configured_input_size0 [range(%input_size, 256, 32768), mul(%input_size, 256), eq(%input_size, %configured_input_size0)] : index, index - %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 8), eq(%route_count, %configured_route_count0)] : index, index - %bounded_route_id_stride = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index - %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 1, 512), eq(%expert_count, %configured_expert_count0)] : index, index - %bounded_output_size, %configured_output_size = index.assume %output_size, %configured_output_size0 [range(%output_size, 128, 4096), mul(%output_size, 128), eq(%output_size, %configured_output_size0)] : index, index - %c8 = index.constant 8 : index - %publishes_output = scalar.constant true : i1 - %body_token = index.constant 0 : index - template.apply<@qwen3_moe.routed_down.q6k_q8_1_x4.body>(%publishes_output, %bounded_token_count, %body_token, %configured_input_size, %configured_route_count, %bounded_route_id_stride, %configured_expert_count, %configured_output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : (i1, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer) - template.apply<@qwen3_moe.routed_down.next_q8_completion>(%c8, %c8, %bounded_token_count, %configured_output_size, %output, %norm_weight, %completion_counter, %next_q8_output) : (index, index, index, index, buffer, buffer, buffer, buffer) - kernel.return -} - -// The production hidden width requires all 256 output tiles to arrive before -// normalization. Compare both residual and packed Q8 publication with the -// ordinary two-dispatch composition, then reuse the same completion word. -check.case public @qwen3_moe_routed_down_q6k_q8_1_x4_next_q8_differential_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(2) : index - %route_id_stride = check.literal value(4) : index - %expert_count = check.literal value(2) : index - %output_size = check.literal value(2048) : index - %routed_row_count = check.literal value(2) : index - %routed_input = check.generate.fill value(0.00390625) : tensor<2x768xf32> - %q8_input = check.generate.fill value(0) : tensor<2x864xi8> - %route_ids = check.generate.iota offset(0) step(1) period(2) : tensor<1x4xi32> - %route_weights = check.generate.fill value(0.5) : tensor<1x2xf32> - %weight = check.generate.fill value(-86) : tensor<2x2048x3x210xi8> - %norm_weight = check.generate.iota offset(-1.0) step(0.0009765625) : tensor<2048xf32> - %expected_output = check.generate.iota offset(-0.5) step(0.00048828125) period(2048) : tensor<1x2048xf32> - %actual_output0 = check.generate.iota offset(-0.5) step(0.00048828125) period(2048) : tensor<1x2048xf32> - %actual_output1 = check.generate.iota offset(-0.5) step(0.00048828125) period(2048) : tensor<1x2048xf32> - %expected_q8 = check.generate.fill value(0) : tensor<2304xi8> - %actual_q8_0 = check.generate.fill value(1) : tensor<2304xi8> - %actual_q8_1 = check.generate.fill value(1) : tensor<2304xi8> - %completion_counter = check.generate.fill value(0) : tensor<1xi32> - %expected_counter = check.generate.fill value(0) : tensor<1xi32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %input_size](%routed_row_count, %input_size, %routed_input, %q8_input) : [index, index](index, index, tensor<2x768xf32>, tensor<2x864xi8>) - kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %expected_output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x864xi8>, tensor<1x4xi32>, tensor<1x2xf32>, tensor<2x2048x3x210xi8>, tensor<1x2048xf32>) - kernel.launch @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4[%token_count](%token_count, %expected_output, %norm_weight, %expected_q8) : [index](index, tensor<1x2048xf32>, tensor<2048xf32>, tensor<2304xi8>) - kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4_next_q8[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %actual_output0, %norm_weight, %completion_counter, %actual_q8_0) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x864xi8>, tensor<1x4xi32>, tensor<1x2xf32>, tensor<2x2048x3x210xi8>, tensor<1x2048xf32>, tensor<2048xf32>, tensor<1xi32>, tensor<2304xi8>) - kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4_next_q8[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %actual_output1, %norm_weight, %completion_counter, %actual_q8_1) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x864xi8>, tensor<1x4xi32>, tensor<1x2xf32>, tensor<2x2048x3x210xi8>, tensor<1x2048xf32>, tensor<2048xf32>, tensor<1xi32>, tensor<2304xi8>) - check.expect.close actual(%actual_output0) expected(%expected_output) atol(0.0) rtol(0.0) nan(same) : tensor<1x2048xf32> - check.expect.close actual(%actual_output1) expected(%expected_output) atol(0.0) rtol(0.0) nan(same) : tensor<1x2048xf32> - check.expect.equal actual(%actual_q8_0) expected(%expected_q8) : tensor<2304xi8> - check.expect.equal actual(%actual_q8_1) expected(%expected_q8) : tensor<2304xi8> - check.expect.equal actual(%completion_counter) expected(%expected_counter) : tensor<1xi32> - check.return -} - -check.case public @qwen3_moe_routed_down_q6k_q8_1_x4_next_q8_benchmark_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(8) : index - %route_id_stride = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(2048) : index - %q8_input = check.generate.fill value(0) : tensor<1x8x864xi8> - %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<1x8xi32> - %route_weights = check.generate.fill value(0.125) : tensor<1x8xf32> - %weight = check.generate.fill value(0) : tensor<128x2048x3x210xi8> - %output = check.generate.fill value(1.0) : tensor<1x2048xf32> - %norm_weight = check.generate.fill value(1.0) : tensor<2048xf32> - %completion_counter = check.generate.fill value(0) : tensor<1xi32> - %next_q8_output = check.generate.fill value(0) : tensor<2304xi8> - kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4_next_q8[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output, %norm_weight, %completion_counter, %next_q8_output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<1x8x864xi8>, tensor<1x8xi32>, tensor<1x8xf32>, tensor<128x2048x3x210xi8>, tensor<1x2048xf32>, tensor<2048xf32>, tensor<1xi32>, tensor<2304xi8>) - check.return -} - -check.case public @qwen3_moe_routed_down_q6k_q8_1_x4_next_q8_composed_benchmark_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(8) : index - %route_id_stride = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(2048) : index - %q8_input = check.generate.fill value(0) : tensor<1x8x864xi8> - %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<1x8xi32> - %route_weights = check.generate.fill value(0.125) : tensor<1x8xf32> - %weight = check.generate.fill value(0) : tensor<128x2048x3x210xi8> - %output = check.generate.fill value(1.0) : tensor<1x2048xf32> - %norm_weight = check.generate.fill value(1.0) : tensor<2048xf32> - %next_q8_output = check.generate.fill value(0) : tensor<2304xi8> - kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<1x8x864xi8>, tensor<1x8xi32>, tensor<1x8xf32>, tensor<128x2048x3x210xi8>, tensor<1x2048xf32>) - kernel.launch @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4[%token_count](%token_count, %output, %norm_weight, %next_q8_output) : [index](index, tensor<1x2048xf32>, tensor<2048xf32>, tensor<2304xi8>) - check.return -} - -// Shared direct-F32 decode schedule. Every subgroup owns one scale frame and -// one output channel. Providers specialize how many Q6_K blocks a subgroup -// covers per pass from the target subgroup width. -template.def<@qwen3_moe.routed_down.q6k_f32.body> device @qwen3_moe_routed_down_q6k_f32_body(%publish_output: i1, %subgroup_count: index, %block_step: index, %scale_stage_bytes: offset, %token_count: index, %token: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) { - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index - %bounded_token, %body_token_count = index.assume %token, %bounded_token_count [lt(%token, %bounded_token_count)] : index, index - %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index - %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index - %bounded_route_id_stride = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index - %bounded_expert_count = index.assume %expert_count [range(%expert_count, 1, 512)] : index - %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144)] : index - %channel_tile = kernel.workgroup.id : index - %subgroup0 = kernel.subgroup.id : index - %lane0 = kernel.subgroup.lane.id : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c16 = index.constant 16 : index - %c210 = index.constant 210 : index - %c256 = index.constant 256 : index - %c0_i32 = scalar.constant 0 : i32 - %c0_f32 = scalar.constant 0.0 : f32 - %c0_offset = index.constant 0 : offset - %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index - %lane = index.assume %lane0 [range(%lane0, 0, 63)] : index - %channel_tile_base = index.mul %channel_tile, %subgroup_count : index - %channel = index.add %channel_tile_base, %subgroup : index - %valid_channel = index.cmp ult, %channel, %bounded_output_size : index - %safe_channel = scf.select %valid_channel, %channel, %c0 : index - %lane_i32 = index.cast %lane : index to i32 - %is_lane_zero = scalar.cmpi eq, %lane_i32, %c0_i32 : i32 - %assignment_count = index.mul %body_token_count, %bounded_route_count : index - %q6_block_count = index.div %bounded_input_size, %c256 : index - %weight_row_byte_count = index.mul %q6_block_count, %c210 : index - %cohort = index.div %lane, %c16 : index - %input_noalias, %route_id_noalias, %route_weight_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %input, %route_ids, %route_weights, %weight, %output : buffer, buffer, buffer, buffer, buffer - %route_id_view = buffer.view %route_id_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_route_id_stride]xi32> - %route_weight_view = buffer.view %route_weight_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_route_count]xf32> - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_output_size]xf32> - %scale_stage = buffer.alloca align(16) %scale_stage_bytes : buffer - %routed_sum0 = scf.if %publish_output -> (f32) { - %active_sum = scf.for %route = [%c0 to %bounded_route_count step %c1](%route_acc = %c0_f32 : f32) -> (f32) unroll { - // Route metadata is uniform across all channels. Keeping the access in - // the unrolled route body exposes that uniformity directly to target - // lowering instead of materializing a dynamic subgroup broadcast. - %bounded_route, %body_route_count = index.assume %route, %bounded_route_count [lt(%route, %bounded_route_count)] : index, index - %route_for_stride, %body_route_id_stride = index.assume %bounded_route, %bounded_route_id_stride [lt(%bounded_route, %bounded_route_id_stride)] : index, index - %expert_i32 = view.load %route_id_view[%bounded_token, %route_for_stride] : view<[%body_token_count]x[%bounded_route_id_stride]xi32> -> i32 - %route_weight = view.load %route_weight_view[%bounded_token, %bounded_route] : view<[%body_token_count]x[%bounded_route_count]xf32> -> f32 - %expert0 = index.cast %expert_i32 : i32 to index - %expert1 = index.assume %expert0 [range(%expert0, 0, 511)] : index - %expert, %weight_expert_count = index.assume %expert1, %bounded_expert_count [lt(%expert1, %bounded_expert_count)] : index, index - %expert_output_base = index.mul %expert, %bounded_output_size : index - %expert_channel = index.add %expert_output_base, %safe_channel : index - %input_row_base = index.mul %bounded_token, %bounded_route_count : index - %input_row = index.add %input_row_base, %bounded_route : index - %lane_sum = scf.for %block_base = [%c0 to %q6_block_count step %block_step](%block_acc = %c0_f32 : f32) -> (f32) { - %block = index.add %block_base, %cohort : index - func.call @ggml_q6k_stage_f32_scales(%bounded_input_size, %expert_channel, %block, %subgroup_count, %subgroup, %lane, %weight_noalias, %scale_stage) : (index, index, index, index, index, index, buffer, buffer) - %input0, %input1, %input2, %input3 = func.call @ggml_q6k_load_f32_block(%assignment_count, %bounded_input_size, %input_row, %block, %lane, %input_noalias) : (index, index, index, index, index, buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) - %contribution = func.call @ggml_q6k_f32_block_row(%bounded_input_size, %expert_channel, %block, %subgroup_count, %subgroup, %lane, %weight_noalias, %scale_stage, %input0, %input1, %input2, %input3) : (index, index, index, index, index, index, buffer, buffer, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (f32) - %next = scalar.addf %block_acc, %contribution : f32 - scf.yield %next : f32 - } - %route_sum = kernel.subgroup.reduce %lane_sum : f32 - %weighted = scalar.mulf %route_sum, %route_weight : f32 - %next = scalar.addf %route_acc, %weighted : f32 - scf.yield %next : f32 - } - scf.yield %active_sum : f32 - } else { - scf.yield %c0_f32 : f32 - } - %routed_sum = scf.select %valid_channel, %routed_sum0, %c0_f32 : f32 - %writes_active_channel = scalar.andi %publish_output, %valid_channel : i1 - %writes_output = scalar.andi %writes_active_channel, %is_lane_zero : i1 - scf.if %writes_output { - %bounded_channel, %body_output_size = index.assume %channel, %bounded_output_size [lt(%channel, %bounded_output_size)] : index, index - %residual = view.load %output_view[%bounded_token, %bounded_channel] : view<[%body_token_count]x[%bounded_output_size]xf32> -> f32 - %result = scalar.addf %residual, %routed_sum : f32 - view.store %result, %output_view[%bounded_token, %bounded_channel] : f32, view<[%body_token_count]x[%bounded_output_size]xf32> - } - template.return -} - -// Decode schedule matching the Vulkan oracle's two adjacent output rows per -// wave. Four wave64 subgroups publish eight channels per workgroup. Keeping the -// route loop rolled prevents its eight iterations from multiplying the two-row -// contraction's register pressure and instruction footprint. -template.def<@qwen3_moe.routed_down.q6k_f32.pair_body> device @qwen3_moe_routed_down_q6k_f32_pair_body(%token_count: index, %token: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) { - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index - %bounded_token, %body_token_count = index.assume %token, %bounded_token_count [lt(%token, %bounded_token_count)] : index, index - %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index - %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index - %bounded_route_id_stride = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index - %bounded_expert_count = index.assume %expert_count [range(%expert_count, 1, 512)] : index - %bounded_output_size = index.assume %output_size [range(%output_size, 2, 4096)] : index - %pair_tile = kernel.workgroup.id : index - %subgroup0 = kernel.subgroup.id : index - %lane0 = kernel.subgroup.lane.id : index - %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 - %c256 = index.constant 256 : index - %c0_f32 = scalar.constant 0.0 : f32 - %c0_offset = index.constant 0 : offset - %scale_stage_bytes = index.constant 2048 : offset - %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 3)] : index - %lane = index.assume %lane0 [range(%lane0, 0, 63)] : index - %pair_tile_base = index.mul %pair_tile, %c4 : index - %pair = index.add %pair_tile_base, %subgroup : index - %row00 = index.mul %pair, %c2 : index - %row0, %body_output_size = index.assume %row00, %bounded_output_size [lt(%row00, %bounded_output_size)] : index, index - %row1 = index.add %row0, %c1 : index - %row1_valid = index.cmp ult, %row1, %body_output_size : index - %safe_row1 = scf.select %row1_valid, %row1, %c0 : index - %assignment_count = index.mul %body_token_count, %bounded_route_count : index - %q6_block_count = index.div %bounded_input_size, %c256 : index - %cohort = index.div %lane, %c16 : index - %scale_frame0 = index.mul %subgroup, %c2 : index - %scale_frame1 = index.add %scale_frame0, %c1 : index - %input_noalias, %route_id_noalias, %route_weight_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %input, %route_ids, %route_weights, %weight, %output : buffer, buffer, buffer, buffer, buffer - %route_id_view = buffer.view %route_id_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_route_id_stride]xi32> - %route_weight_view = buffer.view %route_weight_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_route_count]xf32> - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%body_output_size]xf32> - %scale_stage = buffer.alloca align(16) %scale_stage_bytes : buffer - %routed_sum0, %routed_sum1 = scf.for %route = [%c0 to %bounded_route_count step %c1](%route_acc0 = %c0_f32 : f32, %route_acc1 = %c0_f32 : f32) -> (f32, f32) { - %bounded_route, %body_route_count = index.assume %route, %bounded_route_count [lt(%route, %bounded_route_count)] : index, index - %route_for_stride, %body_route_id_stride = index.assume %bounded_route, %bounded_route_id_stride [lt(%bounded_route, %bounded_route_id_stride)] : index, index - %expert_i32 = view.load %route_id_view[%bounded_token, %route_for_stride] : view<[%body_token_count]x[%bounded_route_id_stride]xi32> -> i32 - %route_weight = view.load %route_weight_view[%bounded_token, %bounded_route] : view<[%body_token_count]x[%bounded_route_count]xf32> -> f32 - %expert0 = index.cast %expert_i32 : i32 to index - %expert1 = index.assume %expert0 [range(%expert0, 0, 511)] : index - %expert, %body_expert_count = index.assume %expert1, %bounded_expert_count [lt(%expert1, %bounded_expert_count)] : index, index - %expert_output_base = index.mul %expert, %body_output_size : index - %expert_row0 = index.add %expert_output_base, %row0 : index - %expert_row1 = index.add %expert_output_base, %safe_row1 : index - %input_row_base = index.mul %bounded_token, %body_route_count : index - %input_row = index.add %input_row_base, %bounded_route : index - %lane_sum0, %lane_sum1 = scf.for %block_base = [%c0 to %q6_block_count step %c4](%block_acc0 = %c0_f32 : f32, %block_acc1 = %c0_f32 : f32) -> (f32, f32) { - %block = index.add %block_base, %cohort : index - func.call @ggml_q6k_stage_f32_scales(%bounded_input_size, %expert_row0, %block, %c8, %scale_frame0, %lane, %weight_noalias, %scale_stage) : (index, index, index, index, index, index, buffer, buffer) - %input00, %input01, %input02, %input03 = func.call @ggml_q6k_load_f32_block(%assignment_count, %bounded_input_size, %input_row, %block, %lane, %input_noalias) : (index, index, index, index, index, buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) - %contribution0 = func.call @ggml_q6k_f32_block_row(%bounded_input_size, %expert_row0, %block, %c8, %scale_frame0, %lane, %weight_noalias, %scale_stage, %input00, %input01, %input02, %input03) : (index, index, index, index, index, index, buffer, buffer, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (f32) - %contribution1 = scf.if %row1_valid -> (f32) { - func.call @ggml_q6k_stage_f32_scales(%bounded_input_size, %expert_row1, %block, %c8, %scale_frame1, %lane, %weight_noalias, %scale_stage) : (index, index, index, index, index, index, buffer, buffer) - // Reloading the shared activations lets the first row's vectors die - // before the second scale stage. Carrying them across that stage - // increases register pressure and is slower on gfx1151. - %input10, %input11, %input12, %input13 = func.call @ggml_q6k_load_f32_block(%assignment_count, %bounded_input_size, %input_row, %block, %lane, %input_noalias) : (index, index, index, index, index, buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) - %row1_contribution = func.call @ggml_q6k_f32_block_row(%bounded_input_size, %expert_row1, %block, %c8, %scale_frame1, %lane, %weight_noalias, %scale_stage, %input10, %input11, %input12, %input13) : (index, index, index, index, index, index, buffer, buffer, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (f32) - scf.yield %row1_contribution : f32 - } else { - scf.yield %c0_f32 : f32 - } - %next0 = scalar.addf %block_acc0, %contribution0 : f32 - %next1 = scalar.addf %block_acc1, %contribution1 : f32 - scf.yield %next0, %next1 : f32, f32 - } - %route_sum0 = kernel.subgroup.reduce %lane_sum0 : f32 - %route_sum1 = kernel.subgroup.reduce %lane_sum1 : f32 - %weighted0 = scalar.mulf %route_sum0, %route_weight : f32 - %weighted1 = scalar.mulf %route_sum1, %route_weight : f32 - %next0 = scalar.addf %route_acc0, %weighted0 : f32 - %next1 = scalar.addf %route_acc1, %weighted1 : f32 - scf.yield %next0, %next1 : f32, f32 - } - %is_lane_zero = index.cmp eq, %lane, %c0 : index - scf.if %is_lane_zero { - %residual0 = view.load %output_view[%bounded_token, %row0] : view<[%body_token_count]x[%body_output_size]xf32> -> f32 - %result0 = scalar.addf %residual0, %routed_sum0 : f32 - view.store %result0, %output_view[%bounded_token, %row0] : f32, view<[%body_token_count]x[%body_output_size]xf32> - scf.if %row1_valid { - %residual1 = view.load %output_view[%bounded_token, %row1] : view<[%body_token_count]x[%body_output_size]xf32> -> f32 - %result1 = scalar.addf %residual1, %routed_sum1 : f32 - view.store %result1, %output_view[%bounded_token, %row1] : f32, view<[%body_token_count]x[%body_output_size]xf32> - } - } - template.return -} - -// Wave64 matches llama.cpp's Vulkan subgroup width and contracts four Q6_K -// blocks per subgroup pass. -kernel.def target(@qwen3_moe_routed_down_q6k_gfx11_wave64) @qwen3_moe_routed_down_q6k_f32_wave64(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index) { - %token_capacity = config.get @qwen3_moe.workload.token_capacity : index - %configured_output_size = config.get @qwen3_moe.routed_down.output_size : index - %c1 = index.constant 1 : index - %c3 = index.constant 3 : index - %c4 = index.constant 4 : index - %c256 = index.constant 256 : index - %padded_output_size = index.add %configured_output_size, %c3 : index - %output_tiles = index.div %padded_output_size, %c4 : index - kernel.launch.config workgroups(%output_tiles, %token_capacity, %c1) workgroup_size(%c256, %c1, %c1) : index -} launch(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) { - %token_capacity = config.get @qwen3_moe.workload.token_capacity : index - %configured_input_size0 = config.get @qwen3_moe.routed_down.input_size : index - %configured_route_count0 = config.get @qwen3_moe.routed_down.route_count : index - %configured_expert_count0 = config.get @qwen3_moe.routed_down.expert_count : index - %configured_output_size0 = config.get @qwen3_moe.routed_down.output_size : index - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index - %bounded_input_size, %configured_input_size = index.assume %input_size, %configured_input_size0 [range(%input_size, 256, 32768), mul(%input_size, 256), eq(%input_size, %configured_input_size0)] : index, index - %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 8), eq(%route_count, %configured_route_count0)] : index, index - %bounded_route_id_stride = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index - %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 1, 512), eq(%expert_count, %configured_expert_count0)] : index, index - %bounded_output_size, %configured_output_size = index.assume %output_size, %configured_output_size0 [range(%output_size, 1, 4096), eq(%output_size, %configured_output_size0)] : index, index - %token0 = kernel.workgroup.id : index - %c0 = index.constant 0 : index - %valid_token = index.cmp ult, %token0, %bounded_token_count : index - %safe_token0 = scf.select %valid_token, %token0, %c0 : index - %safe_token, %body_token_count = index.assume %safe_token0, %bounded_token_count [lt(%safe_token0, %bounded_token_count)] : index, index - %c4 = index.constant 4 : index - %scale_stage_bytes = index.constant 1024 : offset - template.apply<@qwen3_moe.routed_down.q6k_f32.body>(%valid_token, %c4, %c4, %scale_stage_bytes, %body_token_count, %safe_token, %configured_input_size, %configured_route_count, %bounded_route_id_stride, %configured_expert_count, %configured_output_size, %input, %route_ids, %route_weights, %weight, %output) : (i1, index, index, offset, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer) - kernel.return -} - -// Two selected experts exercise the noncompact route-ID stride, normalized -// weighted reduction, in-place residual update, K=768, and the output tile -// tail. Uniform 0xaa Q6_K bytes decode to a deterministic nonzero row with -// negative signed group scales. -check.case public @qwen3_moe_routed_down_q6k_q8_1_x4_nonzero_residual_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(2) : index - %route_id_stride = check.literal value(4) : index - %expert_count = check.literal value(2) : index - %output_size = check.literal value(9) : index - %routed_input = check.generate.fill value(0.00390625) : tensor<2x768xf32> - %q8_input = check.generate.fill value(0) : tensor<2x864xi8> - %route_ids = check.generate.iota offset(0) step(1) period(2) : tensor<1x4xi32> - %route_weights = check.generate.fill value(0.5) : tensor<1x2xf32> - %weight = check.generate.fill value(-86) : tensor<2x9x3x210xi8> - %output = check.generate.fill value(1.0) : tensor<1x9xf32> - %expected = check.generate.fill value(135.31431579589844) : tensor<1x9xf32> - %routed_row_count = check.literal value(2) : index - kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %input_size](%routed_row_count, %input_size, %routed_input, %q8_input) : [index, index](index, index, tensor<2x768xf32>, tensor<2x864xi8>) - kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x864xi8>, tensor<1x4xi32>, tensor<1x2xf32>, tensor<2x9x3x210xi8>, tensor<1x9xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.25) rtol(9.9999999999999995e-07) nan(same) : tensor<1x9xf32> - check.return -} - -// Two tokens use different route-weight sums so token, route, packed-row, and -// output addressing cannot collapse to the single-token case. -check.case public @qwen3_moe_routed_down_q6k_q8_1_x4_nonzero_prefill_case { - %token_count = check.literal value(2) : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(2) : index - %route_id_stride = check.literal value(4) : index - %expert_count = check.literal value(2) : index - %output_size = check.literal value(1) : index - %routed_row_count = check.literal value(4) : index - %routed_input = check.generate.fill value(0.00390625) : tensor<2x2x768xf32> - %q8_input = check.generate.fill value(0) : tensor<2x2x864xi8> - %route_ids = check.generate.iota offset(0) step(1) period(2) : tensor<2x4xi32> - %route_weights = check.generate.iota offset(0.25) step(0.25) period(4) : tensor<2x2xf32> - %weight = check.generate.fill value(-86) : tensor<2x1x3x210xi8> - %output = check.generate.fill value(1.0) : tensor<2x1xf32> - %expected = check.generate.iota offset(101.73573684692383) step(134.31431579589844) : tensor<2x1xf32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %input_size](%routed_row_count, %input_size, %routed_input, %q8_input) : [index, index](index, index, tensor<2x2x768xf32>, tensor<2x2x864xi8>) - kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x2x864xi8>, tensor<2x4xi32>, tensor<2x2xf32>, tensor<2x1x3x210xi8>, tensor<2x1xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.25) rtol(9.9999999999999995e-07) nan(same) : tensor<2x1xf32> - check.return -} - -// The exact-representable routed inputs make direct-F32 and Q8_1 contraction -// comparable while preserving noncompact route IDs and an output-channel tail. -check.case public @qwen3_moe_routed_down_q6k_f32_wave64_differential_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(2) : index - %route_id_stride = check.literal value(4) : index - %expert_count = check.literal value(2) : index - %output_size = check.literal value(9) : index - %routed_row_count = check.literal value(2) : index - %routed_input = check.generate.fill value(0.00390625) : tensor<2x768xf32> - %q8_input = check.generate.fill value(0) : tensor<2x864xi8> - %route_ids = check.generate.iota offset(0) step(1) period(2) : tensor<1x4xi32> - %route_weights = check.generate.fill value(0.5) : tensor<1x2xf32> - %weight = check.generate.fill value(-86) : tensor<2x9x3x210xi8> - %expected = check.generate.fill value(1.0) : tensor<1x9xf32> - %actual_wave64 = check.generate.fill value(1.0) : tensor<1x9xf32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %input_size](%routed_row_count, %input_size, %routed_input, %q8_input) : [index, index](index, index, tensor<2x768xf32>, tensor<2x864xi8>) - kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %expected) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x864xi8>, tensor<1x4xi32>, tensor<1x2xf32>, tensor<2x9x3x210xi8>, tensor<1x9xf32>) - kernel.launch @qwen3_moe_routed_down_q6k_f32_wave64[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %routed_input, %route_ids, %route_weights, %weight, %actual_wave64) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x768xf32>, tensor<1x4xi32>, tensor<1x2xf32>, tensor<2x9x3x210xi8>, tensor<1x9xf32>) - check.expect.close actual(%actual_wave64) expected(%expected) atol(0.25) rtol(0.01) nan(same) : tensor<1x9xf32> - check.return -} - -// Wave64 reference for the row producer used by the fused completion path. -kernel.def target(@qwen3_moe_routed_down_q6k_gfx11_wave64) @qwen3_moe_rmsnorm_quantize_q8_1_x4_wave64_production_check() { - %c1 = index.constant 1 : index - %c256 = index.constant 256 : index - kernel.launch.config workgroups(%c1, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index -} launch(%input: buffer, %norm_weight: buffer, %q8_output: buffer) { - %publish_normalized = scalar.constant false : i1 - %reduction_subgroup_count = index.constant 4 : index - %token_count = index.constant 1 : index - %token = index.constant 0 : index - template.apply<@qwen3_moe.rmsnorm_quantize_q8_1_x4.body>(%publish_normalized, %reduction_subgroup_count, %token_count, %token, %input, %norm_weight, %q8_output, %q8_output) : (i1, index, index, index, buffer, buffer, buffer, buffer) - kernel.return -} - -// The production hidden width requires every direct-F32 output tile to arrive -// before normalization. Nonzero Q6_K data compares the paired-row producer -// against the ordinary wave64 producer and standalone normalization, then -// reuses the completion word so stale completion state is observable. -check.case public @qwen3_moe_routed_down_q6k_f32_wave64_next_q8_differential_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(8) : index - %route_id_stride = check.literal value(8) : index - %expert_count = check.literal value(8) : index - %output_size = check.literal value(2048) : index - %input = check.generate.iota offset(-0.5) step(0.0001627604166666667) period(6144) : tensor<8x768xf32> - %route_ids = check.generate.iota offset(0) step(1) period(8) : tensor<8xi32> - %route_weights = check.generate.iota offset(0.02777777777777778) step(0.02777777777777778) : tensor<8xf32> - %weight = check.generate.fill value(-86) : tensor<8x2048x3x210xi8> - %norm_weight = check.generate.iota offset(-1.0) step(0.0009765625) : tensor<2048xf32> - %expected_output = check.generate.iota offset(-0.5) step(0.00048828125) period(2048) : tensor<2048xf32> - %actual_output0 = check.generate.iota offset(-0.5) step(0.00048828125) period(2048) : tensor<2048xf32> - %actual_output1 = check.generate.iota offset(-0.5) step(0.00048828125) period(2048) : tensor<2048xf32> - %expected_q8 = check.generate.fill value(0) : tensor<2304xi8> - %actual_q8_0 = check.generate.fill value(1) : tensor<2304xi8> - %actual_q8_1 = check.generate.fill value(1) : tensor<2304xi8> - %completion_counter = check.generate.fill value(0) : tensor<1xi32> - %expected_counter = check.generate.fill value(0) : tensor<1xi32> - kernel.launch @qwen3_moe_routed_down_q6k_f32_wave64[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %input, %route_ids, %route_weights, %weight, %expected_output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<8x768xf32>, tensor<8xi32>, tensor<8xf32>, tensor<8x2048x3x210xi8>, tensor<2048xf32>) - kernel.launch @qwen3_moe_rmsnorm_quantize_q8_1_x4_wave64_production_check(%expected_output, %norm_weight, %expected_q8) : (tensor<2048xf32>, tensor<2048xf32>, tensor<2304xi8>) - kernel.launch @ggml_mul_mat_id_vector_weighted_wave64_publish_q8[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %input, %route_ids, %route_weights, %weight, %actual_output0, %norm_weight, %completion_counter, %actual_q8_0) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<8x768xf32>, tensor<8xi32>, tensor<8xf32>, tensor<8x2048x3x210xi8>, tensor<2048xf32>, tensor<2048xf32>, tensor<1xi32>, tensor<2304xi8>) - kernel.launch @ggml_mul_mat_id_vector_weighted_wave64_publish_q8[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %input, %route_ids, %route_weights, %weight, %actual_output1, %norm_weight, %completion_counter, %actual_q8_1) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<8x768xf32>, tensor<8xi32>, tensor<8xf32>, tensor<8x2048x3x210xi8>, tensor<2048xf32>, tensor<2048xf32>, tensor<1xi32>, tensor<2304xi8>) - check.expect.close actual(%actual_output0) expected(%expected_output) atol(0.0) rtol(0.0) nan(same) : tensor<2048xf32> - check.expect.close actual(%actual_output1) expected(%expected_output) atol(0.0) rtol(0.0) nan(same) : tensor<2048xf32> - check.expect.equal actual(%actual_q8_0) expected(%expected_q8) : tensor<2304xi8> - check.expect.equal actual(%actual_q8_1) expected(%expected_q8) : tensor<2304xi8> - check.expect.equal actual(%completion_counter) expected(%expected_counter) : tensor<1xi32> - check.return -} - -check.case public @qwen3_moe_routed_down_q6k_q8_1_x4_benchmark_case { - %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512]) name("token_count") : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(8) : index - %route_id_stride = check.literal value(128) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(2048) : index - %q8_input = check.generate.fill value(0) : tensor<[%token_count]x8x864xi8> - %route_ids = check.generate.iota offset(0) step(1) period(127) : tensor<[%token_count]x128xi32> - %route_weights = check.generate.fill value(0.125) : tensor<[%token_count]x8xf32> - %weight = check.generate.fill value(0) : tensor<128x2048x3x210xi8> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - %expected = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<[%token_count]x8x864xi8>, tensor<[%token_count]x128xi32>, tensor<[%token_count]x8xf32>, tensor<128x2048x3x210xi8>, tensor<[%token_count]x2048xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> - check.return -} - -check.case public @qwen3_moe_routed_down_q6k_q8_1_x4_pipeline_benchmark_case { - %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512]) name("token_count") : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(8) : index - %route_id_stride = check.literal value(128) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(2048) : index - // Each 768-element route contains six complete Q8_1 x4 groups, so packing - // the eight contiguous routes as one 6,144-element token row preserves the - // exact per-route physical layout without a derived testbench scalar. - %routed_input_size = check.literal value(6144) : index - %routed_input = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf32> - %q8_input = check.generate.fill value(1) : tensor<[%token_count]x8x864xi8> - %route_ids = check.generate.iota offset(0) step(1) period(127) : tensor<[%token_count]x128xi32> - %route_weights = check.generate.fill value(0.125) : tensor<[%token_count]x8xf32> - %weight = check.generate.fill value(0) : tensor<128x2048x3x210xi8> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - %expected = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %routed_input_size](%token_count, %routed_input_size, %routed_input, %q8_input) : [index, index](index, index, tensor<[%token_count]x8x768xf32>, tensor<[%token_count]x8x864xi8>) - kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<[%token_count]x8x864xi8>, tensor<[%token_count]x128xi32>, tensor<[%token_count]x8xf32>, tensor<128x2048x3x210xi8>, tensor<[%token_count]x2048xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> - check.return -} - -check.case public @qwen3_moe_routed_down_q6k_f32_wave64_benchmark_case { - %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512]) name("token_count") : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(8) : index - %route_id_stride = check.literal value(128) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(2048) : index - %routed_input = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf32> - %route_ids = check.generate.iota offset(0) step(1) period(127) : tensor<[%token_count]x128xi32> - %route_weights = check.generate.fill value(0.125) : tensor<[%token_count]x8xf32> - %weight = check.generate.fill value(0) : tensor<128x2048x3x210xi8> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - %expected = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - kernel.launch @qwen3_moe_routed_down_q6k_f32_wave64[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %routed_input, %route_ids, %route_weights, %weight, %output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<[%token_count]x8x768xf32>, tensor<[%token_count]x128xi32>, tensor<[%token_count]x8xf32>, tensor<128x2048x3x210xi8>, tensor<[%token_count]x2048xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> - check.return -} - -check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_nonzero_residual_case> @qwen3_moe_routed_down_q6k_q8_1_x4_small - -check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_decode {token_count = 1} - -check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_small_batch_2 {token_count = 2} - -check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_small_batch_4 {token_count = 4} - -check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_small_batch_8 {token_count = 8} - -check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_small_batch_16 {token_count = 16} - -check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_prefill_32 {token_count = 32} - -check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_prefill_128 {token_count = 128} - -check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_prefill_512 {token_count = 512} - -check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_pipeline_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_pipeline_decode {token_count = 1} - -check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_pipeline_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_pipeline_prefill_32 {token_count = 32} - -check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_pipeline_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_pipeline_prefill_128 {token_count = 128} - -check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_pipeline_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_pipeline_prefill_512 {token_count = 512} - -check.benchmark<@qwen3_moe_routed_down_q6k_f32_wave64_benchmark_case> @qwen3_moe_routed_down_q6k_f32_wave64_decode {token_count = 1} - -check.benchmark<@qwen3_moe_routed_down_q6k_f32_wave64_benchmark_case> @qwen3_moe_routed_down_q6k_f32_wave64_small_batch_2 {token_count = 2} - -check.benchmark<@qwen3_moe_routed_down_q6k_f32_wave64_benchmark_case> @qwen3_moe_routed_down_q6k_f32_wave64_small_batch_4 {token_count = 4} - -check.benchmark<@qwen3_moe_routed_down_q6k_f32_wave64_benchmark_case> @qwen3_moe_routed_down_q6k_f32_wave64_small_batch_8 {token_count = 8} - -check.benchmark<@qwen3_moe_routed_down_q6k_f32_wave64_benchmark_case> @qwen3_moe_routed_down_q6k_f32_wave64_small_batch_16 {token_count = 16} - -check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_next_q8_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_next_q8_decode - -check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_next_q8_composed_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_next_q8_composed_decode diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_quantized_f16_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_quantized_f16_wmma.loom index 9a39d2ce8534..4b1a879bde48 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_quantized_f16_wmma.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_quantized_f16_wmma.loom @@ -1,555 +1,16 @@ -// Expert-grouped Q4_K and Q6_K down projections for gfx11 prefill. +// Weighted reduction of compact FP16 routed projections into an FP32 residual. // -// Each two-wave workgroup decodes 64 output channels for 32 routed rows of one -// expert. Raw GGUF weights and FP16 routed SwiGLU rows are staged as FP16 WMMA -// operands. Results remain FP16 and are scattered into compact -// [token, route, hidden] order without collisions. A following reduction -// widens each route, applies its normalized weight, and accumulates the -// residual in FP32. -// -// Top-k routing selects an expert at most once per token. The expert table -// therefore contains at most token_count assignments per expert, while each -// assignment ordinal still identifies the compact [token, route] activation -// and route-weight row. -template.decl @qwen3_moe.routed_down.quantized_prefill.body(%weight_format: index, %token_count: index, %input: buffer, %expert_table: buffer, %weight: buffer, %output: buffer) +// One workitem owns four adjacent output channels. All top-k routes are +// accumulated in FP32 before the residual is loaded and published once. -amdgpu.target @qwen3_moe_routed_down_gfx11_wave64 {subgroup_size = 64} +template.decl @ggml.routing.weighted_reduce_f16_f32(%token_count: index, %route_weights: buffer, %routed_output: buffer, %residual_input: buffer, %output: buffer) -config.decl @qwen3_moe.routed_down.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] +amdgpu.target @qwen3_moe_routed_down_gfx11_wave64 {subgroup_size = 64} config.decl @qwen3_moe.routed_down.route_count : %value: index where [range(%value, 1, 8)] -config.decl @qwen3_moe.routed_down.expert_count : %value: index where [range(%value, 1, 512)] - config.decl @qwen3_moe.routed_down.output_size : %value: index where [range(%value, 1, 4096)] -kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count$10: index, %input_size$11: index) launch(%token_count$12: index, %input_size$13: index, %input: buffer, %output: buffer) - -kernel.decl @ggml_moe_build_expert_table(%token_count$16: index, %route_count$17: index, %route_stride$18: index, %expert_count$19: index) launch(%token_count$20: index, %route_count$21: index, %route_stride$22: index, %expert_count$23: index, %route_ids: buffer, %expert_table: buffer) - -kernel.decl @ggml_moe_build_expert_partition_table(%token_count$26: index, %route_count$27: index, %expert_count$28: index) launch(%token_count$29: index, %route_count$30: index, %expert_count$31: index, %expert_table: buffer, %partition_table: buffer) - -kernel.decl @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma(%token_count$34: index) launch(%token_count$35: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) - -kernel.decl @ggml_mul_mat_id_vector_pair_binary_publish_f32(%token_count$42: index, %route_count$43: index, %route_stride$44: index, %expert_count$45: index, %output_size$46: index) launch(%token_count$47: index, %route_count$48: index, %route_stride$49: index, %expert_count$50: index, %output_size$51: index, %q8_input: buffer, %route_ids: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) - -func.decl @qwen3_moe_q4k_wmma_vector4(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group: index, %packet: index) -> (vector<4xf16>) - -func.decl @qwen3_moe_q4k_wmma_load_header(%weight: buffer, %row_byte_base: offset, %q4_block: index) -> (vector<4xi32>) - -func.decl @qwen3_moe_q4k_wmma_load_code(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group_pair: index, %packet: index) -> (vector<1xi32>) - -func.decl @qwen3_moe_q4k_wmma_vector4_from_header_code(%q4_group: index, %header_words: vector<4xi32>, %q_word: vector<1xi32>) -> (vector<4xf16>) - -func.decl @ggml_q6k_f16_vector4(%weight: buffer, %weight_row_byte_base: offset, %q6_block: index, %q6_group: index, %packet: index) -> (vector<4xf16>) - -kernel.decl @qwen3_moe_routed_down_q4k_q8_1_x4(%token_count$83: index, %input_size$84: index, %route_count$85: index, %route_id_stride$86: index, %expert_count$87: index, %output_size$88: index) launch(%token_count$89: index, %input_size$90: index, %route_count$91: index, %route_id_stride$92: index, %expert_count$93: index, %output_size$94: index, %q8_input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) - -kernel.decl @qwen3_moe_routed_down_q6k_q8_1_x4(%token_count$100: index, %input_size$101: index, %route_count$102: index, %route_id_stride$103: index, %expert_count$104: index, %output_size$105: index) launch(%token_count$106: index, %input_size$107: index, %route_count$108: index, %route_id_stride$109: index, %expert_count$110: index, %output_size$111: index, %q8_input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) - -// Acquires the packed words shared by one four-group half of a Q6_K block. -// Each QL word supplies two groups and the QH word supplies all four. -func.def inline @qwen3_moe_q6k_wmma_load_half_codes(%weight: buffer, %weight_row_byte_base: offset, %q6_block: index, %q6_half: index, %packet: index) -> (vector<1xi32>, vector<1xi32>, vector<1xi32>) { - %c8 = index.constant 8 : index - %c16 = index.constant 16 : index - %block_bytes = index.constant 210 : offset - %qh_byte_add = index.constant 128 : offset - %bounded_half = index.assume %q6_half [range(%q6_half, 0, 1)] : index - %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index - %block_byte_add = index.scale %q6_block, %block_bytes : index, offset -> offset - %block_byte_base = index.add %weight_row_byte_base, %block_byte_add : offset - %qh_byte_base = index.add %block_byte_base, %qh_byte_add : offset - %ql_view = buffer.view %weight[%block_byte_base] : buffer -> view<32xi32> - %qh_view = buffer.view %weight[%qh_byte_base] : buffer -> view<16xi32> - %ql_half_word_base = index.mul %bounded_half, %c16 : index - %ql0_word_index = index.add %ql_half_word_base, %bounded_packet : index - %ql1_word_base = index.add %ql_half_word_base, %c8 : index - %ql1_word_index = index.add %ql1_word_base, %bounded_packet : index - %qh_half_word_base = index.mul %bounded_half, %c8 : index - %qh_word_index = index.add %qh_half_word_base, %bounded_packet : index - %ql0_word = vector.load %ql_view[%ql0_word_index] : view<32xi32> -> vector<1xi32> - %ql1_word = vector.load %ql_view[%ql1_word_index] : view<32xi32> -> vector<1xi32> - %qh_word = vector.load %qh_view[%qh_word_index] : view<16xi32> -> vector<1xi32> - func.return %ql0_word, %ql1_word, %qh_word : vector<1xi32>, vector<1xi32>, vector<1xi32> -} - -// Decodes four adjacent values after the surrounding schedule has selected -// the scale and retained the packed code words at their natural lifetimes. -func.def inline @qwen3_moe_q6k_wmma_vector4_from_scale_codes(%q6_group: index, %scale_i8: i8, %d_f16: f16, %ql_word: vector<1xi32>, %qh_word: vector<1xi32>) -> (vector<4xf16>) { - %c2 = index.constant 2 : index - %c4 = index.constant 4 : index - %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 - %group_in_half = index.rem %bounded_group, %c4 : 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> - %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 = 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_f32 = vector.mulf %centered, %combined_scale_vector : vector<4xf32> - %values = vector.fptrunc %values_f32 : vector<4xf32> to vector<4xf16> - func.return %values : vector<4xf16> -} - -// Loads only the group scale and block multiplier while reusing packed words -// retained by an enclosing four-group half-block schedule. -func.def inline @qwen3_moe_q6k_wmma_vector4_from_half_codes(%weight: buffer, %weight_row_byte_base: offset, %q6_block: index, %q6_group: index, %packet: index, %ql0_word: vector<1xi32>, %ql1_word: vector<1xi32>, %qh_word: vector<1xi32>) -> (vector<4xf16>) { - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c2 = index.constant 2 : index - %c4 = index.constant 4 : index - %block_bytes = index.constant 210 : offset - %scale_byte_add = index.constant 192 : offset - %d_byte_add = index.constant 208 : offset - %bounded_group = index.assume %q6_group [range(%q6_group, 0, 7)] : index - %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index - %block_byte_add = index.scale %q6_block, %block_bytes : index, offset -> offset - %block_byte_base = index.add %weight_row_byte_base, %block_byte_add : offset - %scale_byte_base = index.add %block_byte_base, %scale_byte_add : offset - %d_byte_base = index.add %block_byte_base, %d_byte_add : offset - %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<16xi8> - %d_view = buffer.view %weight[%d_byte_base] : buffer -> view<1xf16> - %group_in_half = index.rem %bounded_group, %c4 : index - %ql_side = index.rem %group_in_half, %c2 : index - %uses_ql1 = index.cmp eq, %ql_side, %c1 : index - %ql_word = scf.select %uses_ql1, %ql1_word, %ql0_word : vector<1xi32> - %scale_packet_half = index.div %bounded_packet, %c4 : index - %scale_group_base = index.mul %bounded_group, %c2 : index - %scale_index = index.add %scale_group_base, %scale_packet_half : index - %scale_i8 = view.load %scale_view[%scale_index] : view<16xi8> -> i8 - %d_f16 = view.load %d_view[%c0] : view<1xf16> -> f16 - %values = func.call @qwen3_moe_q6k_wmma_vector4_from_scale_codes(%bounded_group, %scale_i8, %d_f16, %ql_word, %qh_word) : (index, i8, f16, vector<1xi32>, vector<1xi32>) -> (vector<4xf16>) - func.return %values : vector<4xf16> -} - -// Shared raw-quantized matrix schedule. Entry points pass a literal weight -// format so linking and JIT specialization erase the inactive packed decoder. -template.def<@qwen3_moe.routed_down.quantized_prefill.body> device @qwen3_moe_routed_down_quantized_f16_wmma_body(%weight_format: index, %token_count: index, %input: buffer, %expert_table: buffer, %weight: buffer, %output: buffer) { - %input_size = config.get @qwen3_moe.routed_down.input_size : index - %route_count = config.get @qwen3_moe.routed_down.route_count : index - %expert_count = config.get @qwen3_moe.routed_down.expert_count : index - %output_size = config.get @qwen3_moe.routed_down.output_size : index - %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index - %bounded_expert_count = index.assume %expert_count [range(%expert_count, 1, 128)] : index - %bounded_output_size = index.assume %output_size [range(%output_size, 1, 4096)] : index - %channel_tile = kernel.workgroup.id : index - %route_tile = kernel.workgroup.id : index - %expert = kernel.workgroup.id : index - %workitem = kernel.workitem.id : index - %subgroup0 = kernel.subgroup.id : index - %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 1)] : index - %lane = kernel.subgroup.lane.id : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c2 = index.constant 2 : index - %c4 = index.constant 4 : index - %c6 = index.constant 6 : index - %c8 = index.constant 8 : index - %c16 = index.constant 16 : index - %c32 = index.constant 32 : index - %c40 = index.constant 40 : index - %c48 = index.constant 48 : index - %c63 = index.constant 63 : index - %c64 = index.constant 64 : index - %c256 = index.constant 256 : index - %c0_offset = index.constant 0 : offset - %c4_bytes = index.constant 4 : offset - %q4_block_bytes = index.constant 144 : offset - %q6_block_bytes = index.constant 210 : offset - %weight_stage_bytes = index.constant 5120 : offset - %activation_stage_bytes = index.constant 2560 : offset - %route_stage_bytes = index.constant 128 : offset - %wave_result_stage_bytes = index.constant 512 : offset - %result_stage_bytes = index.constant 1024 : offset - %c0_i32 = scalar.constant 0 : i32 - %cn1_i32 = scalar.constant -1 : i32 - %c0_i32x1 = vector.constant 0 : vector<1xi32> - %c0_i32x4 = vector.constant 0 : vector<4xi32> - %c0_f16x4 = vector.constant 0.0 : vector<4xf16> - %zero_accumulator = vector.constant 0.0 : vector<8xf16> - %m = index.constant 16 : index - %n = index.constant 16 : index - %k = index.constant 16 : index - %assignment_count = index.mul %token_count, %bounded_route_count : index - %assignment_table_byte_base = index.scale %bounded_expert_count, %c4_bytes : index, offset -> offset - %is_q4 = index.cmp eq, %weight_format, %c4 : index - %is_q6 = index.cmp eq, %weight_format, %c6 : index - %quant_block_bytes = scf.select %is_q6, %q6_block_bytes, %q4_block_bytes : offset - %quant_block_count = index.div %input_size, %c256 : index - %weight_row_bytes = index.scale %quant_block_count, %quant_block_bytes : index, offset -> offset - %weight_expert_bytes = index.scale %bounded_output_size, %weight_row_bytes : index, offset -> offset - %input_noalias, %expert_table_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %input, %expert_table, %weight, %output : buffer, buffer, buffer, buffer - %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%assignment_count]x[%input_size]xf16> - %count_view = buffer.view %expert_table_noalias[%c0_offset] : buffer -> view<[%bounded_expert_count]xi32> - %assignment_view = buffer.view %expert_table_noalias[%assignment_table_byte_base] : buffer -> view<[%bounded_expert_count]x[%token_count]xi32> - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%assignment_count]x[%bounded_output_size]xf16> - %weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer - %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer - %route_stage = buffer.alloca align(16) %route_stage_bytes : buffer - %result_stage = buffer.alloca align(16) %result_stage_bytes : buffer - %weight_stage_view = buffer.view %weight_stage[%c0_offset] : buffer -> view<64x40xf16> - %activation_stage_physical_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x40xf16> - %activation_fragment_layout = encoding.layout.strided [1, %c40] : encoding - %activation_fragment_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x32xf16, %activation_fragment_layout> - %route_stage_view = buffer.view %route_stage[%c0_offset] : buffer -> view<32xi32> - %wave_result_stage_offset = index.scale %subgroup, %wave_result_stage_bytes : index, offset -> offset - %result_fragment_layout = encoding.layout.strided [1, %c16] : encoding - %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 - %initial_route_tile_base = index.mul %route_tile, %c32 : index - // The launch workload fixes the interleaved route partitions for this exact - // command-program specialization. - %padded_token_count = index.add %token_count, %c63 : index - %route_partition_count = index.div %padded_token_count, %c64 : index - %route_partition_step = index.mul %route_partition_count, %c32 : index - %bounded_expert, %table_expert_count = index.assume %expert, %bounded_expert_count [lt(%expert, %bounded_expert_count)] : index, index - %is_workitem_zero = index.cmp eq, %workitem, %c0 : index - %lane_expert_route_count = scf.if %is_workitem_zero -> (i32) { - %loaded = view.load %count_view[%bounded_expert] : view<[%bounded_expert_count]xi32> -> i32 - scf.yield %loaded : i32 - } else { - scf.yield %c0_i32 : i32 - } - %expert_route_count_reduced = kernel.workgroup.reduce %lane_expert_route_count : i32 - %expert_route_count_i32 = kernel.subgroup.broadcast.first %expert_route_count_reduced : i32 - %expert_route_count0 = index.cast %expert_route_count_i32 : i32 to index - %expert_route_count = index.assume %expert_route_count0 [range(%expert_route_count0, 0, 2048)] : index - scf.for %route_tile_base = [%initial_route_tile_base to %expert_route_count step %route_partition_step] { - // Snapshot this expert's compact assignment map once, then reuse it for - // every K group and output-channel tile. - %loads_route = index.cmp ult, %workitem, %c32 : index - scf.if %loads_route { - %local_route = index.assume %workitem [range(%workitem, 0, 31)] : index - %assignment_ordinal = index.add %route_tile_base, %local_route : index - %valid_row = index.cmp ult, %assignment_ordinal, %expert_route_count : index - %assignment_i32 = scf.if %valid_row -> (i32) { - %bounded_assignment_ordinal, %table_token_count = index.assume %assignment_ordinal, %token_count [lt(%assignment_ordinal, %token_count)] : index, index - %loaded = view.load %assignment_view[%bounded_expert, %bounded_assignment_ordinal] : view<[%bounded_expert_count]x[%token_count]xi32> -> i32 - scf.yield %loaded : i32 - } else { - scf.yield %cn1_i32 : i32 - } - view.store %assignment_i32, %route_stage_view[%local_route] : i32, view<32xi32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %load_packet = index.rem %workitem, %c8 : index - %load_k = index.mul %load_packet, %c4 : index - %load_row0 = index.div %workitem, %c8 : index - %load_row = index.assume %load_row0 [range(%load_row0, 0, 15)] : index - %expert_byte_base = index.scale %bounded_expert, %weight_expert_bytes : index, offset -> offset - %subgroup_channel_add = index.mul %subgroup, %c32 : index - %subgroup_channel1 = index.add %subgroup_channel_add, %c16 : index - %init00 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> - %init01 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> - %init10 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> - %init11 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> - %result00, %result01, %result10, %result11 = scf.for %quant_block = [%c0 to %quant_block_count 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>) { - // Four lanes cover the 64 output rows owned by one load packet. Snapshot - // each Q4_K block header once and retain it across all eight quant groups. - // The Q6 specialization erases this loop with %is_q4=false. - %header0, %header1, %header2, %header3 = scf.for %header_row_offset = [%c0 to %c64 step %c16](%prior_header0 = %c0_i32x4 : vector<4xi32>, %prior_header1 = %c0_i32x4 : vector<4xi32>, %prior_header2 = %c0_i32x4 : vector<4xi32>, %prior_header3 = %c0_i32x4 : vector<4xi32>) -> (vector<4xi32>, vector<4xi32>, vector<4xi32>, vector<4xi32>) unroll { - %header_local_row0 = index.add %load_row, %header_row_offset : index - %header_local_row = index.assume %header_local_row0 [range(%header_local_row0, 0, 63)] : index - %header_channel = index.add %channel_tile_base, %header_local_row : index - %valid_header_channel = index.cmp ult, %header_channel, %bounded_output_size : index - %loads_q4_header = scalar.andi %is_q4, %valid_header_channel : i1 - %loaded_header = scf.if %loads_q4_header -> (vector<4xi32>) { - %header_channel_byte_add = index.scale %header_channel, %weight_row_bytes : index, offset -> offset - %header_row_byte_base = index.add %expert_byte_base, %header_channel_byte_add : offset - %header_words = func.call @qwen3_moe_q4k_wmma_load_header(%weight_noalias, %header_row_byte_base, %quant_block) : (buffer, offset, index) -> (vector<4xi32>) - scf.yield %header_words : vector<4xi32> - } else { - scf.yield %c0_i32x4 : vector<4xi32> - } - %updates_header0 = index.cmp eq, %header_row_offset, %c0 : index - %updates_header1 = index.cmp eq, %header_row_offset, %c16 : index - %updates_header2 = index.cmp eq, %header_row_offset, %c32 : index - %updates_header3 = index.cmp eq, %header_row_offset, %c48 : index - %next_header0 = scf.select %updates_header0, %loaded_header, %prior_header0 : vector<4xi32> - %next_header1 = scf.select %updates_header1, %loaded_header, %prior_header1 : vector<4xi32> - %next_header2 = scf.select %updates_header2, %loaded_header, %prior_header2 : vector<4xi32> - %next_header3 = scf.select %updates_header3, %loaded_header, %prior_header3 : vector<4xi32> - scf.yield %next_header0, %next_header1, %next_header2, %next_header3 : vector<4xi32>, vector<4xi32>, vector<4xi32>, vector<4xi32> - } - %quant_group_outer_count = scf.select %is_q4, %c4, %c2 : index - %groups_per_outer = scf.select %is_q4, %c2, %c4 : index - %block_result00, %block_result01, %block_result10, %block_result11 = scf.for %quant_group_outer = [%c0 to %quant_group_outer_count step %c1](%acc00 = %block_acc00 : vector<8xf16>, %acc01 = %block_acc01 : vector<8xf16>, %acc10 = %block_acc10 : vector<8xf16>, %acc11 = %block_acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { - // Q4_K retains one packed code word across its adjacent low/high group - // pair. The Q6 specialization erases this acquisition path. - %q_word0, %q_word1, %q_word2, %q_word3 = scf.for %code_row_offset = [%c0 to %c64 step %c16](%prior_q_word0 = %c0_i32x1 : vector<1xi32>, %prior_q_word1 = %c0_i32x1 : vector<1xi32>, %prior_q_word2 = %c0_i32x1 : vector<1xi32>, %prior_q_word3 = %c0_i32x1 : vector<1xi32>) -> (vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>) unroll { - %code_local_row0 = index.add %load_row, %code_row_offset : index - %code_local_row = index.assume %code_local_row0 [range(%code_local_row0, 0, 63)] : index - %code_channel = index.add %channel_tile_base, %code_local_row : index - %valid_code_channel = index.cmp ult, %code_channel, %bounded_output_size : index - %loads_q4_code = scalar.andi %is_q4, %valid_code_channel : i1 - %loaded_q_word = scf.if %loads_q4_code -> (vector<1xi32>) { - %q4_group_pair = index.assume %quant_group_outer [range(%quant_group_outer, 0, 3)] : index - %code_channel_byte_add = index.scale %code_channel, %weight_row_bytes : index, offset -> offset - %code_row_byte_base = index.add %expert_byte_base, %code_channel_byte_add : offset - %q_word = func.call @qwen3_moe_q4k_wmma_load_code(%weight_noalias, %code_row_byte_base, %quant_block, %q4_group_pair, %load_packet) : (buffer, offset, index, index, index) -> (vector<1xi32>) - scf.yield %q_word : vector<1xi32> - } else { - scf.yield %c0_i32x1 : vector<1xi32> - } - %updates_q_word0 = index.cmp eq, %code_row_offset, %c0 : index - %updates_q_word1 = index.cmp eq, %code_row_offset, %c16 : index - %updates_q_word2 = index.cmp eq, %code_row_offset, %c32 : index - %updates_q_word3 = index.cmp eq, %code_row_offset, %c48 : index - %next_q_word0 = scf.select %updates_q_word0, %loaded_q_word, %prior_q_word0 : vector<1xi32> - %next_q_word1 = scf.select %updates_q_word1, %loaded_q_word, %prior_q_word1 : vector<1xi32> - %next_q_word2 = scf.select %updates_q_word2, %loaded_q_word, %prior_q_word2 : vector<1xi32> - %next_q_word3 = scf.select %updates_q_word3, %loaded_q_word, %prior_q_word3 : vector<1xi32> - scf.yield %next_q_word0, %next_q_word1, %next_q_word2, %next_q_word3 : vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32> - } - // Four Q6_K groups in one half-block share a QH word and consume two - // nibbles from each of two QL words. Retain the three packed words for - // all four groups. The Q4 specialization erases this acquisition path. - %q6_ql00, %q6_ql10, %q6_qh0, %q6_ql01, %q6_ql11, %q6_qh1, %q6_ql02, %q6_ql12, %q6_qh2, %q6_ql03, %q6_ql13, %q6_qh3 = scf.for %q6_code_row_offset = [%c0 to %c64 step %c16](%prior_q6_ql00 = %c0_i32x1 : vector<1xi32>, %prior_q6_ql10 = %c0_i32x1 : vector<1xi32>, %prior_q6_qh0 = %c0_i32x1 : vector<1xi32>, %prior_q6_ql01 = %c0_i32x1 : vector<1xi32>, %prior_q6_ql11 = %c0_i32x1 : vector<1xi32>, %prior_q6_qh1 = %c0_i32x1 : vector<1xi32>, %prior_q6_ql02 = %c0_i32x1 : vector<1xi32>, %prior_q6_ql12 = %c0_i32x1 : vector<1xi32>, %prior_q6_qh2 = %c0_i32x1 : vector<1xi32>, %prior_q6_ql03 = %c0_i32x1 : vector<1xi32>, %prior_q6_ql13 = %c0_i32x1 : vector<1xi32>, %prior_q6_qh3 = %c0_i32x1 : vector<1xi32>) -> (vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>) unroll { - %q6_code_local_row0 = index.add %load_row, %q6_code_row_offset : index - %q6_code_local_row = index.assume %q6_code_local_row0 [range(%q6_code_local_row0, 0, 63)] : index - %q6_code_channel = index.add %channel_tile_base, %q6_code_local_row : index - %valid_q6_code_channel = index.cmp ult, %q6_code_channel, %bounded_output_size : index - %loads_q6_code = scalar.andi %is_q6, %valid_q6_code_channel : i1 - %loaded_q6_ql0, %loaded_q6_ql1, %loaded_q6_qh = scf.if %loads_q6_code -> (vector<1xi32>, vector<1xi32>, vector<1xi32>) { - %q6_half = index.assume %quant_group_outer [range(%quant_group_outer, 0, 1)] : index - %q6_code_channel_byte_add = index.scale %q6_code_channel, %weight_row_bytes : index, offset -> offset - %q6_code_row_byte_base = index.add %expert_byte_base, %q6_code_channel_byte_add : offset - %ql0_word, %ql1_word, %qh_word = func.call @qwen3_moe_q6k_wmma_load_half_codes(%weight_noalias, %q6_code_row_byte_base, %quant_block, %q6_half, %load_packet) : (buffer, offset, index, index, index) -> (vector<1xi32>, vector<1xi32>, vector<1xi32>) - scf.yield %ql0_word, %ql1_word, %qh_word : vector<1xi32>, vector<1xi32>, vector<1xi32> - } else { - scf.yield %c0_i32x1, %c0_i32x1, %c0_i32x1 : vector<1xi32>, vector<1xi32>, vector<1xi32> - } - %updates_q6_code0 = index.cmp eq, %q6_code_row_offset, %c0 : index - %updates_q6_code1 = index.cmp eq, %q6_code_row_offset, %c16 : index - %updates_q6_code2 = index.cmp eq, %q6_code_row_offset, %c32 : index - %updates_q6_code3 = index.cmp eq, %q6_code_row_offset, %c48 : index - %next_q6_ql00 = scf.select %updates_q6_code0, %loaded_q6_ql0, %prior_q6_ql00 : vector<1xi32> - %next_q6_ql10 = scf.select %updates_q6_code0, %loaded_q6_ql1, %prior_q6_ql10 : vector<1xi32> - %next_q6_qh0 = scf.select %updates_q6_code0, %loaded_q6_qh, %prior_q6_qh0 : vector<1xi32> - %next_q6_ql01 = scf.select %updates_q6_code1, %loaded_q6_ql0, %prior_q6_ql01 : vector<1xi32> - %next_q6_ql11 = scf.select %updates_q6_code1, %loaded_q6_ql1, %prior_q6_ql11 : vector<1xi32> - %next_q6_qh1 = scf.select %updates_q6_code1, %loaded_q6_qh, %prior_q6_qh1 : vector<1xi32> - %next_q6_ql02 = scf.select %updates_q6_code2, %loaded_q6_ql0, %prior_q6_ql02 : vector<1xi32> - %next_q6_ql12 = scf.select %updates_q6_code2, %loaded_q6_ql1, %prior_q6_ql12 : vector<1xi32> - %next_q6_qh2 = scf.select %updates_q6_code2, %loaded_q6_qh, %prior_q6_qh2 : vector<1xi32> - %next_q6_ql03 = scf.select %updates_q6_code3, %loaded_q6_ql0, %prior_q6_ql03 : vector<1xi32> - %next_q6_ql13 = scf.select %updates_q6_code3, %loaded_q6_ql1, %prior_q6_ql13 : vector<1xi32> - %next_q6_qh3 = scf.select %updates_q6_code3, %loaded_q6_qh, %prior_q6_qh3 : vector<1xi32> - scf.yield %next_q6_ql00, %next_q6_ql10, %next_q6_qh0, %next_q6_ql01, %next_q6_ql11, %next_q6_qh1, %next_q6_ql02, %next_q6_ql12, %next_q6_qh2, %next_q6_ql03, %next_q6_ql13, %next_q6_qh3 : vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32> - } - %outer_result00, %outer_result01, %outer_result10, %outer_result11 = scf.for %group_within_outer = [%c0 to %groups_per_outer step %c1](%group_acc00 = %acc00 : vector<8xf16>, %group_acc01 = %acc01 : vector<8xf16>, %group_acc10 = %acc10 : vector<8xf16>, %group_acc11 = %acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { - %quant_group_base = index.mul %quant_group_outer, %groups_per_outer : index - %quant_group0 = index.add %quant_group_base, %group_within_outer : index - %quant_group = index.assume %quant_group0 [range(%quant_group0, 0, 7)] : index - %block_k_base = index.mul %quant_block, %c256 : index - %group_k_add = index.mul %quant_group, %c32 : index - %k_origin = index.add %block_k_base, %group_k_add : index - scf.for %row_offset = [%c0 to %c64 step %c16] unroll { - %local_row0 = index.add %load_row, %row_offset : index - %local_row = index.assume %local_row0 [range(%local_row0, 0, 63)] : index - %selects_row0 = index.cmp eq, %row_offset, %c0 : index - %selects_row2 = index.cmp eq, %row_offset, %c32 : index - %selects_low_row_pair = index.cmp ult, %row_offset, %c32 : index - %selected_header01 = scf.select %selects_row0, %header0, %header1 : vector<4xi32> - %selected_header23 = scf.select %selects_row2, %header2, %header3 : vector<4xi32> - %selected_header = scf.select %selects_low_row_pair, %selected_header01, %selected_header23 : vector<4xi32> - %selected_q_word01 = scf.select %selects_row0, %q_word0, %q_word1 : vector<1xi32> - %selected_q_word23 = scf.select %selects_row2, %q_word2, %q_word3 : vector<1xi32> - %selected_q_word = scf.select %selects_low_row_pair, %selected_q_word01, %selected_q_word23 : vector<1xi32> - %selected_q6_ql0_01 = scf.select %selects_row0, %q6_ql00, %q6_ql01 : vector<1xi32> - %selected_q6_ql0_23 = scf.select %selects_row2, %q6_ql02, %q6_ql03 : vector<1xi32> - %selected_q6_ql0 = scf.select %selects_low_row_pair, %selected_q6_ql0_01, %selected_q6_ql0_23 : vector<1xi32> - %selected_q6_ql1_01 = scf.select %selects_row0, %q6_ql10, %q6_ql11 : vector<1xi32> - %selected_q6_ql1_23 = scf.select %selects_row2, %q6_ql12, %q6_ql13 : vector<1xi32> - %selected_q6_ql1 = scf.select %selects_low_row_pair, %selected_q6_ql1_01, %selected_q6_ql1_23 : vector<1xi32> - %selected_q6_qh01 = scf.select %selects_row0, %q6_qh0, %q6_qh1 : vector<1xi32> - %selected_q6_qh23 = scf.select %selects_row2, %q6_qh2, %q6_qh3 : vector<1xi32> - %selected_q6_qh = scf.select %selects_low_row_pair, %selected_q6_qh01, %selected_q6_qh23 : vector<1xi32> - %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>) { - %decoded = scf.if %is_q6 -> (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 - %q6_values = func.call @qwen3_moe_q6k_wmma_vector4_from_half_codes(%weight_noalias, %row_byte_base, %quant_block, %quant_group, %load_packet, %selected_q6_ql0, %selected_q6_ql1, %selected_q6_qh) : (buffer, offset, index, index, index, vector<1xi32>, vector<1xi32>, vector<1xi32>) -> (vector<4xf16>) - scf.yield %q6_values : vector<4xf16> - } else { - %q4_values = func.call @qwen3_moe_q4k_wmma_vector4_from_header_code(%quant_group, %selected_header, %selected_q_word) : (index, vector<4xi32>, vector<1xi32>) -> (vector<4xf16>) - scf.yield %q4_values : vector<4xf16> - } - scf.yield %decoded : vector<4xf16> - } else { - scf.yield %c0_f16x4 : vector<4xf16> - } - %is_activation_row = index.cmp ult, %local_row, %c32 : index - vector.store %weight_values, %weight_stage_view[%local_row, %load_k] : vector<4xf16>, view<64x40xf16> - scf.if %is_activation_row { - %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>) { - %assignment0 = index.cast %assignment_i32 : i32 to index - %assignment = index.assume %assignment0 [range(%assignment0, 0, 16383)] : index - %bounded_assignment, %bounded_assignment_count = index.assume %assignment, %assignment_count [lt(%assignment, %assignment_count)] : index, index - %input_k = index.add %k_origin, %load_k : index - %loaded = vector.load %input_view[%bounded_assignment, %input_k] : view<[%assignment_count]x[%input_size]xf16> -> vector<4xf16> - scf.yield %loaded : vector<4xf16> - } else { - scf.yield %c0_f16x4 : vector<4xf16> - } - vector.store %activation_values, %activation_stage_physical_view[%activation_row, %load_k] : vector<4xf16>, view<32x40xf16> - } - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %next00, %next01, %next10, %next11 = scf.for %k_half = [%c0 to %c32 step %c16](%half_acc00 = %group_acc00 : vector<8xf16>, %half_acc01 = %group_acc01 : vector<8xf16>, %half_acc10 = %group_acc10 : vector<8xf16>, %half_acc11 = %group_acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) unroll { - %lhs0 = vector.fragment.load %weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> - %lhs1 = vector.fragment.load %weight_stage_view[%subgroup_channel1, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> - %rhs0 = vector.fragment.load %activation_fragment_view[%k_half, %c0] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> - %rhs1 = vector.fragment.load %activation_fragment_view[%k_half, %c16] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> - %half_next00 = vector.mma %lhs0, %rhs0, %half_acc00 : vector<16xf16>, vector<16xf16>, vector<8xf16> - %half_next01 = vector.mma %lhs0, %rhs1, %half_acc01 : vector<16xf16>, vector<16xf16>, vector<8xf16> - %half_next10 = vector.mma %lhs1, %rhs0, %half_acc10 : vector<16xf16>, vector<16xf16>, vector<8xf16> - %half_next11 = vector.mma %lhs1, %rhs1, %half_acc11 : vector<16xf16>, vector<16xf16>, vector<8xf16> - scf.yield %half_next00, %half_next01, %half_next10, %half_next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - scf.yield %next00, %next01, %next10, %next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> - } - scf.yield %outer_result00, %outer_result01, %outer_result10, %outer_result11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> - } - scf.yield %block_result00, %block_result01, %block_result10, %block_result11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> - } - // WMMA produces [channel][route] fragments. Transpose through wave-private - // LDS so each lane publishes four adjacent channels for one assignment. - %publish_route0 = index.div %lane, %c4 : index - %publish_route = index.assume %publish_route0 [range(%publish_route0, 0, 15)] : index - %publish_packet0 = index.rem %lane, %c4 : index - %publish_packet = index.assume %publish_packet0 [range(%publish_packet0, 0, 3)] : index - %publish_channel_add = index.mul %publish_packet, %c4 : index - %local_route1 = index.add %c16, %publish_route : index - %assignment0_i32 = view.load %route_stage_view[%publish_route] : view<32xi32> -> i32 - %assignment1_i32 = view.load %route_stage_view[%local_route1] : view<32xi32> -> i32 - %assignment0_nonnegative = scalar.cmpi sge, %assignment0_i32, %c0_i32 : i32 - %assignment1_nonnegative = scalar.cmpi sge, %assignment1_i32, %c0_i32 : i32 - %safe_assignment0_i32 = scf.select %assignment0_nonnegative, %assignment0_i32, %c0_i32 : i32 - %safe_assignment1_i32 = scf.select %assignment1_nonnegative, %assignment1_i32, %c0_i32 : i32 - %safe_assignment0_0 = index.cast %safe_assignment0_i32 : i32 to index - %safe_assignment1_0 = index.cast %safe_assignment1_i32 : i32 to index - %safe_assignment0 = index.assume %safe_assignment0_0 [range(%safe_assignment0_0, 0, 16383)] : index - %safe_assignment1 = index.assume %safe_assignment1_0 [range(%safe_assignment1_0, 0, 16383)] : index - %bounded_assignment0, %bounded_assignment_count0 = index.assume %safe_assignment0, %assignment_count [lt(%safe_assignment0, %assignment_count)] : index, index - %bounded_assignment1, %bounded_assignment_count1 = index.assume %safe_assignment1, %assignment_count [lt(%safe_assignment1, %assignment_count)] : index, index - %subgroup_channel_base = index.add %channel_tile_base, %subgroup_channel_add : index - %channel0 = index.add %subgroup_channel_base, %publish_channel_add : index - %channel1_base = index.add %subgroup_channel_base, %c16 : index - %channel1 = index.add %channel1_base, %publish_channel_add : index - %valid_channel0 = index.cmp ult, %channel0, %bounded_output_size : index - %valid_channel1 = index.cmp ult, %channel1, %bounded_output_size : index - %writes00 = scalar.andi %assignment0_nonnegative, %valid_channel0 : i1 - %writes01 = scalar.andi %assignment1_nonnegative, %valid_channel0 : i1 - %writes10 = scalar.andi %assignment0_nonnegative, %valid_channel1 : i1 - %writes11 = scalar.andi %assignment1_nonnegative, %valid_channel1 : i1 - vector.fragment.store %result00, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> - kernel.barrier scope(subgroup) ordering(acq_rel) - scf.if %writes00 { - %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> - %mask = vector.mask.range [%channel0 to %bounded_output_size step %c1] : index -> vector<4xi1> - vector.store.mask %values, %output_view[%bounded_assignment0, %channel0], %mask : vector<4xf16>, view<[%assignment_count]x[%bounded_output_size]xf16>, vector<4xi1> - } - kernel.barrier scope(subgroup) ordering(acq_rel) - vector.fragment.store %result01, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> - kernel.barrier scope(subgroup) ordering(acq_rel) - scf.if %writes01 { - %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> - %mask = vector.mask.range [%channel0 to %bounded_output_size step %c1] : index -> vector<4xi1> - vector.store.mask %values, %output_view[%bounded_assignment1, %channel0], %mask : vector<4xf16>, view<[%assignment_count]x[%bounded_output_size]xf16>, vector<4xi1> - } - kernel.barrier scope(subgroup) ordering(acq_rel) - vector.fragment.store %result10, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> - kernel.barrier scope(subgroup) ordering(acq_rel) - scf.if %writes10 { - %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> - %mask = vector.mask.range [%channel1 to %bounded_output_size step %c1] : index -> vector<4xi1> - vector.store.mask %values, %output_view[%bounded_assignment0, %channel1], %mask : vector<4xf16>, view<[%assignment_count]x[%bounded_output_size]xf16>, vector<4xi1> - } - kernel.barrier scope(subgroup) ordering(acq_rel) - vector.fragment.store %result11, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> - kernel.barrier scope(subgroup) ordering(acq_rel) - scf.if %writes11 { - %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> - %mask = vector.mask.range [%channel1 to %bounded_output_size step %c1] : index -> vector<4xi1> - vector.store.mask %values, %output_view[%bounded_assignment1, %channel1], %mask : vector<4xf16>, view<[%assignment_count]x[%bounded_output_size]xf16>, vector<4xi1> - } - // Route, operand, and result stages are reused by the next concentrated - // routing partition. All waves must finish publication before reuse. - kernel.barrier scope(workgroup) ordering(acq_rel) - } - template.return -} - -// Q4_K and Q6_K retain separate entry points while sharing the complete matrix -// schedule. This keeps format routing outside the hot kernel. -kernel.def target(@qwen3_moe_routed_down_gfx11_wave64) @qwen3_moe_routed_down_q4k_f16_wmma_grouped(%token_count: index) { - %expert_count = config.get @qwen3_moe.routed_down.expert_count : index - %output_size = config.get @qwen3_moe.routed_down.output_size : index - %c1 = index.constant 1 : index - %c63 = index.constant 63 : index - %c64 = index.constant 64 : index - %c128 = index.constant 128 : index - %padded_output_size = index.add %output_size, %c63 : index - %output_tiles = index.div %padded_output_size, %c64 : index - %padded_token_count = index.add %token_count, %c63 : index - %route_tiles = index.div %padded_token_count, %c64 : index - kernel.launch.config workgroups(%output_tiles, %route_tiles, %expert_count) workgroup_size(%c128, %c1, %c1) : index -} launch(%token_count: index, %input: buffer, %expert_table: buffer, %weight: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { - %q4 = index.constant 4 : index - template.apply<@qwen3_moe.routed_down.quantized_prefill.body>(%q4, %token_count, %input, %expert_table, %weight, %output) : (index, index, buffer, buffer, buffer, buffer) - kernel.return -} - -kernel.def target(@qwen3_moe_routed_down_gfx11_wave64) @qwen3_moe_routed_down_q6k_f16_wmma_grouped(%token_count: index) { - %expert_count = config.get @qwen3_moe.routed_down.expert_count : index - %output_size = config.get @qwen3_moe.routed_down.output_size : index - %c1 = index.constant 1 : index - %c63 = index.constant 63 : index - %c64 = index.constant 64 : index - %c128 = index.constant 128 : index - %padded_output_size = index.add %output_size, %c63 : index - %output_tiles = index.div %padded_output_size, %c64 : index - %padded_token_count = index.add %token_count, %c63 : index - %route_tiles = index.div %padded_token_count, %c64 : index - kernel.launch.config workgroups(%output_tiles, %route_tiles, %expert_count) workgroup_size(%c128, %c1, %c1) : index -} launch(%token_count: index, %input: buffer, %expert_table: buffer, %weight: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { - %q6 = index.constant 6 : index - template.apply<@qwen3_moe.routed_down.quantized_prefill.body>(%q6, %token_count, %input, %expert_table, %weight, %output) : (index, index, buffer, buffer, buffer, buffer) - kernel.return -} - -// Reduces compact routed projections into the residual in logical token order. -// -// One workitem owns four adjacent output channels. All top-k routes are -// accumulated in FP32 before the residual is loaded and published once. The -// preceding grouped projection is the only producer of each FP16 route row, so -// this boundary requires no global atomics. kernel.def target(@qwen3_moe_routed_down_gfx11_wave64) @qwen3_moe_routed_down_weighted_reduce_f16_f32(%token_count: index) { %output_size = config.get @qwen3_moe.routed_down.output_size : index %c1 = index.constant 1 : index @@ -560,328 +21,23 @@ kernel.def target(@qwen3_moe_routed_down_gfx11_wave64) @qwen3_moe_routed_down_we %output_tiles = index.div %rounded_output_size, %c256 : index kernel.launch.config workgroups(%output_tiles, %token_count, %c1) workgroup_size(%c64, %c1, %c1) : index } launch(%token_count: index, %route_weights: buffer, %routed_output: buffer, %residual_input: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { - %route_count = config.get @qwen3_moe.routed_down.route_count : index - %output_size = config.get @qwen3_moe.routed_down.output_size : index - %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index - %bounded_output_size = index.assume %output_size [range(%output_size, 1, 4096)] : index - %channel_tile = kernel.workgroup.id : index - %token0 = kernel.workgroup.id : index - %lane = kernel.workitem.id : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c4 = index.constant 4 : index - %c256 = index.constant 256 : index - %c0_offset = index.constant 0 : offset - %c0_f16x4 = vector.constant 0.0 : vector<4xf16> - %c0_f32x4 = vector.constant 0.0 : vector<4xf32> - %assignment_count = index.mul %token_count, %bounded_route_count : index - %channel_tile_base = index.mul %channel_tile, %c256 : index - %lane_channel_add = index.mul %lane, %c4 : index - %channel = index.add %channel_tile_base, %lane_channel_add : index - %active_token = index.cmp ult, %token0, %token_count : index - %valid_channel = index.cmp ult, %channel, %bounded_output_size : index - %publishes_output = scalar.andi %active_token, %valid_channel : i1 - %route_weights_noalias, %routed_output_noalias, %residual_input_noalias, %output_noalias = buffer.assume.noalias %route_weights, %routed_output, %residual_input, %output : buffer, buffer, buffer, buffer - %route_weight_view = buffer.view %route_weights_noalias[%c0_offset] : buffer -> view<[%token_count]x[%bounded_route_count]xf32> - %routed_output_view = buffer.view %routed_output_noalias[%c0_offset] : buffer -> view<[%assignment_count]x[%bounded_output_size]xf16> - %residual_input_view = buffer.view %residual_input_noalias[%c0_offset] : buffer -> view<[%token_count]x[%bounded_output_size]xf32> - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%token_count]x[%bounded_output_size]xf32> - scf.if %publishes_output { - %token, %output_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index - %mask = vector.mask.range [%channel to %bounded_output_size step %c1] : index -> vector<4xi1> - %weighted_sum = scf.for %route = [%c0 to %bounded_route_count step %c1](%sum = %c0_f32x4 : vector<4xf32>) -> (vector<4xf32>) unroll { - %assignment = index.madd %token, %bounded_route_count, %route : index - %routed = vector.load.mask %routed_output_view[%assignment, %channel], %mask, %c0_f16x4 : view<[%assignment_count]x[%bounded_output_size]xf16>, vector<4xi1>, vector<4xf16> - %wide = vector.extf %routed : vector<4xf16> to vector<4xf32> - %route_weight = view.load %route_weight_view[%token, %route] : view<[%token_count]x[%bounded_route_count]xf32> -> f32 - %route_weight_x4 = vector.splat %route_weight : vector<4xf32> - %weighted = vector.mulf %wide, %route_weight_x4 : vector<4xf32> - %next = vector.addf %sum, %weighted : vector<4xf32> - scf.yield %next : vector<4xf32> - } - %residual = vector.load.mask %residual_input_view[%token, %channel], %mask, %c0_f32x4 : view<[%token_count]x[%bounded_output_size]xf32>, vector<4xi1>, vector<4xf32> - %result = vector.addf %residual, %weighted_sum : vector<4xf32> - vector.store.mask %result, %output_view[%token, %channel], %mask : vector<4xf32>, view<[%token_count]x[%bounded_output_size]xf32>, vector<4xi1> - } + template.apply<@ggml.routing.weighted_reduce_f16_f32>(%token_count, %route_weights, %routed_output, %residual_input, %output) : (index, buffer, buffer, buffer, buffer) kernel.return } -// Complete expert-MLP coverage crosses route partitions and output tails. The -// established Q8_1 gate/up and Q4_K down providers form an independent -// reference while the production path keeps the routed SwiGLU handoff in F16. -// The complete-chain tolerance covers both independent quantized accumulation -// regimes; the focused gate/up and down cases constrain each boundary to 1%. -check.case public @qwen3_moe_expert_mlp_q4k_f16_handoff_differential_case { - %token_count = check.literal value(67) : index - %gate_input_size = check.literal value(512) : index - %gate_output_size = check.literal value(256) : index - %down_output_size = check.literal value(33) : index - %routed_row_count = check.literal value(134) : index - %route_count = check.literal value(2) : index - %route_stride = check.literal value(4) : index - %expert_count = check.literal value(4) : index - %input = check.generate.fill value(0.00390625) : tensor<67x512xf32> - %q8_input = check.generate.fill value(0) : tensor<67x576xi8> - %route_ids = check.generate.iota offset(0) step(1) period(4) : tensor<67x4xi32> - %route_weights = check.generate.fill value(0.5) : tensor<67x2xf32> - %expert_table = check.generate.fill value(-1) : tensor<272xi32> - %partition_table = check.generate.fill value(-1) : tensor<10xi32> - %gate_weight = check.generate.fill value(34) : tensor<4x256x2x144xi8> - %up_weight = check.generate.fill value(35) : tensor<4x256x2x144xi8> - %down_weight = check.generate.fill value(-86) : tensor<4x33x1x144xi8> - %reference_gate_output = check.generate.fill value(1.0) : tensor<67x2x256xf32> - %q8_down_input = check.generate.fill value(0) : tensor<67x2x288xi8> - %expected_output = check.generate.fill value(1.0) : tensor<67x33xf32> - %actual_gate_output = check.generate.fill value(1.0) : tensor<67x2x256xf16> - %actual_routed_output = check.generate.fill value(0.0) : tensor<67x2x33xf16> - %actual_output = check.generate.fill value(1.0) : tensor<67x33xf32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %gate_input_size](%token_count, %gate_input_size, %input, %q8_input) : [index, index](index, index, tensor<67x512xf32>, tensor<67x576xi8>) - kernel.launch @ggml_mul_mat_id_vector_pair_binary_publish_f32[%token_count, %route_count, %route_stride, %expert_count, %gate_output_size](%token_count, %route_count, %route_stride, %expert_count, %gate_output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %reference_gate_output) : [index, index, index, index, index](index, index, index, index, index, tensor<67x576xi8>, tensor<67x4xi32>, tensor<4x256x2x144xi8>, tensor<4x256x2x144xi8>, tensor<67x2x256xf32>) - kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %gate_output_size](%routed_row_count, %gate_output_size, %reference_gate_output, %q8_down_input) : [index, index](index, index, tensor<67x2x256xf32>, tensor<67x2x288xi8>) - kernel.launch @qwen3_moe_routed_down_q4k_q8_1_x4[%token_count, %gate_output_size, %route_count, %route_stride, %expert_count, %down_output_size](%token_count, %gate_output_size, %route_count, %route_stride, %expert_count, %down_output_size, %q8_down_input, %route_ids, %route_weights, %down_weight, %expected_output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<67x2x288xi8>, tensor<67x4xi32>, tensor<67x2xf32>, tensor<4x33x1x144xi8>, tensor<67x33xf32>) - kernel.launch @ggml_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<67x4xi32>, tensor<272xi32>) - kernel.launch @ggml_moe_build_expert_partition_table[%token_count, %route_count, %expert_count](%token_count, %route_count, %expert_count, %expert_table, %partition_table) : [index, index, index](index, index, index, tensor<272xi32>, tensor<10xi32>) - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma[%token_count](%token_count, %input, %expert_table, %partition_table, %gate_weight, %up_weight, %actual_gate_output) : [index](index, tensor<67x512xf32>, tensor<272xi32>, tensor<10xi32>, tensor<4x256x2x144xi8>, tensor<4x256x2x144xi8>, tensor<67x2x256xf16>) - kernel.launch @qwen3_moe_routed_down_q4k_f16_wmma_grouped[%token_count](%token_count, %actual_gate_output, %expert_table, %down_weight, %actual_routed_output) : [index](index, tensor<67x2x256xf16>, tensor<272xi32>, tensor<4x33x1x144xi8>, tensor<67x2x33xf16>) - kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %actual_routed_output, %actual_output, %actual_output) : [index](index, tensor<67x2xf32>, tensor<67x2x33xf16>, tensor<67x33xf32>, tensor<67x33xf32>) - check.expect.close actual(%actual_output) expected(%expected_output) atol(0.5) rtol(0.05) nan(same) : tensor<67x33xf32> +// Compare separate output storage with the in-place residual/output binding +// used by composed routed-down schedules. +check.case public @qwen3_moe_routed_down_weighted_reduce_f16_f32_alias_case { + %token_count = check.literal value(1) : index + %route_weights = check.generate.iota offset(0.03125) step(0.015625) period(8) : tensor<1x8xf32> + %routed_output = check.generate.iota offset(-0.25) step(0.015625) period(31) : tensor<1x8x2048xf16> + %separate_residual = check.generate.iota offset(-0.5) step(0.0078125) period(17) : tensor<1x2048xf32> + %in_place_residual = check.generate.iota offset(-0.5) step(0.0078125) period(17) : tensor<1x2048xf32> + %separate_output = check.generate.fill value(0.0) : tensor<1x2048xf32> + kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %routed_output, %separate_residual, %separate_output) : [index](index, tensor<1x8xf32>, tensor<1x8x2048xf16>, tensor<1x2048xf32>, tensor<1x2048xf32>) + kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %routed_output, %in_place_residual, %in_place_residual) : [index](index, tensor<1x8xf32>, tensor<1x8x2048xf16>, tensor<1x2048xf32>, tensor<1x2048xf32>) + check.expect.equal actual(%in_place_residual) expected(%separate_output) : tensor<1x2048xf32> check.return } -// Q6_K shares the production matrix schedule but has an independent packed -// decoder and Q8_1 reference provider. -check.case public @qwen3_moe_expert_mlp_q6k_f16_handoff_differential_case { - %token_count = check.literal value(67) : index - %gate_input_size = check.literal value(512) : index - %gate_output_size = check.literal value(256) : index - %down_output_size = check.literal value(33) : index - %routed_row_count = check.literal value(134) : index - %route_count = check.literal value(2) : index - %route_stride = check.literal value(4) : index - %expert_count = check.literal value(4) : index - %input = check.generate.fill value(0.00390625) : tensor<67x512xf32> - %q8_input = check.generate.fill value(0) : tensor<67x576xi8> - %route_ids = check.generate.iota offset(0) step(1) period(4) : tensor<67x4xi32> - %route_weights = check.generate.fill value(0.5) : tensor<67x2xf32> - %expert_table = check.generate.fill value(-1) : tensor<272xi32> - %partition_table = check.generate.fill value(-1) : tensor<10xi32> - %gate_weight = check.generate.fill value(34) : tensor<4x256x2x144xi8> - %up_weight = check.generate.fill value(35) : tensor<4x256x2x144xi8> - %down_weight = check.generate.fill value(-86) : tensor<4x33x1x210xi8> - %reference_gate_output = check.generate.fill value(1.0) : tensor<67x2x256xf32> - %q8_down_input = check.generate.fill value(0) : tensor<67x2x288xi8> - %expected_output = check.generate.fill value(1.0) : tensor<67x33xf32> - %actual_gate_output = check.generate.fill value(1.0) : tensor<67x2x256xf16> - %actual_routed_output = check.generate.fill value(0.0) : tensor<67x2x33xf16> - %actual_output = check.generate.fill value(1.0) : tensor<67x33xf32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %gate_input_size](%token_count, %gate_input_size, %input, %q8_input) : [index, index](index, index, tensor<67x512xf32>, tensor<67x576xi8>) - kernel.launch @ggml_mul_mat_id_vector_pair_binary_publish_f32[%token_count, %route_count, %route_stride, %expert_count, %gate_output_size](%token_count, %route_count, %route_stride, %expert_count, %gate_output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %reference_gate_output) : [index, index, index, index, index](index, index, index, index, index, tensor<67x576xi8>, tensor<67x4xi32>, tensor<4x256x2x144xi8>, tensor<4x256x2x144xi8>, tensor<67x2x256xf32>) - kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %gate_output_size](%routed_row_count, %gate_output_size, %reference_gate_output, %q8_down_input) : [index, index](index, index, tensor<67x2x256xf32>, tensor<67x2x288xi8>) - kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4[%token_count, %gate_output_size, %route_count, %route_stride, %expert_count, %down_output_size](%token_count, %gate_output_size, %route_count, %route_stride, %expert_count, %down_output_size, %q8_down_input, %route_ids, %route_weights, %down_weight, %expected_output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<67x2x288xi8>, tensor<67x4xi32>, tensor<67x2xf32>, tensor<4x33x1x210xi8>, tensor<67x33xf32>) - kernel.launch @ggml_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<67x4xi32>, tensor<272xi32>) - kernel.launch @ggml_moe_build_expert_partition_table[%token_count, %route_count, %expert_count](%token_count, %route_count, %expert_count, %expert_table, %partition_table) : [index, index, index](index, index, index, tensor<272xi32>, tensor<10xi32>) - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma[%token_count](%token_count, %input, %expert_table, %partition_table, %gate_weight, %up_weight, %actual_gate_output) : [index](index, tensor<67x512xf32>, tensor<272xi32>, tensor<10xi32>, tensor<4x256x2x144xi8>, tensor<4x256x2x144xi8>, tensor<67x2x256xf16>) - kernel.launch @qwen3_moe_routed_down_q6k_f16_wmma_grouped[%token_count](%token_count, %actual_gate_output, %expert_table, %down_weight, %actual_routed_output) : [index](index, tensor<67x2x256xf16>, tensor<272xi32>, tensor<4x33x1x210xi8>, tensor<67x2x33xf16>) - kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %actual_routed_output, %actual_output, %actual_output) : [index](index, tensor<67x2xf32>, tensor<67x2x33xf16>, tensor<67x33xf32>, tensor<67x33xf32>) - check.expect.close actual(%actual_output) expected(%expected_output) atol(0.5) rtol(0.05) nan(same) : tensor<67x33xf32> - check.return -} - -// Differential coverage crosses token, route, expert, channel, and output-tail -// boundaries. The direct Q8_1 provider supplies an independent raw-Q4_K -// reference while the grouped provider exercises collision-free FP16 -// publication and a separate weighted FP32 residual reduction. -check.case public @qwen3_moe_routed_down_q4k_f16_wmma_differential_case { - %token_count = check.literal value(67) : index - %input_size = check.literal value(512) : index - %route_count = check.literal value(2) : index - %route_stride = check.literal value(4) : index - %expert_count = check.literal value(4) : index - %output_size = check.literal value(33) : index - %routed_row_count = check.literal value(134) : index - %input_f32 = check.generate.fill value(0.00390625) : tensor<67x2x512xf32> - %input_f16 = check.generate.fill value(0.00390625) : tensor<67x2x512xf16> - %q8_input = check.generate.fill value(0) : tensor<67x2x576xi8> - %route_ids = check.generate.iota offset(0) step(1) period(4) : tensor<67x4xi32> - %route_weights = check.generate.fill value(0.5) : tensor<67x2xf32> - %expert_table = check.generate.fill value(-1) : tensor<272xi32> - %weight = check.generate.fill value(-86) : tensor<4x33x2x144xi8> - %expected = check.generate.fill value(1.0) : tensor<67x33xf32> - %routed_output = check.generate.fill value(0.0) : tensor<67x2x33xf16> - %actual = check.generate.fill value(1.0) : tensor<67x33xf32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %input_size](%routed_row_count, %input_size, %input_f32, %q8_input) : [index, index](index, index, tensor<67x2x512xf32>, tensor<67x2x576xi8>) - kernel.launch @qwen3_moe_routed_down_q4k_q8_1_x4[%token_count, %input_size, %route_count, %route_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %expected) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<67x2x576xi8>, tensor<67x4xi32>, tensor<67x2xf32>, tensor<4x33x2x144xi8>, tensor<67x33xf32>) - kernel.launch @ggml_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<67x4xi32>, tensor<272xi32>) - kernel.launch @qwen3_moe_routed_down_q4k_f16_wmma_grouped[%token_count](%token_count, %input_f16, %expert_table, %weight, %routed_output) : [index](index, tensor<67x2x512xf16>, tensor<272xi32>, tensor<4x33x2x144xi8>, tensor<67x2x33xf16>) - kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %routed_output, %actual, %actual) : [index](index, tensor<67x2xf32>, tensor<67x2x33xf16>, tensor<67x33xf32>, tensor<67x33xf32>) - check.expect.close actual(%actual) expected(%expected) atol(0.5) rtol(0.01) nan(same) : tensor<67x33xf32> - check.return -} - -// Q6_K follows the identical routed matrix schedule while selecting its -// independent packed-row decoder. The same irregular shape prevents format -// specialization from weakening routing or tail coverage. -check.case public @qwen3_moe_routed_down_q6k_f16_wmma_differential_case { - %token_count = check.literal value(67) : index - %input_size = check.literal value(512) : index - %route_count = check.literal value(2) : index - %route_stride = check.literal value(4) : index - %expert_count = check.literal value(4) : index - %output_size = check.literal value(33) : index - %routed_row_count = check.literal value(134) : index - %input_f32 = check.generate.fill value(0.00390625) : tensor<67x2x512xf32> - %input_f16 = check.generate.fill value(0.00390625) : tensor<67x2x512xf16> - %q8_input = check.generate.fill value(0) : tensor<67x2x576xi8> - %route_ids = check.generate.iota offset(0) step(1) period(4) : tensor<67x4xi32> - %route_weights = check.generate.fill value(0.5) : tensor<67x2xf32> - %expert_table = check.generate.fill value(-1) : tensor<272xi32> - %weight = check.generate.fill value(-86) : tensor<4x33x2x210xi8> - %expected = check.generate.fill value(1.0) : tensor<67x33xf32> - %routed_output = check.generate.fill value(0.0) : tensor<67x2x33xf16> - %actual = check.generate.fill value(1.0) : tensor<67x33xf32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %input_size](%routed_row_count, %input_size, %input_f32, %q8_input) : [index, index](index, index, tensor<67x2x512xf32>, tensor<67x2x576xi8>) - kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4[%token_count, %input_size, %route_count, %route_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %expected) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<67x2x576xi8>, tensor<67x4xi32>, tensor<67x2xf32>, tensor<4x33x2x210xi8>, tensor<67x33xf32>) - kernel.launch @ggml_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<67x4xi32>, tensor<272xi32>) - kernel.launch @qwen3_moe_routed_down_q6k_f16_wmma_grouped[%token_count](%token_count, %input_f16, %expert_table, %weight, %routed_output) : [index](index, tensor<67x2x512xf16>, tensor<272xi32>, tensor<4x33x2x210xi8>, tensor<67x2x33xf16>) - kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %routed_output, %actual, %actual) : [index](index, tensor<67x2xf32>, tensor<67x2x33xf16>, tensor<67x33xf32>, tensor<67x33xf32>) - check.expect.close actual(%actual) expected(%expected) atol(0.5) rtol(0.01) nan(same) : tensor<67x33xf32> - check.return -} - -check.case public @qwen3_moe_routed_down_q4k_f16_wmma_benchmark_case { - %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512]) name("token_count") : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(128) : index - %expert_count = check.literal value(128) : index - %route_ids = check.generate.iota offset(0) step(1) period(127) : tensor<[%token_count]x128xi32> - %route_weights = check.generate.fill value(0.125) : tensor<[%token_count]x8xf32> - %expert_table = check.generate.fill value(-1) : tensor<65664xi32> - %input = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf16> - %weight = check.generate.fill value(0) : tensor<128x2048x3x144xi8> - %routed_output = check.generate.fill value(0.0) : tensor<[%token_count]x8x2048xf16> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - %expected = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - kernel.launch @ggml_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x128xi32>, tensor<65664xi32>) - kernel.launch @qwen3_moe_routed_down_q4k_f16_wmma_grouped[%token_count](%token_count, %input, %expert_table, %weight, %routed_output) : [index](index, tensor<[%token_count]x8x768xf16>, tensor<65664xi32>, tensor<128x2048x3x144xi8>, tensor<[%token_count]x8x2048xf16>) - kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %routed_output, %output, %output) : [index](index, tensor<[%token_count]x8xf32>, tensor<[%token_count]x8x2048xf16>, tensor<[%token_count]x2048xf32>, tensor<[%token_count]x2048xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> - check.return -} - -check.case public @qwen3_moe_routed_down_q6k_f16_wmma_benchmark_case { - %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512]) name("token_count") : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(128) : index - %expert_count = check.literal value(128) : index - %route_ids = check.generate.iota offset(0) step(1) period(127) : tensor<[%token_count]x128xi32> - %route_weights = check.generate.fill value(0.125) : tensor<[%token_count]x8xf32> - %expert_table = check.generate.fill value(-1) : tensor<65664xi32> - %input = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf16> - %weight = check.generate.fill value(0) : tensor<128x2048x3x210xi8> - %routed_output = check.generate.fill value(0.0) : tensor<[%token_count]x8x2048xf16> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - %expected = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - kernel.launch @ggml_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x128xi32>, tensor<65664xi32>) - kernel.launch @qwen3_moe_routed_down_q6k_f16_wmma_grouped[%token_count](%token_count, %input, %expert_table, %weight, %routed_output) : [index](index, tensor<[%token_count]x8x768xf16>, tensor<65664xi32>, tensor<128x2048x3x210xi8>, tensor<[%token_count]x8x2048xf16>) - kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %routed_output, %output, %output) : [index](index, tensor<[%token_count]x8xf32>, tensor<[%token_count]x8x2048xf16>, tensor<[%token_count]x2048xf32>, tensor<[%token_count]x2048xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> - check.return -} - -// Maximally diverse decode-batch controls. Compact route storage makes the -// flattened 0..127 iota assign every M=16 route to a different expert. -check.case public @qwen3_moe_routed_down_q4k_f16_wmma_diverse_benchmark_case { - %token_count = check.param.choice values([1, 2, 4, 8, 16]) name("token_count") : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<[%token_count]x8xi32> - %route_weights = check.generate.fill value(0.125) : tensor<[%token_count]x8xf32> - %expert_table = check.generate.fill value(-1) : tensor<65664xi32> - %input = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf16> - %weight = check.generate.fill value(0) : tensor<128x2048x3x144xi8> - %routed_output = check.generate.fill value(0.0) : tensor<[%token_count]x8x2048xf16> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - %expected = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - kernel.launch @ggml_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x8xi32>, tensor<65664xi32>) - kernel.launch @qwen3_moe_routed_down_q4k_f16_wmma_grouped[%token_count](%token_count, %input, %expert_table, %weight, %routed_output) : [index](index, tensor<[%token_count]x8x768xf16>, tensor<65664xi32>, tensor<128x2048x3x144xi8>, tensor<[%token_count]x8x2048xf16>) - kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %routed_output, %output, %output) : [index](index, tensor<[%token_count]x8xf32>, tensor<[%token_count]x8x2048xf16>, tensor<[%token_count]x2048xf32>, tensor<[%token_count]x2048xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> - check.return -} - -check.case public @qwen3_moe_routed_down_q6k_f16_wmma_diverse_benchmark_case { - %token_count = check.param.choice values([1, 2, 4, 8, 16]) name("token_count") : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<[%token_count]x8xi32> - %route_weights = check.generate.fill value(0.125) : tensor<[%token_count]x8xf32> - %expert_table = check.generate.fill value(-1) : tensor<65664xi32> - %input = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf16> - %weight = check.generate.fill value(0) : tensor<128x2048x3x210xi8> - %routed_output = check.generate.fill value(0.0) : tensor<[%token_count]x8x2048xf16> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - %expected = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - kernel.launch @ggml_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x8xi32>, tensor<65664xi32>) - kernel.launch @qwen3_moe_routed_down_q6k_f16_wmma_grouped[%token_count](%token_count, %input, %expert_table, %weight, %routed_output) : [index](index, tensor<[%token_count]x8x768xf16>, tensor<65664xi32>, tensor<128x2048x3x210xi8>, tensor<[%token_count]x8x2048xf16>) - kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %routed_output, %output, %output) : [index](index, tensor<[%token_count]x8xf32>, tensor<[%token_count]x8x2048xf16>, tensor<[%token_count]x2048xf32>, tensor<[%token_count]x2048xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> - check.return -} - -check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_differential_case> @qwen3_moe_routed_down_q4k_f16_wmma_small - -check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_decode {token_count = 1} - -check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_small_batch_2 {token_count = 2} - -check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_small_batch_4 {token_count = 4} - -check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_small_batch_8 {token_count = 8} - -check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_small_batch_16 {token_count = 16} - -check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_diverse_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_diverse_decode {token_count = 1} - -check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_diverse_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_diverse_small_batch_2 {token_count = 2} - -check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_diverse_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_diverse_small_batch_4 {token_count = 4} - -check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_diverse_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_diverse_small_batch_8 {token_count = 8} - -check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_diverse_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_diverse_small_batch_16 {token_count = 16} - -check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_prefill_32 {token_count = 32} - -check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_prefill_128 {token_count = 128} - -check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_prefill_512 {token_count = 512} - -check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_differential_case> @qwen3_moe_routed_down_q6k_f16_wmma_small - -check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_decode {token_count = 1} - -check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_small_batch_2 {token_count = 2} - -check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_small_batch_4 {token_count = 4} - -check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_small_batch_8 {token_count = 8} - -check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_small_batch_16 {token_count = 16} - -check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_diverse_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_diverse_decode {token_count = 1} - -check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_diverse_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_diverse_small_batch_2 {token_count = 2} - -check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_diverse_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_diverse_small_batch_4 {token_count = 4} - -check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_diverse_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_diverse_small_batch_8 {token_count = 8} - -check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_diverse_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_diverse_small_batch_16 {token_count = 16} - -check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_prefill_32 {token_count = 32} - -check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_prefill_128 {token_count = 128} - -check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_prefill_512 {token_count = 512} +check.benchmark<@qwen3_moe_routed_down_weighted_reduce_f16_f32_alias_case> @qwen3_moe_routed_down_weighted_reduce_f16_f32_alias diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom index cf71ac51b464..0c2850abc6e8 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom @@ -5,6 +5,8 @@ // residual results across the workgroup RMS reduction, then publishes both // the updated hidden state and the learned-weight normalized input consumed // by the next layer. +template.decl @ggml.routing.weighted_reduce_next_rmsnorm_f32(%token_count: index, %route_weights: buffer, %routed_output: buffer, %hidden_state: buffer, %next_norm_weight: buffer, %next_projection_input: buffer) + amdgpu.target @qwen3_moe_routed_down_next_norm_gfx11_wave32 {subgroup_size = 32} config.decl @qwen3_moe.model.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] @@ -24,92 +26,7 @@ kernel.def target(@qwen3_moe_routed_down_next_norm_gfx11_wave32) @qwen3_moe_rout %c256 = index.constant 256 : index kernel.launch.config workgroups(%token_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index } launch(%token_count: index, %route_weights: buffer, %routed_output: buffer, %hidden_state: buffer, %next_norm_weight: buffer, %next_projection_input: buffer) where [range(%token_count, 1, 2048)] { - %route_count0 = config.get @qwen3_moe.routed_down.route_count : index - %output_size0 = config.get @qwen3_moe.routed_down.output_size : index - %hidden_size0 = config.get @qwen3_moe.model.hidden_size : index - %epsilon = config.get @qwen3_moe.model.rms_epsilon : f32 - %route_count = index.assume %route_count0 [range(%route_count0, 8, 8)] : index - %output_size, %hidden_size = index.assume %output_size0, %hidden_size0 [range(%output_size0, 2048, 2048), range(%hidden_size0, 2048, 2048), eq(%output_size0, %hidden_size0)] : index, index - %token0 = kernel.workgroup.id : index - %workitem0 = kernel.workitem.id : index - %workitem = index.assume %workitem0 [range(%workitem0, 0, 255)] : index - %subgroup0 = kernel.subgroup.id : index - %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index - %lane = kernel.subgroup.lane.id : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c8 = index.constant 8 : index - %channel = index.mul %workitem, %c8 : index - %scratch_bytes = index.constant 32 : offset - %c0_offset = index.constant 0 : offset - %c0_f32 = scalar.constant 0.0 : f32 - %c0_f32x8 = vector.constant 0.0 : vector<8xf32> - %active_token = index.cmp ult, %token0, %token_count : index - %safe_token0 = scf.select %active_token, %token0, %c0 : index - %token, %launch_token_count = index.assume %safe_token0, %token_count [lt(%safe_token0, %token_count)] : index, index - %hidden_size_i32 = index.cast %hidden_size : index to i32 - %hidden_size_f32 = scalar.sitofp %hidden_size_i32 : i32 to f32 - %assignment_count = index.mul %launch_token_count, %route_count : index - %route_weights_noalias, %routed_output_noalias, %hidden_state_noalias, %next_norm_weight_noalias, %next_projection_input_noalias = buffer.assume.noalias %route_weights, %routed_output, %hidden_state, %next_norm_weight, %next_projection_input : buffer, buffer, buffer, buffer, buffer - %route_weights_view = buffer.view %route_weights_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%route_count]xf32> - %routed_output_view = buffer.view %routed_output_noalias[%c0_offset] : buffer -> view<[%assignment_count]x[%output_size]xf16> - %hidden_state_view = buffer.view %hidden_state_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> - %next_norm_weight_view = buffer.view %next_norm_weight_noalias[%c0_offset] : buffer -> view<[%hidden_size]xf32> - %next_projection_input_view = buffer.view %next_projection_input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> - %scratch = buffer.alloca align(16) %scratch_bytes : buffer - %scratch_view = buffer.view %scratch[%c0_offset] : buffer -> view<8xf32> - - // The predicate depends only on the workgroup ID, so every workitem either - // executes the complete barrier-bearing row reduction or skips it together. - scf.if %active_token { - %weighted_sum = scf.for %route = [%c0 to %route_count step %c1](%sum = %c0_f32x8 : vector<8xf32>) -> (vector<8xf32>) unroll { - %assignment = index.madd %token, %route_count, %route : index - %routed_f16 = vector.load %routed_output_view[%assignment, %channel] : view<[%assignment_count]x[%output_size]xf16> -> vector<8xf16> - %routed = vector.extf %routed_f16 : vector<8xf16> to vector<8xf32> - %route_weight = view.load %route_weights_view[%token, %route] : view<[%launch_token_count]x[%route_count]xf32> -> f32 - %route_weight_x8 = vector.splat %route_weight : vector<8xf32> - %weighted = vector.mulf %routed, %route_weight_x8 : vector<8xf32> - %next = vector.addf %sum, %weighted : vector<8xf32> - scf.yield %next : vector<8xf32> - } - %residual = vector.load %hidden_state_view[%token, %channel] : view<[%launch_token_count]x[%hidden_size]xf32> -> vector<8xf32> - %result = vector.addf %residual, %weighted_sum : vector<8xf32> - %squares = vector.mulf %result, %result : vector<8xf32> - %thread_sum = vector.reduce %squares, %c0_f32 : vector<8xf32>, f32 - %subgroup_sum = kernel.subgroup.reduce %thread_sum : f32 - - %is_subgroup_leader = index.cmp eq, %lane, %c0 : index - scf.if %is_subgroup_leader { - view.store %subgroup_sum, %scratch_view[%subgroup] : f32, view<8xf32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %is_reduction_subgroup = index.cmp eq, %subgroup, %c0 : index - %is_reduction_lane = index.cmp ult, %lane, %c8 : index - %loads_subgroup_sum = scalar.andi %is_reduction_subgroup, %is_reduction_lane : i1 - %subgroup_partial = scf.if %loads_subgroup_sum -> (f32) { - %value = view.load %scratch_view[%lane] : view<8xf32> -> f32 - scf.yield %value : f32 - } else { - scf.yield %c0_f32 : f32 - } - %row_sum = kernel.subgroup.reduce %subgroup_partial : f32 - %writes_scale = scalar.andi %is_reduction_subgroup, %is_subgroup_leader : i1 - scf.if %writes_scale { - %mean = scalar.divf %row_sum, %hidden_size_f32 : f32 - %biased_mean = scalar.addf %mean, %epsilon : f32 - %scale = scalar.rsqrtf %biased_mean : f32 - view.store %scale, %scratch_view[%c0] : f32, view<8xf32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - - %scale = view.load %scratch_view[%c0] : view<8xf32> -> f32 - %scale_x8 = vector.splat %scale : vector<8xf32> - %learned_weight = vector.load %next_norm_weight_view[%channel] : view<[%hidden_size]xf32> -> vector<8xf32> - %normalized = vector.mulf %result, %scale_x8 : vector<8xf32> - %next_projection = vector.mulf %normalized, %learned_weight : vector<8xf32> - vector.store %result, %hidden_state_view[%token, %channel] : vector<8xf32>, view<[%launch_token_count]x[%hidden_size]xf32> - vector.store %next_projection, %next_projection_input_view[%token, %channel] : vector<8xf32>, view<[%launch_token_count]x[%hidden_size]xf32> - } + template.apply<@ggml.routing.weighted_reduce_next_rmsnorm_f32>(%token_count, %route_weights, %routed_output, %hidden_state, %next_norm_weight, %next_projection_input) : (index, buffer, buffer, buffer, buffer, buffer) kernel.return } @@ -158,3 +75,5 @@ check.case public @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_compos check.benchmark<@qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_fused_benchmark_case> @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_fused_prefill_512 check.benchmark<@qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_composed_benchmark_case> @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_composed_prefill_512 + +check.benchmark<@qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_differential_case> @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_differential diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_q8_1_x4.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_q8_1_x4.loom deleted file mode 100644 index 9d452beccb71..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_q8_1_x4.loom +++ /dev/null @@ -1,269 +0,0 @@ -// Fused grouped-prefill routed residual reduction, next-layer RMSNorm, and -// GGML Q8_1 x4 publication. -// -// One 256-workitem workgroup, comprising eight wave32 subgroups, owns one -// complete 2048-channel token row. Each workitem retains eight adjacent -// residual results across the RMS reduction, updates hidden state, and -// quantizes the learned-weight normalized values directly into the physical -// row consumed by quantized QKV. This is the interlayer producer for the -// quantized attention schedule; no F32 normalized row is materialized. -// -// Four neighboring workitems own one logical 32-element Q8_1 block. Each -// workitem contributes two packed four-value words. LDS preserves the same -// eight-word max and sum reduction order as the standalone GGML packer so the -// fused and decomposed paths produce identical metadata and payload bytes. -amdgpu.target @qwen3_moe_routed_down_next_norm_q8_gfx11_wave32 {subgroup_size = 32} - -config.decl @qwen3_moe.model.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] - -config.decl @qwen3_moe.model.rms_epsilon : f32 - -config.decl @qwen3_moe.routed_down.route_count : %value: index where [range(%value, 1, 8)] - -config.decl @qwen3_moe.routed_down.output_size : %value: index where [range(%value, 1, 4096)] - -config.decl @qwen3_moe.workload.token_capacity : %value: index where [range(%value, 1, 2048)] - -kernel.decl @qwen3_moe_routed_down_weighted_reduce_f16_f32(%token_count$5: index) launch(%token_count$6: index, %route_weights: buffer, %routed_output: buffer, %residual_input: buffer, %output: buffer) - -kernel.decl @qwen3_moe_rmsnorm_f32(%token_count$10: index) launch(%token_count$11: index, %input: buffer, %weight: buffer, %output: buffer) - -kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count$15: index, %input_size$16: index) launch(%token_count$17: index, %input_size$18: index, %input: buffer, %output: buffer) - -kernel.def target(@qwen3_moe_routed_down_next_norm_q8_gfx11_wave32) @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4(%token_count: index) { - %token_capacity = config.get @qwen3_moe.workload.token_capacity : index - %c1 = index.constant 1 : index - %c256 = index.constant 256 : index - kernel.launch.config workgroups(%token_capacity, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index -} launch(%token_count: index, %route_weights: buffer, %routed_output: buffer, %hidden_state: buffer, %next_norm_weight: buffer, %next_projection_input: buffer) { - %route_count0 = config.get @qwen3_moe.routed_down.route_count : index - %output_size0 = config.get @qwen3_moe.routed_down.output_size : index - %hidden_size0 = config.get @qwen3_moe.model.hidden_size : index - %epsilon = config.get @qwen3_moe.model.rms_epsilon : f32 - %token_capacity = config.get @qwen3_moe.workload.token_capacity : index - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index - %route_count = index.assume %route_count0 [range(%route_count0, 8, 8)] : index - %output_size, %hidden_size = index.assume %output_size0, %hidden_size0 [range(%output_size0, 2048, 2048), range(%hidden_size0, 2048, 2048), eq(%output_size0, %hidden_size0)] : index, index - %token0 = kernel.workgroup.id : index - %workitem0 = kernel.workitem.id : index - %workitem = index.assume %workitem0 [range(%workitem0, 0, 255)] : index - %subgroup0 = kernel.subgroup.id : index - %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index - %lane = kernel.subgroup.lane.id : index - %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 - %c128 = index.constant 128 : index - %group_bytes = index.constant 144 : offset - %payload_byte_add = index.constant 16 : offset - %scratch_d_byte_add = index.constant 2048 : offset - %scratch_bytes = index.constant 2304 : offset - %c0_offset = index.constant 0 : offset - %c0_f32 = scalar.constant 0.0 : f32 - %c1_f32 = scalar.constant 1.0 : f32 - %c127 = scalar.constant 127.0 : f32 - %c0_f32x8 = vector.constant 0.0 : vector<8xf32> - %active_token = index.cmp ult, %token0, %bounded_token_count : index - %safe_token0 = scf.select %active_token, %token0, %c0 : index - %token, %launch_token_count = index.assume %safe_token0, %bounded_token_count [lt(%safe_token0, %bounded_token_count)] : index, index - %hidden_size_i32 = index.cast %hidden_size : index to i32 - %hidden_size_f32 = scalar.sitofp %hidden_size_i32 : i32 to f32 - %assignment_count = index.mul %launch_token_count, %route_count : index - %physical_group_count = index.div %hidden_size, %c128 : index - %row_bytes = index.scale %physical_group_count, %group_bytes : index, offset -> offset - %token_output_byte_base = index.scale %token, %row_bytes : index, offset -> offset - %route_weights_noalias, %routed_output_noalias, %hidden_state_noalias, %next_norm_weight_noalias, %next_projection_input_noalias = buffer.assume.noalias %route_weights, %routed_output, %hidden_state, %next_norm_weight, %next_projection_input : buffer, buffer, buffer, buffer, buffer - %route_weights_view = buffer.view %route_weights_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%route_count]xf32> - %routed_output_view = buffer.view %routed_output_noalias[%c0_offset] : buffer -> view<[%assignment_count]x[%output_size]xf16> - %hidden_state_view = buffer.view %hidden_state_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> - %next_norm_weight_view = buffer.view %next_norm_weight_noalias[%c0_offset] : buffer -> view<[%hidden_size]xf32> - %scratch = buffer.alloca align(16) %scratch_bytes : buffer - %scratch_values = buffer.view %scratch[%c0_offset] : buffer -> view<512xf32> - %scratch_d = buffer.view %scratch[%scratch_d_byte_add] : buffer -> view<64xf32> - // The token predicate is workgroup-uniform, keeping every RMSNorm and Q8_1 - // barrier either active for the whole workgroup or skipped by all workitems. - scf.if %active_token { - %channel = index.mul %workitem, %c8 : index - %weighted_sum = scf.for %route = [%c0 to %route_count step %c1](%sum = %c0_f32x8 : vector<8xf32>) -> (vector<8xf32>) unroll { - %assignment = index.madd %token, %route_count, %route : index - %routed_f16 = vector.load %routed_output_view[%assignment, %channel] : view<[%assignment_count]x[%output_size]xf16> -> vector<8xf16> - %routed = vector.extf %routed_f16 : vector<8xf16> to vector<8xf32> - %route_weight = view.load %route_weights_view[%token, %route] : view<[%launch_token_count]x[%route_count]xf32> -> f32 - %route_weight_x8 = vector.splat %route_weight : vector<8xf32> - %weighted = vector.mulf %routed, %route_weight_x8 : vector<8xf32> - %next = vector.addf %sum, %weighted : vector<8xf32> - scf.yield %next : vector<8xf32> - } - %residual = vector.load %hidden_state_view[%token, %channel] : view<[%launch_token_count]x[%hidden_size]xf32> -> vector<8xf32> - %result = vector.addf %residual, %weighted_sum : vector<8xf32> - %squares = vector.mulf %result, %result : vector<8xf32> - %thread_sum = vector.reduce %squares, %c0_f32 : vector<8xf32>, f32 - %subgroup_sum = kernel.subgroup.reduce %thread_sum : f32 - - %is_subgroup_leader = index.cmp eq, %lane, %c0 : index - scf.if %is_subgroup_leader { - view.store %subgroup_sum, %scratch_values[%subgroup] : f32, view<512xf32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %is_reduction_subgroup = index.cmp eq, %subgroup, %c0 : index - %is_reduction_lane = index.cmp ult, %lane, %c8 : index - %loads_subgroup_sum = scalar.andi %is_reduction_subgroup, %is_reduction_lane : i1 - %subgroup_partial = scf.if %loads_subgroup_sum -> (f32) { - %value = view.load %scratch_values[%lane] : view<512xf32> -> f32 - scf.yield %value : f32 - } else { - scf.yield %c0_f32 : f32 - } - %row_sum = kernel.subgroup.reduce %subgroup_partial : f32 - %writes_scale = scalar.andi %is_reduction_subgroup, %is_subgroup_leader : i1 - scf.if %writes_scale { - %mean = scalar.divf %row_sum, %hidden_size_f32 : f32 - %biased_mean = scalar.addf %mean, %epsilon : f32 - %scale = scalar.rsqrtf %biased_mean : f32 - view.store %scale, %scratch_values[%c0] : f32, view<512xf32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - - %row_scale = view.load %scratch_values[%c0] : view<512xf32> -> f32 - // All waves must capture the row scale before scratch is reused for the - // independent Q8 block reductions. - kernel.barrier scope(workgroup) ordering(acq_rel) - %row_scale_vector = vector.splat %row_scale : vector<8xf32> - %learned_weight = vector.load %next_norm_weight_view[%channel] : view<[%hidden_size]xf32> -> vector<8xf32> - %normalized0 = vector.mulf %result, %row_scale_vector : vector<8xf32> - %normalized = vector.mulf %normalized0, %learned_weight : vector<8xf32> - vector.store %result, %hidden_state_view[%token, %channel] : vector<8xf32>, view<[%launch_token_count]x[%hidden_size]xf32> - - %low_values = vector.slice %normalized[0] : vector<8xf32> -> vector<4xf32> - %high_values = vector.slice %normalized[4] : vector<8xf32> -> vector<4xf32> - %low_absolute_values = vector.absf %low_values : vector<4xf32> - %high_absolute_values = vector.absf %high_values : vector<4xf32> - %low_max = vector.reduce %low_absolute_values, %c0_f32 : vector<4xf32>, f32 - %high_max = vector.reduce %high_absolute_values, %c0_f32 : vector<4xf32>, f32 - %partial_base = index.mul %workitem, %c2 : index - %partial_high = index.add %partial_base, %c1 : index - view.store %low_max, %scratch_values[%partial_base] : f32, view<512xf32> - view.store %high_max, %scratch_values[%partial_high] : f32, view<512xf32> - kernel.barrier scope(workgroup) ordering(acq_rel) - - %block_in_row = index.div %workitem, %c4 : index - %workitem_in_block = index.rem %workitem, %c4 : index - %is_block_leader = index.cmp eq, %workitem_in_block, %c0 : index - scf.if %is_block_leader { - %block_partial_base = index.mul %block_in_row, %c8 : index - %block_maxima = vector.load %scratch_values[%block_partial_base] : view<512xf32> -> vector<8xf32> - %amax = vector.reduce %block_maxima, %c0_f32 : vector<8xf32>, f32 - %d = scalar.divf %amax, %c127 : f32 - view.store %d, %scratch_d[%block_in_row] : f32, view<64xf32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - - %d = view.load %scratch_d[%block_in_row] : view<64xf32> -> f32 - %d_nonzero = scalar.cmpf one, %d, %c0_f32 : f32 - %d_inverse = scf.if %d_nonzero -> (f32) { - %inverse = scalar.divf %c1_f32, %d : f32 - scf.yield %inverse : f32 - } else { - scf.yield %c0_f32 : f32 - } - %d_inverse_vector = vector.splat %d_inverse : vector<8xf32> - %scaled_values = vector.mulf %normalized, %d_inverse_vector : vector<8xf32> - %rounded_values = vector.roundf %scaled_values : vector<8xf32> - %quantized_values = vector.fptosi %rounded_values : vector<8xf32> to vector<8xi8> - %packed_words = vector.bitcast %quantized_values : vector<8xi8> to vector<2xi32> - %physical_group = index.div %block_in_row, %c4 : index - %block_in_group = index.rem %block_in_row, %c4 : index - %group_byte_add = index.scale %physical_group, %group_bytes : index, offset -> offset - %group_byte_offset = index.add %token_output_byte_base, %group_byte_add : offset - %payload_byte_offset = index.add %group_byte_offset, %payload_byte_add : offset - %group_ds = buffer.view %next_projection_input_noalias[%group_byte_offset] : buffer -> view<8xf16> - %group_qs = buffer.view %next_projection_input_noalias[%payload_byte_offset] : buffer -> view<32xi32> - %block_word_base = index.mul %block_in_group, %c8 : index - %workitem_word_add = index.mul %workitem_in_block, %c2 : index - %packed_word_index0 = index.add %block_word_base, %workitem_word_add : index - %packed_word_index = index.assume %packed_word_index0 [range(%packed_word_index0, 0, 31)] : index - vector.store %packed_words, %group_qs[%packed_word_index] : vector<2xi32>, view<32xi32> - - %low_rounded_values = vector.slice %rounded_values[0] : vector<8xf32> -> vector<4xf32> - %high_rounded_values = vector.slice %rounded_values[4] : vector<8xf32> -> vector<4xf32> - %low_quantized_sum = vector.reduce %low_rounded_values, %c0_f32 : vector<4xf32>, f32 - %high_quantized_sum = vector.reduce %high_rounded_values, %c0_f32 : vector<4xf32>, f32 - view.store %low_quantized_sum, %scratch_values[%partial_base] : f32, view<512xf32> - view.store %high_quantized_sum, %scratch_values[%partial_high] : f32, view<512xf32> - kernel.barrier scope(workgroup) ordering(acq_rel) - scf.if %is_block_leader { - %block_partial_base = index.mul %block_in_row, %c8 : index - %block_sums = vector.load %scratch_values[%block_partial_base] : view<512xf32> -> vector<8xf32> - %quantized_sum = vector.reduce %block_sums, %c0_f32 : vector<8xf32>, f32 - %s = scalar.mulf %quantized_sum, %d : f32 - %d_f16 = scalar.fptrunc %d : f32 to f16 - %s_f16 = scalar.fptrunc %s : f32 to f16 - %ds_index = index.mul %block_in_group, %c2 : index - %s_index = index.add %ds_index, %c1 : index - view.store %d_f16, %group_ds[%ds_index] : f16, view<8xf16> - view.store %s_f16, %group_ds[%s_index] : f16, view<8xf16> - } - } - kernel.return -} - -// Fourteen independent token rows lock the complete interlayer contract. The -// fused producer must match ordinary weighted reduction for hidden state and -// RMSNorm followed by the generic GGML packer for every Q8 metadata and -// payload byte. -check.case public @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_differential_case { - %token_count = check.literal value(14) : index - %hidden_size = check.literal value(2048) : index - %route_weights = check.generate.iota offset(0.03125) step(0.015625) period(8) : tensor<14x8xf32> - %routed_seed = check.param.seed base(5858425849430430008) count(1) : i64 - %routed_output = check.generate.random.uniform seed(%routed_seed) range(-0.5 to 0.5) : tensor<14x8x2048xf16> - %expected_hidden_state = check.generate.iota offset(-0.5) step(0.0078125) period(31) : tensor<14x2048xf32> - %actual_hidden_state = check.generate.iota offset(-0.5) step(0.0078125) period(31) : tensor<14x2048xf32> - %next_norm_weight = check.generate.iota offset(0.5) step(0.0078125) period(29) : tensor<2048xf32> - %normalized = check.generate.fill value(0.0) : tensor<14x2048xf32> - %expected_q8 = check.generate.fill value(0) : tensor<14x2304xi8> - %actual_q8 = check.generate.fill value(1) : tensor<14x2304xi8> - kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %routed_output, %expected_hidden_state, %expected_hidden_state) : [index](index, tensor<14x8xf32>, tensor<14x8x2048xf16>, tensor<14x2048xf32>, tensor<14x2048xf32>) - kernel.launch @qwen3_moe_rmsnorm_f32[%token_count](%token_count, %expected_hidden_state, %next_norm_weight, %normalized) : [index](index, tensor<14x2048xf32>, tensor<2048xf32>, tensor<14x2048xf32>) - kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %hidden_size](%token_count, %hidden_size, %normalized, %expected_q8) : [index, index](index, index, tensor<14x2048xf32>, tensor<14x2304xi8>) - kernel.launch @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4[%token_count](%token_count, %route_weights, %routed_output, %actual_hidden_state, %next_norm_weight, %actual_q8) : [index](index, tensor<14x8xf32>, tensor<14x8x2048xf16>, tensor<14x2048xf32>, tensor<2048xf32>, tensor<14x2304xi8>) - check.expect.close actual(%actual_hidden_state) expected(%expected_hidden_state) atol(0.0001) rtol(0.0001) nan(same) : tensor<14x2048xf32> - check.expect.equal actual(%actual_q8) expected(%expected_q8) : tensor<14x2304xi8> - check.return -} - -check.case public @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_fused_benchmark_case { - %token_count = check.param.choice values([14, 32, 128, 512]) name("token_count") : index - %route_weights = check.generate.fill value(0.125) : tensor<[%token_count]x8xf32> - %routed_output = check.generate.fill value(0.0) : tensor<[%token_count]x8x2048xf16> - %hidden_state = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - %next_norm_weight = check.generate.fill value(1.0) : tensor<2048xf32> - %next_projection_input = check.generate.fill value(0) : tensor<[%token_count]x2304xi8> - kernel.launch @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4[%token_count](%token_count, %route_weights, %routed_output, %hidden_state, %next_norm_weight, %next_projection_input) : [index](index, tensor<[%token_count]x8xf32>, tensor<[%token_count]x8x2048xf16>, tensor<[%token_count]x2048xf32>, tensor<2048xf32>, tensor<[%token_count]x2304xi8>) - check.return -} - -check.case public @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_composed_benchmark_case { - %token_count = check.param.choice values([14, 32, 128, 512]) name("token_count") : index - %hidden_size = check.literal value(2048) : index - %route_weights = check.generate.fill value(0.125) : tensor<[%token_count]x8xf32> - %routed_output = check.generate.fill value(0.0) : tensor<[%token_count]x8x2048xf16> - %hidden_state = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - %next_norm_weight = check.generate.fill value(1.0) : tensor<2048xf32> - %normalized = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> - %next_projection_input = check.generate.fill value(0) : tensor<[%token_count]x2304xi8> - kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %routed_output, %hidden_state, %hidden_state) : [index](index, tensor<[%token_count]x8xf32>, tensor<[%token_count]x8x2048xf16>, tensor<[%token_count]x2048xf32>, tensor<[%token_count]x2048xf32>) - kernel.launch @qwen3_moe_rmsnorm_f32[%token_count](%token_count, %hidden_state, %next_norm_weight, %normalized) : [index](index, tensor<[%token_count]x2048xf32>, tensor<2048xf32>, tensor<[%token_count]x2048xf32>) - kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %hidden_size](%token_count, %hidden_size, %normalized, %next_projection_input) : [index, index](index, index, tensor<[%token_count]x2048xf32>, tensor<[%token_count]x2304xi8>) - check.return -} - -check.benchmark<@qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_fused_benchmark_case> @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_fused_prefill_14 {token_count = 14} - -check.benchmark<@qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_fused_benchmark_case> @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_fused_prefill_512 {token_count = 512} - -check.benchmark<@qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_composed_benchmark_case> @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_composed_prefill_14 {token_count = 14} - -check.benchmark<@qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_composed_benchmark_case> @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_composed_prefill_512 {token_count = 512} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_gate_up_swiglu_q4k.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_gate_up_swiglu_q4k.loom deleted file mode 100644 index 9164d5ef7d7c..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_gate_up_swiglu_q4k.loom +++ /dev/null @@ -1,1239 +0,0 @@ -// Fuses the Qwen3 MoE gate and up expert projections with the SwiGLU -// epilogue. Weights are consumed directly from GGUF's raw Q4_K layout: -// [expert][output channel][K / 256][144 bytes]. Activations use GGML's Q8_1 -// x4 layout produced by the shared quantization kernel. -// -// Route IDs have a logical route_count and an independent physical -// route_stride. This admits llama.cpp's [token][128] argsort storage while -// operating on only the selected top-8 entries, and also admits compact route -// buffers produced by a future fused router. -// -// Two schedules share this contract. Small token counts use one wave per -// [token, route, output channel]. Larger token counts first build a compact -// assignment table per expert, then group up to 32 selected rows with 32 -// output channels so raw weights and quantized activations can be reused -// through workgroup memory. -template.decl @ggml.matmul_id.vector_pair.body(%publish_output: i1, %binary_op: index, %input_size: index, %token_count: index, %token: index, %route_count: index, %route: index, %route_stride: index, %expert_count: index, %output_size: index, %channel: index, %lane: index, %q8_input: buffer, %route_ids: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) - -template.decl @qwen3_moe.routed_gate_up.quantize_q8_1_x4.subgroup_body(%publish_output: i1, %group_count0: index, %group0: index, %input: buffer, %output: buffer) - -template.decl @ggml.matmul_id.vector_pair.q8_1_x4.completion(%token_count: index, %route_count: index, %output_size: index, %token: index, %route: index, %channel_base: index, %subgroup: index, %workitem: index, %output: buffer, %completion_counters: buffer, %next_q8_output: buffer) - -config.decl @qwen3_moe.routed_gate_up.input_size : %value: index where [range(%value, 512, 32768), mul(%value, 512)] - -config.decl @qwen3_moe.routed_gate_up.expert_count : %value: index where [range(%value, 1, 512)] - -config.decl @qwen3_moe.routed_gate_up.route_count : %value: index where [range(%value, 1, 8)] - -config.decl @qwen3_moe.routed_gate_up.output_size : %value: index where [range(%value, 1, 4096)] - -config.decl @qwen3_moe.workload.token_capacity : %value: index where [range(%value, 1, 2048)] - -kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count$25: index, %input_size$26: index) launch(%token_count$27: index, %input_size$28: index, %input: buffer, %output: buffer) - -kernel.decl @ggml_moe_build_expert_table(%token_count$expert_table: index, %route_count$expert_table: index, %route_stride$expert_table: index, %expert_count$expert_table: index) launch(%token_count$expert_table_launch: index, %route_count$expert_table_launch: index, %route_stride$expert_table_launch: index, %expert_count$expert_table_launch: index, %route_ids$expert_table: buffer, %expert_table$expert_table: buffer) - -kernel.decl @ggml_moe_build_expert_partition_table(%token_count$partition_table: index, %route_count$partition_table: index, %expert_count$partition_table: index) launch(%token_count$partition_table_launch: index, %route_count$partition_table_launch: index, %expert_count$partition_table_launch: index, %expert_table$partition_table: buffer, %partition_table$partition_table: buffer) - -kernel.decl @ggml_mul_mat_id_vector_pair_binary_publish_f32(%token_count$29: index, %route_count$30: index, %route_stride$31: index, %expert_count$32: index, %output_size$33: index) launch(%token_count$34: index, %route_count$35: index, %route_stride$36: index, %expert_count$37: index, %output_size$38: index, %q8_input: buffer, %route_ids: buffer, %lhs_weight: buffer, %rhs_weight: buffer, %output: buffer) - -kernel.decl @ggml_mul_mat_id_vector_pair_binary_publish_q8(%token_count$39: index, %route_count$40: index, %route_stride$41: index, %expert_count$42: index, %output_size$43: index) launch(%token_count$44: index, %route_count$45: index, %route_stride$46: index, %expert_count$47: index, %output_size$48: index, %q8_input: buffer, %route_ids: buffer, %lhs_weight: buffer, %rhs_weight: buffer, %output: buffer, %completion_counters: buffer, %next_q8_output: buffer) - -func.decl @ggml_q8_1_x4_block(%q8_input: buffer, %row_byte_base: offset, %q8_block: index) -> (vector<32xi8>, f32, f32) - -// Decodes one 16-element Q4_K chunk while preserving its unsigned nibbles for -// AMDGPU's mixed u8*s8 dot4 instruction. The returned scale terms are shared -// by every routed activation row consuming the same weight chunk. -func.def inline @qwen3_moe_q4k_chunk_local(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group: index, %q4_half: index) -> (vector<16xi8>, f32, f32) { - %c4 = index.constant 4 : index - %c8 = index.constant 8 : index - %c16 = index.constant 16 : index - %block_bytes = index.constant 144 : offset - %header_byte_count = index.constant 4 : index - %code_byte_add = index.constant 16 : offset - %c4_i32 = scalar.constant 4 : i32 - %c6_i32 = scalar.constant 6 : i32 - %c15_i32 = scalar.constant 15 : i32 - %c63_i32 = scalar.constant 63 : i32 - %nibble_mask = vector.constant 252645135 : vector<4xi32> - %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index - %bounded_half = index.assume %q4_half [range(%q4_half, 0, 1)] : index - %block_byte_add = index.scale %q4_block, %block_bytes : index, offset -> offset - %block_byte_base = index.add %row_byte_base, %block_byte_add : offset - %code_byte_base = index.add %block_byte_base, %code_byte_add : offset - %header_view = buffer.view %weight[%block_byte_base] : buffer -> view<4xi32> - %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<32xi32> - %header_words = vector.load %header_view[0] : view<4xi32> -> vector<4xi32> - %header_halves = vector.bitcast %header_words : vector<4xi32> to vector<8xf16> - %d_f16 = vector.extract %header_halves[0] : vector<8xf16> -> f16 - %dmin_f16 = vector.extract %header_halves[1] : vector<8xf16> -> f16 - %d = scalar.extf %d_f16 : f16 to f32 - %dmin = scalar.extf %dmin_f16 : f16 to f32 - %header_bytes = vector.bitcast %header_words : vector<4xi32> to vector<16xi8> - %iqs_group = index.mul %bounded_group, %c8 : index - %iqs_half = index.mul %bounded_half, %c4 : index - %iqs = index.add %iqs_group, %iqs_half : index - %qs_page0 = index.div %iqs, %c16 : index - %qs_page = index.mul %qs_page0, %c8 : index - %qs_lane = index.rem %iqs, %c8 : index - %qs_index0 = index.add %qs_page, %qs_lane : index - %qs_index = index.assume %qs_index0 [range(%qs_index0, 0, 28)] : index - %iqs_mod16 = index.rem %iqs, %c16 : index - %nibble_page = index.div %iqs_mod16, %c8 : index - %nibble_shift_index = index.mul %nibble_page, %c4 : index - %nibble_shift_i32 = index.cast %nibble_shift_index : index to i32 - %nibble_shift = vector.splat %nibble_shift_i32 : vector<4xi32> - %packed_codes = vector.load %code_view[%qs_index] : view<32xi32> -> vector<4xi32> - %shifted_codes = vector.shrui %packed_codes, %nibble_shift : vector<4xi32> - %masked_codes = vector.andi %shifted_codes, %nibble_mask : vector<4xi32> - %q4_values = vector.bitcast %masked_codes : vector<4xi32> to vector<16xi8> - %is_low_group = index.cmp ult, %bounded_group, %c4 : index - %scale, %minimum = scf.if %is_low_group -> (i32, i32) { - %scale_index = index.add %bounded_group, %header_byte_count : index - %minimum_index0 = index.add %bounded_group, %c4 : index - %minimum_index = index.add %minimum_index0, %header_byte_count : index - %scale_i8 = vector.extract %header_bytes[%scale_index] : vector<16xi8> -> i8 - %minimum_i8 = vector.extract %header_bytes[%minimum_index] : vector<16xi8> -> i8 - %scale_u8 = scalar.extui %scale_i8 : i8 to i32 - %minimum_u8 = scalar.extui %minimum_i8 : i8 to i32 - %scale_low6 = scalar.andi %scale_u8, %c63_i32 : i32 - %minimum_low6 = scalar.andi %minimum_u8, %c63_i32 : i32 - scf.yield %scale_low6, %minimum_low6 : i32, i32 - } else { - %packed_index0 = index.add %bounded_group, %c4 : index - %packed_index = index.add %packed_index0, %header_byte_count : index - %scale_high_index = index.sub %bounded_group, %c4 : index - %scale_high_header_index = index.add %scale_high_index, %header_byte_count : index - %minimum_high_header_index = index.add %bounded_group, %header_byte_count : index - %packed_i8 = vector.extract %header_bytes[%packed_index] : vector<16xi8> -> i8 - %scale_high_i8 = vector.extract %header_bytes[%scale_high_header_index] : vector<16xi8> -> i8 - %minimum_high_i8 = vector.extract %header_bytes[%minimum_high_header_index] : vector<16xi8> -> i8 - %packed = scalar.extui %packed_i8 : i8 to i32 - %scale_high = scalar.extui %scale_high_i8 : i8 to i32 - %minimum_high = scalar.extui %minimum_high_i8 : i8 to i32 - %scale_low4 = scalar.andi %packed, %c15_i32 : i32 - %minimum_low4 = scalar.shrui %packed, %c4_i32 : i32 - %scale_high2_raw = scalar.shrui %scale_high, %c6_i32 : i32 - %minimum_high2_raw = scalar.shrui %minimum_high, %c6_i32 : i32 - %scale_high2 = scalar.shli %scale_high2_raw, %c4_i32 : i32 - %minimum_high2 = scalar.shli %minimum_high2_raw, %c4_i32 : i32 - %scale_value = scalar.ori %scale_low4, %scale_high2 : i32 - %minimum_value = scalar.ori %minimum_low4, %minimum_high2 : i32 - scf.yield %scale_value, %minimum_value : i32, i32 - } - %scale_f32 = scalar.uitofp %scale : i32 to f32 - %minimum_f32 = scalar.uitofp %minimum : i32 to f32 - %d_scale = scalar.mulf %d, %scale_f32 : f32 - %dmin_scale = scalar.mulf %dmin, %minimum_f32 : f32 - func.return %q4_values, %d_scale, %dmin_scale : vector<16xi8>, f32, f32 -} - -// Global-weight variant of the Q4_K chunk decoder. Scalar scale-byte loads are -// legal and cheaper for the one-wave schedule, while the LDS schedule uses the -// packed-header variant above because AMDGPU has no local i8 load descriptor. -func.def inline @qwen3_moe_q4k_chunk_global(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group: index, %q4_half: index) -> (vector<16xi8>, f32, f32) { - %c4 = index.constant 4 : index - %c8 = index.constant 8 : index - %c16 = index.constant 16 : index - %block_bytes = index.constant 144 : offset - %scale_byte_add = index.constant 4 : offset - %code_byte_add = index.constant 16 : offset - %c4_i32 = scalar.constant 4 : i32 - %c6_i32 = scalar.constant 6 : i32 - %c15_i32 = scalar.constant 15 : i32 - %c63_i32 = scalar.constant 63 : i32 - %nibble_mask = vector.constant 252645135 : vector<4xi32> - %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index - %bounded_half = index.assume %q4_half [range(%q4_half, 0, 1)] : index - %block_byte_add = index.scale %q4_block, %block_bytes : index, offset -> offset - %block_byte_base = index.add %row_byte_base, %block_byte_add : offset - %scale_byte_base = index.add %block_byte_base, %scale_byte_add : offset - %code_byte_base = index.add %block_byte_base, %code_byte_add : offset - %half_view = buffer.view %weight[%block_byte_base] : buffer -> view<2xf16> - %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<12xi8> - %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<32xi32> - %d_f16 = view.load %half_view[0] : view<2xf16> -> f16 - %dmin_f16 = view.load %half_view[1] : view<2xf16> -> f16 - %d = scalar.extf %d_f16 : f16 to f32 - %dmin = scalar.extf %dmin_f16 : f16 to f32 - %iqs_group = index.mul %bounded_group, %c8 : index - %iqs_half = index.mul %bounded_half, %c4 : index - %iqs = index.add %iqs_group, %iqs_half : index - %qs_page0 = index.div %iqs, %c16 : index - %qs_page = index.mul %qs_page0, %c8 : index - %qs_lane = index.rem %iqs, %c8 : index - %qs_index0 = index.add %qs_page, %qs_lane : index - %qs_index = index.assume %qs_index0 [range(%qs_index0, 0, 28)] : index - %iqs_mod16 = index.rem %iqs, %c16 : index - %nibble_page = index.div %iqs_mod16, %c8 : index - %nibble_shift_index = index.mul %nibble_page, %c4 : index - %nibble_shift_i32 = index.cast %nibble_shift_index : index to i32 - %nibble_shift = vector.splat %nibble_shift_i32 : vector<4xi32> - %packed_codes = vector.load %code_view[%qs_index] : view<32xi32> -> vector<4xi32> - %shifted_codes = vector.shrui %packed_codes, %nibble_shift : vector<4xi32> - %masked_codes = vector.andi %shifted_codes, %nibble_mask : vector<4xi32> - %q4_values = vector.bitcast %masked_codes : vector<4xi32> to vector<16xi8> - %is_low_group = index.cmp ult, %bounded_group, %c4 : index - %scale, %minimum = scf.if %is_low_group -> (i32, i32) { - %minimum_index = index.add %bounded_group, %c4 : index - %scale_i8 = view.load %scale_view[%bounded_group] : view<12xi8> -> i8 - %minimum_i8 = view.load %scale_view[%minimum_index] : view<12xi8> -> i8 - %scale_u8 = scalar.extui %scale_i8 : i8 to i32 - %minimum_u8 = scalar.extui %minimum_i8 : i8 to i32 - %scale_low6 = scalar.andi %scale_u8, %c63_i32 : i32 - %minimum_low6 = scalar.andi %minimum_u8, %c63_i32 : i32 - scf.yield %scale_low6, %minimum_low6 : i32, i32 - } else { - %packed_index = index.add %bounded_group, %c4 : index - %scale_high_index = index.sub %bounded_group, %c4 : index - %packed_i8 = view.load %scale_view[%packed_index] : view<12xi8> -> i8 - %scale_high_i8 = view.load %scale_view[%scale_high_index] : view<12xi8> -> i8 - %minimum_high_i8 = view.load %scale_view[%bounded_group] : view<12xi8> -> i8 - %packed = scalar.extui %packed_i8 : i8 to i32 - %scale_high = scalar.extui %scale_high_i8 : i8 to i32 - %minimum_high = scalar.extui %minimum_high_i8 : i8 to i32 - %scale_low4 = scalar.andi %packed, %c15_i32 : i32 - %minimum_low4 = scalar.shrui %packed, %c4_i32 : i32 - %scale_high2_raw = scalar.shrui %scale_high, %c6_i32 : i32 - %minimum_high2_raw = scalar.shrui %minimum_high, %c6_i32 : i32 - %scale_high2 = scalar.shli %scale_high2_raw, %c4_i32 : i32 - %minimum_high2 = scalar.shli %minimum_high2_raw, %c4_i32 : i32 - %scale_value = scalar.ori %scale_low4, %scale_high2 : i32 - %minimum_value = scalar.ori %minimum_low4, %minimum_high2 : i32 - scf.yield %scale_value, %minimum_value : i32, i32 - } - %scale_f32 = scalar.uitofp %scale : i32 to f32 - %minimum_f32 = scalar.uitofp %minimum : i32 to f32 - %d_scale = scalar.mulf %d, %scale_f32 : f32 - %dmin_scale = scalar.mulf %dmin, %minimum_f32 : f32 - func.return %q4_values, %d_scale, %dmin_scale : vector<16xi8>, f32, f32 -} - -// Decodes one Q4_K scale/minimum pair from a header loaded as three packed -// scale words. Keeping the complete header in registers lets paired-group -// contractions share both the header and packed-code load. -func.def inline @qwen3_moe_q4k_scale_from_header(%scale0: i32, %scale1: i32, %scale2: i32, %q4_group: index) -> (i32, i32) { - %c4 = index.constant 4 : index - %c8 = index.constant 8 : index - %c2_i32 = scalar.constant 2 : i32 - %c4_i32 = scalar.constant 4 : i32 - %c15_i32 = scalar.constant 15 : i32 - %c48_i32 = scalar.constant 48 : i32 - %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index - %is_low_group = index.cmp ult, %bounded_group, %c4 : index - %scale_lane = index.rem %bounded_group, %c4 : index - %scale_shift_index = index.mul %scale_lane, %c8 : index - %scale_shift = index.cast %scale_shift_index : index to i32 - %high_shift = scalar.addi %scale_shift, %c2_i32 : i32 - %minimum_shift = scalar.addi %scale_shift, %c4_i32 : i32 - %selected_scale_source = scf.select %is_low_group, %scale0, %scale2 : i32 - %selected_minimum_source = scf.select %is_low_group, %scale1, %scale2 : i32 - %selected_scale_high_shift = scf.select %is_low_group, %scale_shift, %high_shift : i32 - %selected_minimum_low_shift = scf.select %is_low_group, %scale_shift, %minimum_shift : i32 - %scale_low0 = scalar.shrui %selected_scale_source, %scale_shift : i32 - %scale_low = scalar.andi %scale_low0, %c15_i32 : i32 - %scale_high0 = scalar.shrui %scale0, %selected_scale_high_shift : i32 - %scale_high = scalar.andi %scale_high0, %c48_i32 : i32 - %scale = scalar.ori %scale_low, %scale_high : i32 - %minimum_low0 = scalar.shrui %selected_minimum_source, %selected_minimum_low_shift : i32 - %minimum_low = scalar.andi %minimum_low0, %c15_i32 : i32 - %minimum_high0 = scalar.shrui %scale1, %selected_scale_high_shift : i32 - %minimum_high = scalar.andi %minimum_high0, %c48_i32 : i32 - %minimum = scalar.ori %minimum_low, %minimum_high : i32 - func.return %scale, %minimum : i32, i32 -} - -// Decodes two adjacent Q4_K groups from the low and high nibbles of one -// packed-code load. The caller supplies the already-loaded 16-byte header. -func.def inline @qwen3_moe_q4k_chunk_pair_global(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group_pair: index, %q4_half: index, %header_words: vector<4xi32>) -> (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) { - %c1 = index.constant 1 : index - %c2 = index.constant 2 : index - %c4 = index.constant 4 : index - %c8 = index.constant 8 : index - %block_bytes = index.constant 144 : offset - %code_byte_add = index.constant 16 : offset - %c4_i32 = scalar.constant 4 : i32 - %nibble_mask = vector.constant 252645135 : vector<4xi32> - %bounded_pair = index.assume %q4_group_pair [range(%q4_group_pair, 0, 3)] : index - %bounded_half = index.assume %q4_half [range(%q4_half, 0, 1)] : index - %block_byte_add = index.scale %q4_block, %block_bytes : index, offset -> offset - %block_byte_base = index.add %row_byte_base, %block_byte_add : offset - %code_byte_base = index.add %block_byte_base, %code_byte_add : offset - %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<32xi32> - %header_halves = vector.bitcast %header_words : vector<4xi32> to vector<8xf16> - %d_f16 = vector.extract %header_halves[0] : vector<8xf16> -> f16 - %dmin_f16 = vector.extract %header_halves[1] : vector<8xf16> -> f16 - %scale0 = vector.extract %header_words[1] : vector<4xi32> -> i32 - %scale1 = vector.extract %header_words[2] : vector<4xi32> -> i32 - %scale2 = vector.extract %header_words[3] : vector<4xi32> -> i32 - %d = scalar.extf %d_f16 : f16 to f32 - %dmin = scalar.extf %dmin_f16 : f16 to f32 - %pair_code_base = index.mul %bounded_pair, %c8 : index - %half_code_add = index.mul %bounded_half, %c4 : index - %code_index0 = index.add %pair_code_base, %half_code_add : index - %code_index = index.assume %code_index0 [range(%code_index0, 0, 28)] : index - %packed_codes = vector.load %code_view[%code_index] : view<32xi32> -> vector<4xi32> - %low_codes = vector.andi %packed_codes, %nibble_mask : vector<4xi32> - %c4_i32v = vector.splat %c4_i32 : vector<4xi32> - %high_shifted = vector.shrui %packed_codes, %c4_i32v : vector<4xi32> - %high_codes = vector.andi %high_shifted, %nibble_mask : vector<4xi32> - %q4_low = vector.bitcast %low_codes : vector<4xi32> to vector<16xi8> - %q4_high = vector.bitcast %high_codes : vector<4xi32> to vector<16xi8> - %low_group = index.mul %bounded_pair, %c2 : index - %high_group = index.add %low_group, %c1 : index - %low_scale, %low_minimum = func.call @qwen3_moe_q4k_scale_from_header(%scale0, %scale1, %scale2, %low_group) : (i32, i32, i32, index) -> (i32, i32) - %high_scale, %high_minimum = func.call @qwen3_moe_q4k_scale_from_header(%scale0, %scale1, %scale2, %high_group) : (i32, i32, i32, index) -> (i32, i32) - %low_scale_f32 = scalar.uitofp %low_scale : i32 to f32 - %low_minimum_f32 = scalar.uitofp %low_minimum : i32 to f32 - %high_scale_f32 = scalar.uitofp %high_scale : i32 to f32 - %high_minimum_f32 = scalar.uitofp %high_minimum : i32 to f32 - %low_d_scale = scalar.mulf %d, %low_scale_f32 : f32 - %low_dmin_scale = scalar.mulf %dmin, %low_minimum_f32 : f32 - %high_d_scale = scalar.mulf %d, %high_scale_f32 : f32 - %high_dmin_scale = scalar.mulf %dmin, %high_minimum_f32 : f32 - func.return %q4_low, %low_d_scale, %low_dmin_scale, %q4_high, %high_d_scale, %high_dmin_scale : vector<16xi8>, f32, f32, vector<16xi8>, f32, f32 -} - -// Contracts one decoded Q4_K half-group with a Q8_1 half-block. Both -// half-groups apply half of the Q8_1 block-sum correction. -func.def inline @qwen3_moe_q4k_q8_1_dot(%q4_values: vector<16xi8>, %d_scale: f32, %dmin_scale: f32, %q8_values: vector<16xi8>, %q8_d: f32, %q8_s: f32) -> (f32) { - %c0_i32 = scalar.constant 0 : i32 - %c0_i32v = vector.constant 0 : vector<4xi32> - %half_f32 = scalar.constant 0.5 : f32 - %partial_dots = vector.dot4i %q4_values, %q8_values, %c0_i32v : vector<16xi8>, vector<16xi8>, vector<4xi32> - %q_sum = vector.reduce %partial_dots, %c0_i32 : vector<4xi32>, i32 - %q_sum_f32 = scalar.sitofp %q_sum : i32 to f32 - %scaled_dot0 = scalar.mulf %q8_d, %d_scale : f32 - %scaled_dot = scalar.mulf %scaled_dot0, %q_sum_f32 : f32 - %q8_half_sum = scalar.mulf %q8_s, %half_f32 : f32 - %minimum_correction = scalar.mulf %dmin_scale, %q8_half_sum : f32 - %contribution = scalar.subf %scaled_dot, %minimum_correction : f32 - func.return %contribution : f32 -} - -// Contracts one decoded Q4_K half-group against four routed Q8_1 rows. The -// leading row axis stays explicit through dot4 and scale correction so the -// grouped schedule can reuse each weight decode without scalarizing the rows. -func.def inline @qwen3_moe_q4k_q8_1_dot4_rows(%q4_values: vector<16xi8>, %d_scale: f32, %dmin_scale: f32, %q8_values: vector<4x16xi8>, %q8_d: vector<4xf32>, %q8_s: vector<4xf32>) -> (vector<4xf32>) { - %q8_values0 = vector.extract %q8_values[0] : vector<4x16xi8> -> vector<16xi8> - %q8_values1 = vector.extract %q8_values[1] : vector<4x16xi8> -> vector<16xi8> - %q8_values2 = vector.extract %q8_values[2] : vector<4x16xi8> -> vector<16xi8> - %q8_values3 = vector.extract %q8_values[3] : vector<4x16xi8> -> vector<16xi8> - %q8_d0 = vector.extract %q8_d[0] : vector<4xf32> -> f32 - %q8_d1 = vector.extract %q8_d[1] : vector<4xf32> -> f32 - %q8_d2 = vector.extract %q8_d[2] : vector<4xf32> -> f32 - %q8_d3 = vector.extract %q8_d[3] : vector<4xf32> -> f32 - %q8_s0 = vector.extract %q8_s[0] : vector<4xf32> -> f32 - %q8_s1 = vector.extract %q8_s[1] : vector<4xf32> -> f32 - %q8_s2 = vector.extract %q8_s[2] : vector<4xf32> -> f32 - %q8_s3 = vector.extract %q8_s[3] : vector<4xf32> -> f32 - %contribution0 = func.call @qwen3_moe_q4k_q8_1_dot(%q4_values, %d_scale, %dmin_scale, %q8_values0, %q8_d0, %q8_s0) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) - %contribution1 = func.call @qwen3_moe_q4k_q8_1_dot(%q4_values, %d_scale, %dmin_scale, %q8_values1, %q8_d1, %q8_s1) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) - %contribution2 = func.call @qwen3_moe_q4k_q8_1_dot(%q4_values, %d_scale, %dmin_scale, %q8_values2, %q8_d2, %q8_s2) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) - %contribution3 = func.call @qwen3_moe_q4k_q8_1_dot(%q4_values, %d_scale, %dmin_scale, %q8_values3, %q8_d3, %q8_s3) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) - %contribution = vector.from_elements %contribution0, %contribution1, %contribution2, %contribution3 : vector<4xf32> - func.return %contribution : vector<4xf32> -} - -// Convenience wrapper used by the one-wave schedule. -func.def inline @qwen3_moe_q4k_q8_1_chunk(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group: index, %q4_half: index, %q8_values: vector<16xi8>, %q8_d: f32, %q8_s: f32) -> (f32) { - %q4_values, %d_scale, %dmin_scale = func.call @qwen3_moe_q4k_chunk_global(%weight, %row_byte_base, %q4_block, %q4_group, %q4_half) : (buffer, offset, index, index, index) -> (vector<16xi8>, f32, f32) - %contribution = func.call @qwen3_moe_q4k_q8_1_dot(%q4_values, %d_scale, %dmin_scale, %q8_values, %q8_d, %q8_s) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) - func.return %contribution : f32 -} - -// Computes one lane's partial for a complete Q4_K row. Dense and routed -// providers choose output ownership independently, then share this exact -// packed-row contraction before reducing across a wave32 subgroup. -func.def inline @qwen3_moe_q4k_q8_1_x4_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32) { - %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index - %bounded_lane = index.assume %lane [range(%lane, 0, 31)] : index - %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 - %c128 = index.constant 128 : index - %c144 = index.constant 144 : index - %c256 = index.constant 256 : index - %c511 = index.constant 511 : index - %c512 = index.constant 512 : index - %q8_group_bytes = index.constant 144 : offset - %q8_payload_byte_add = index.constant 16 : offset - %c0_f32 = scalar.constant 0.0 : f32 - %q4_block_count = index.div %bounded_input_size, %c256 : index - %q8_group_count = index.div %bounded_input_size, %c128 : index - %padded_input_size = index.add %bounded_input_size, %c511 : index - %iteration_count = index.div %padded_input_size, %c512 : index - %q4_group0 = index.div %bounded_lane, %c2 : index - %q4_group1 = index.rem %q4_group0, %c8 : index - %q4_group = index.assume %q4_group1 [range(%q4_group1, 0, 7)] : index - %q4_half0 = index.rem %bounded_lane, %c2 : index - %q4_half = index.assume %q4_half0 [range(%q4_half0, 0, 1)] : index - %lane_q4_block = index.div %bounded_lane, %c16 : index - %lane_q8_group = index.div %bounded_lane, %c8 : index - %q8_inner_block0 = index.div %bounded_lane, %c2 : index - %q8_inner_block1 = index.rem %q8_inner_block0, %c4 : index - %q8_inner_block = index.assume %q8_inner_block1 [range(%q8_inner_block1, 0, 3)] : index - %q8_inner_word_base = index.mul %q8_inner_block, %c8 : index - %q8_half_word_add = index.mul %q4_half, %c4 : index - %q8_word_index0 = index.add %q8_inner_word_base, %q8_half_word_add : index - %q8_word_index = index.assume %q8_word_index0 [range(%q8_word_index0, 0, 28)] : index - %q8_ds_index = index.mul %q8_inner_block, %c2 : index - %q8_s_index = index.add %q8_ds_index, %c1 : index - %sum = scf.for %iteration = [%c0 to %iteration_count step %c1](%iteration_acc = %c0_f32 : f32) -> (f32) unroll { - %iteration_q4_block = index.mul %iteration, %c2 : index - %q4_block0 = index.add %iteration_q4_block, %lane_q4_block : index - %valid_q4_block = index.cmp ult, %q4_block0, %q4_block_count : index - %contribution = scf.if %valid_q4_block -> (f32) { - %q4_block, %bounded_q4_block_count = index.assume %q4_block0, %q4_block_count [lt(%q4_block0, %q4_block_count)] : index, index - %iteration_q8_group = index.mul %iteration, %c4 : index - %q8_group0 = index.add %iteration_q8_group, %lane_q8_group : index - %q8_group, %bounded_q8_group_count = index.assume %q8_group0, %q8_group_count [lt(%q8_group0, %q8_group_count)] : index, index - %q8_group_byte_add = index.scale %q8_group, %q8_group_bytes : index, offset -> offset - %q8_group_byte_base = index.add %q8_row_byte_base, %q8_group_byte_add : offset - %q8_payload_byte_base = index.add %q8_group_byte_base, %q8_payload_byte_add : offset - %q8_ds_view = buffer.view %q8_input[%q8_group_byte_base] : buffer -> view<8xf16> - %q8_words_view = buffer.view %q8_input[%q8_payload_byte_base] : buffer -> view<32xi32> - %q8_d_f16 = view.load %q8_ds_view[%q8_ds_index] : view<8xf16> -> f16 - %q8_s_f16 = view.load %q8_ds_view[%q8_s_index] : view<8xf16> -> f16 - %q8_d = scalar.extf %q8_d_f16 : f16 to f32 - %q8_s = scalar.extf %q8_s_f16 : f16 to f32 - %q8_words = vector.load %q8_words_view[%q8_word_index] : view<32xi32> -> vector<4xi32> - %q8_values = vector.bitcast %q8_words : vector<4xi32> to vector<16xi8> - %dot = func.call @qwen3_moe_q4k_q8_1_chunk(%weight, %weight_row_byte_base, %q4_block, %q4_group, %q4_half, %q8_values, %q8_d, %q8_s) : (buffer, offset, index, index, index, vector<16xi8>, f32, f32) -> (f32) - scf.yield %dot : f32 - } else { - scf.yield %c0_f32 : f32 - } - %next = scalar.addf %iteration_acc, %contribution : f32 - scf.yield %next : f32 - } - func.return %sum : f32 -} - -// Computes one lane's two adjacent-group contributions within one Q4_K block. -// Eight block lanes consume both nibbles from every packed code load and cover -// the block's full 256-element input extent. -func.def inline @qwen3_moe_q4k_q8_1_x4_paired_block_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %q4_block: index, %block_lane: index) -> (f32) { - %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index - %bounded_q4_block0 = index.assume %q4_block [range(%q4_block, 0, 127)] : index - %bounded_block_lane = index.assume %block_lane [range(%block_lane, 0, 7)] : index - %c1 = index.constant 1 : index - %c2 = index.constant 2 : index - %c4 = index.constant 4 : index - %c8 = index.constant 8 : index - %c128 = index.constant 128 : index - %c256 = index.constant 256 : index - %q4_block_bytes = index.constant 144 : offset - %q8_group_bytes = index.constant 144 : offset - %q8_payload_byte_add = index.constant 16 : offset - %q4_block_count = index.div %bounded_input_size, %c256 : index - %q8_group_count = index.div %bounded_input_size, %c128 : index - %bounded_q4_block, %bounded_q4_block_count = index.assume %bounded_q4_block0, %q4_block_count [lt(%bounded_q4_block0, %q4_block_count)] : index, index - %q4_group_pair0 = index.div %bounded_block_lane, %c2 : index - %q4_group_pair = index.assume %q4_group_pair0 [range(%q4_group_pair0, 0, 3)] : index - %q4_half0 = index.rem %bounded_block_lane, %c2 : index - %q4_half = index.assume %q4_half0 [range(%q4_half0, 0, 1)] : index - %q8_group_in_block0 = index.div %q4_group_pair, %c2 : index - %q8_group_in_block = index.assume %q8_group_in_block0 [range(%q8_group_in_block0, 0, 1)] : index - %pair_in_q8_group0 = index.rem %q4_group_pair, %c2 : index - %pair_in_q8_group = index.assume %pair_in_q8_group0 [range(%pair_in_q8_group0, 0, 1)] : index - %q8_low_inner_block0 = index.mul %pair_in_q8_group, %c2 : index - %q8_low_inner_block = index.assume %q8_low_inner_block0 [range(%q8_low_inner_block0, 0, 2)] : index - %q8_high_inner_block0 = index.add %q8_low_inner_block, %c1 : index - %q8_high_inner_block = index.assume %q8_high_inner_block0 [range(%q8_high_inner_block0, 1, 3)] : index - %q8_half_word_add = index.mul %q4_half, %c4 : index - %q8_low_inner_word_base = index.mul %q8_low_inner_block, %c8 : index - %q8_low_word_index0 = index.add %q8_low_inner_word_base, %q8_half_word_add : index - %q8_low_word_index = index.assume %q8_low_word_index0 [range(%q8_low_word_index0, 0, 20)] : index - %q8_high_inner_word_base = index.mul %q8_high_inner_block, %c8 : index - %q8_high_word_index0 = index.add %q8_high_inner_word_base, %q8_half_word_add : index - %q8_high_word_index = index.assume %q8_high_word_index0 [range(%q8_high_word_index0, 8, 28)] : index - %q8_low_ds_index0 = index.mul %q8_low_inner_block, %c2 : index - %q8_low_ds_index = index.assume %q8_low_ds_index0 [range(%q8_low_ds_index0, 0, 4)] : index - %q8_block_group_base = index.mul %bounded_q4_block, %c2 : index - %q8_group0 = index.add %q8_block_group_base, %q8_group_in_block : index - %q8_group, %bounded_q8_group_count = index.assume %q8_group0, %q8_group_count [lt(%q8_group0, %q8_group_count)] : index, index - %q8_group_byte_add = index.scale %q8_group, %q8_group_bytes : index, offset -> offset - %q8_group_byte_base = index.add %q8_row_byte_base, %q8_group_byte_add : offset - %q8_payload_byte_base = index.add %q8_group_byte_base, %q8_payload_byte_add : offset - %q8_ds_view = buffer.view %q8_input[%q8_group_byte_base] : buffer -> view<8xf16> - %q8_words_view = buffer.view %q8_input[%q8_payload_byte_base] : buffer -> view<32xi32> - %q8_ds = vector.load %q8_ds_view[%q8_low_ds_index] : view<8xf16> -> vector<4xf16> - %q8_low_d_f16 = vector.extract %q8_ds[0] : vector<4xf16> -> f16 - %q8_low_s_f16 = vector.extract %q8_ds[1] : vector<4xf16> -> f16 - %q8_high_d_f16 = vector.extract %q8_ds[2] : vector<4xf16> -> f16 - %q8_high_s_f16 = vector.extract %q8_ds[3] : vector<4xf16> -> f16 - %q8_low_d = scalar.extf %q8_low_d_f16 : f16 to f32 - %q8_low_s = scalar.extf %q8_low_s_f16 : f16 to f32 - %q8_high_d = scalar.extf %q8_high_d_f16 : f16 to f32 - %q8_high_s = scalar.extf %q8_high_s_f16 : f16 to f32 - %q8_low_words = vector.load %q8_words_view[%q8_low_word_index] : view<32xi32> -> vector<4xi32> - %q8_high_words = vector.load %q8_words_view[%q8_high_word_index] : view<32xi32> -> vector<4xi32> - %q8_low_values = vector.bitcast %q8_low_words : vector<4xi32> to vector<16xi8> - %q8_high_values = vector.bitcast %q8_high_words : vector<4xi32> to vector<16xi8> - %q4_block_byte_add = index.scale %bounded_q4_block, %q4_block_bytes : index, offset -> offset - %q4_block_byte_base = index.add %weight_row_byte_base, %q4_block_byte_add : offset - %q4_header_view = buffer.view %weight[%q4_block_byte_base] : buffer -> view<4xi32> - %q4_header_words = vector.load %q4_header_view[0] : view<4xi32> -> vector<4xi32> - %q4_low, %low_d_scale, %low_dmin_scale, %q4_high, %high_d_scale, %high_dmin_scale = func.call @qwen3_moe_q4k_chunk_pair_global(%weight, %weight_row_byte_base, %bounded_q4_block, %q4_group_pair, %q4_half, %q4_header_words) : (buffer, offset, index, index, index, vector<4xi32>) -> (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) - %low = func.call @qwen3_moe_q4k_q8_1_dot(%q4_low, %low_d_scale, %low_dmin_scale, %q8_low_values, %q8_low_d, %q8_low_s) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) - %high = func.call @qwen3_moe_q4k_q8_1_dot(%q4_high, %high_d_scale, %high_dmin_scale, %q8_high_values, %q8_high_d, %q8_high_s) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) - %pair = scalar.addf %low, %high : f32 - func.return %pair : f32 -} - -// Computes one lane's partial while consuming both nibbles of each packed -// Q4_K code load. Eight lanes cover one 256-element Q4_K block, so a wave32 -// advances through four blocks per iteration and reduces the adjacent-group -// contributions together. -func.def inline @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32) { - %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index - %bounded_lane = index.assume %lane [range(%lane, 0, 31)] : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c4 = index.constant 4 : index - %c8 = index.constant 8 : index - %c256 = index.constant 256 : index - %c1023 = index.constant 1023 : index - %c1024 = index.constant 1024 : index - %c0_f32 = scalar.constant 0.0 : f32 - %q4_block_count = index.div %bounded_input_size, %c256 : index - %padded_input_size = index.add %bounded_input_size, %c1023 : index - %iteration_count = index.div %padded_input_size, %c1024 : index - %lane_q4_block = index.div %bounded_lane, %c8 : index - %block_lane0 = index.rem %bounded_lane, %c8 : index - %block_lane = index.assume %block_lane0 [range(%block_lane0, 0, 7)] : index - %sum = scf.for %iteration = [%c0 to %iteration_count step %c1](%iteration_acc = %c0_f32 : f32) -> (f32) unroll { - %iteration_q4_block = index.mul %iteration, %c4 : index - %q4_block0 = index.add %iteration_q4_block, %lane_q4_block : index - %valid_q4_block = index.cmp ult, %q4_block0, %q4_block_count : index - %contribution = scf.if %valid_q4_block -> (f32) { - %q4_block, %bounded_q4_block_count = index.assume %q4_block0, %q4_block_count [lt(%q4_block0, %q4_block_count)] : index, index - %pair = func.call @qwen3_moe_q4k_q8_1_x4_paired_block_lane(%bounded_input_size, %weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %q4_block, %block_lane) : (index, buffer, offset, buffer, offset, index, index) -> (f32) - scf.yield %pair : f32 - } else { - scf.yield %c0_f32 : f32 - } - %next = scalar.addf %iteration_acc, %contribution : f32 - scf.yield %next : f32 - } - func.return %sum : f32 -} - -// Computes one lane's partial for a row owned by an eight-lane cohort. Every -// cohort lane consumes both nibbles for one adjacent Q4_K group pair while the -// cohort walks all blocks in the row. This schedule lets separate cohorts -// contract independent routed rows concurrently. -func.def inline @qwen3_moe_q4k_q8_1_x4_cohort_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %cohort_lane: index) -> (f32) { - %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index - %bounded_cohort_lane = index.assume %cohort_lane [range(%cohort_lane, 0, 7)] : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c256 = index.constant 256 : index - %c0_f32 = scalar.constant 0.0 : f32 - %q4_block_count = index.div %bounded_input_size, %c256 : index - %sum = scf.for %q4_block0 = [%c0 to %q4_block_count step %c1](%block_acc = %c0_f32 : f32) -> (f32) unroll { - %q4_block, %bounded_q4_block_count = index.assume %q4_block0, %q4_block_count [lt(%q4_block0, %q4_block_count)] : index, index - %pair = func.call @qwen3_moe_q4k_q8_1_x4_paired_block_lane(%bounded_input_size, %weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %q4_block, %bounded_cohort_lane) : (index, buffer, offset, buffer, offset, index, index) -> (f32) - %next = scalar.addf %block_acc, %pair : f32 - scf.yield %next : f32 - } - func.return %sum : f32 -} - -// Decodes the packed queue descriptor shared by expert-grouped projections. -// The producer stores a 7-bit expert ordinal, a 6-bit 32-row partition -// ordinal, and a 5-bit row count minus one. -func.def inline @qwen3_moe_unpack_expert_partition_descriptor(%descriptor: i32) -> (index, index, index) { - %c1_i32 = scalar.constant 1 : i32 - %c5_i32 = scalar.constant 5 : i32 - %c7_i32 = scalar.constant 7 : i32 - %c13_i32 = scalar.constant 13 : i32 - %c31_i32 = scalar.constant 31 : i32 - %c63_i32 = scalar.constant 63 : i32 - %c127_i32 = scalar.constant 127 : i32 - %expert_i32 = scalar.andi %descriptor, %c127_i32 : i32 - %partition_shifted_i32 = scalar.shrui %descriptor, %c7_i32 : i32 - %partition_i32 = scalar.andi %partition_shifted_i32, %c63_i32 : i32 - %route_tile_base_i32 = scalar.shli %partition_i32, %c5_i32 : i32 - %row_count_shifted_i32 = scalar.shrui %descriptor, %c13_i32 : i32 - %row_count_minus_one_i32 = scalar.andi %row_count_shifted_i32, %c31_i32 : i32 - %partition_row_count_i32 = scalar.addi %row_count_minus_one_i32, %c1_i32 : i32 - %expert0 = index.cast %expert_i32 : i32 to index - %expert = index.assume %expert0 [range(%expert0, 0, 127)] : index - %route_tile_base0 = index.cast %route_tile_base_i32 : i32 to index - %route_tile_base = index.assume %route_tile_base0 [range(%route_tile_base0, 0, 2016)] : index - %partition_row_count0 = index.cast %partition_row_count_i32 : i32 to index - %partition_row_count = index.assume %partition_row_count0 [range(%partition_row_count0, 1, 32)] : index - func.return %expert, %route_tile_base, %partition_row_count : index, index, index -} - -// Packs one Q8_1 x4 physical group with one selected subgroup. Each lane owns -// four adjacent values and contributes its maximum and sum to one of four -// eight-lane logical blocks. One-hot vectors let ordinary subgroup reductions -// compute all four block aggregates without LDS or workgroup barriers. -template.def<@qwen3_moe.routed_gate_up.quantize_q8_1_x4.subgroup_body> device @qwen3_moe_routed_gate_up_quantize_q8_1_x4_subgroup_body(%publish_output: i1, %group_count0: index, %group0: index, %input: buffer, %output: buffer) { - %group_count, %group = index.assume %group_count0, %group0 [range(%group_count0, 1, 524288), lt(%group0, %group_count0)] : index, index - %lane0 = kernel.subgroup.lane.id : index - %lane = index.assume %lane0 [range(%lane0, 0, 31)] : index - %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 - %c128 = index.constant 128 : index - %group_bytes = index.constant 144 : offset - %payload_byte_add = index.constant 16 : offset - %c0_f32 = scalar.constant 0.0 : f32 - %c1_f32 = scalar.constant 1.0 : f32 - %c127_f32 = scalar.constant 127.0 : f32 - %c0_f32x4 = vector.constant 0.0 : vector<4xf32> - %c0_offset = index.constant 0 : offset - %launched_element_count = index.mul %group_count, %c128 : index - %group_element_base = index.mul %group, %c128 : index - %lane_element_add = index.mul %lane, %c4 : index - %input_index = index.add %group_element_base, %lane_element_add : index - %input_noalias, %output_noalias = buffer.assume.noalias %input, %output : buffer, buffer - %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launched_element_count]xf32> - %input_values = vector.load %input_view[%input_index] : view<[%launched_element_count]xf32> -> vector<4xf32> - %absolute_values = vector.absf %input_values : vector<4xf32> - %thread_maximum = vector.reduce %absolute_values, %c0_f32 : vector<4xf32>, f32 - %block0 = index.div %lane, %c8 : index - %block = index.assume %block0 [range(%block0, 0, 3)] : index - %block_maxima = vector.insert %thread_maximum into %c0_f32x4[%block] : f32, vector<4xf32> - %reduced_maxima = kernel.subgroup.reduce %block_maxima : vector<4xf32> - %amax = vector.extract %reduced_maxima[%block] : vector<4xf32> -> f32 - %d = scalar.divf %amax, %c127_f32 : f32 - %d_nonzero = scalar.cmpf one, %d, %c0_f32 : f32 - %d_inverse = scf.if %d_nonzero -> (f32) { - %inverse = scalar.divf %c1_f32, %d : f32 - scf.yield %inverse : f32 - } else { - scf.yield %c0_f32 : f32 - } - %d_inverse_vector = vector.splat %d_inverse : vector<4xf32> - %scaled_values = vector.mulf %input_values, %d_inverse_vector : vector<4xf32> - %rounded_values = vector.roundf %scaled_values : vector<4xf32> - %quantized_values = vector.fptosi %rounded_values : vector<4xf32> to vector<4xi8> - %packed_word = vector.bitcast %quantized_values : vector<4xi8> to vector<1xi32> - %group_byte_offset = index.scale %group, %group_bytes : index, offset -> offset - %payload_byte_offset = index.add %group_byte_offset, %payload_byte_add : offset - %group_ds = buffer.view %output_noalias[%group_byte_offset] : buffer -> view<8xf16> - %group_qs = buffer.view %output_noalias[%payload_byte_offset] : buffer -> view<32xi32> - scf.if %publish_output { - vector.store %packed_word, %group_qs[%lane] : vector<1xi32>, view<32xi32> - } - %thread_sum = vector.reduce %rounded_values, %c0_f32 : vector<4xf32>, f32 - %block_sums0 = vector.insert %thread_sum into %c0_f32x4[%block] : f32, vector<4xf32> - %reduced_sums = kernel.subgroup.reduce %block_sums0 : vector<4xf32> - %quantized_sum = vector.extract %reduced_sums[%block] : vector<4xf32> -> f32 - %s = scalar.mulf %quantized_sum, %d : f32 - %word_in_block = index.rem %lane, %c8 : index - %is_block_leader = index.cmp eq, %word_in_block, %c0 : index - %publishes_metadata = scalar.andi %publish_output, %is_block_leader : i1 - scf.if %publishes_metadata { - %d_f16 = scalar.fptrunc %d : f32 to f16 - %s_f16 = scalar.fptrunc %s : f32 to f16 - %ds_index = index.mul %block, %c2 : index - %s_index = index.add %ds_index, %c1 : index - view.store %d_f16, %group_ds[%ds_index] : f16, view<8xf16> - view.store %s_f16, %group_ds[%s_index] : f16, view<8xf16> - } - template.return -} - -// Large-token schedule for routed Q4_K gate/up projections. Each workgroup -// owns one expert, 32 output channels, and up to 32 selected rows. It -// reads those rows from the compact expert table, stages one K=256 slice of -// both raw Q4_K projections and Q8_1 activations in LDS, and reuses each -// decoded weight chunk across four rows before a fused SwiGLU epilogue. -kernel.def @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped(%token_count: index, %route_count: index, %expert_count: index, %output_size: index) { - %token_capacity = config.get @qwen3_moe.workload.token_capacity : index - %configured_expert_count = config.get @qwen3_moe.routed_gate_up.expert_count : index - %configured_output_size = config.get @qwen3_moe.routed_gate_up.output_size : index - %c1 = index.constant 1 : index - %c4 = index.constant 4 : index - %c31 = index.constant 31 : index - %c32 = index.constant 32 : index - %workgroup_size = index.constant 256 : index - %padded_output_size = index.add %configured_output_size, %c31 : index - %output_tiles = index.div %padded_output_size, %c32 : index - %padded_token_count = index.add %token_capacity, %c31 : index - %route_tiles = index.div %padded_token_count, %c32 : index - %route_partitions = index.min %route_tiles, %c4 : index - kernel.launch.config workgroups(%output_tiles, %route_partitions, %configured_expert_count) workgroup_size(%workgroup_size, %c1, %c1) : index -} launch(%token_count: index, %route_count: index, %expert_count: index, %output_size: index, %q8_input: buffer, %expert_table: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) { - %input_size = config.get @qwen3_moe.routed_gate_up.input_size : index - %token_capacity = config.get @qwen3_moe.workload.token_capacity : index - %configured_route_count0 = config.get @qwen3_moe.routed_gate_up.route_count : index - %configured_expert_count0 = config.get @qwen3_moe.routed_gate_up.expert_count : index - %configured_output_size0 = config.get @qwen3_moe.routed_gate_up.output_size : index - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 512), le(%token_count, %token_capacity)] : index - %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 8), eq(%route_count, %configured_route_count0)] : index, index - %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 1, 128), eq(%expert_count, %configured_expert_count0)] : index, index - %bounded_output_size, %configured_output_size = index.assume %output_size, %configured_output_size0 [range(%output_size, 1, 4096), eq(%output_size, %configured_output_size0)] : index, index - %channel_tile = kernel.workgroup.id : index - %route_tile = kernel.workgroup.id : index - %expert0 = kernel.workgroup.id : index - %lane = kernel.workitem.id : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c4 = index.constant 4 : index - %c8 = index.constant 8 : index - %c32 = index.constant 32 : index - %c36 = index.constant 36 : index - %c72 = index.constant 72 : index - %c128 = index.constant 128 : index - %c256 = index.constant 256 : index - %gate_stage_word_count = index.constant 1152 : index - %q8_stage_word_count = index.constant 2304 : index - %c0_i32 = scalar.constant 0 : i32 - %c0_f32_rows = vector.constant 0.0 : vector<4xf32> - %c0_offset = index.constant 0 : offset - %c4_bytes = index.constant 4 : offset - %q4_block_bytes = index.constant 144 : offset - %q8_row_bytes = index.constant 288 : offset - %row_id_bytes = index.constant 128 : offset - %weight_stage_bytes = index.constant 4608 : offset - %q8_stage_bytes = index.constant 9216 : offset - %assignment_count = index.mul %bounded_token_count, %bounded_route_count : index - %expert, %table_expert_count = index.assume %expert0, %bounded_expert_count [lt(%expert0, %bounded_expert_count)] : index, index - %assignment_table_byte_base = index.scale %table_expert_count, %c4_bytes : index, offset -> offset - %q4_block_count = index.div %input_size, %c256 : index - %output_route_count = index.mul %bounded_token_count, %bounded_route_count : index - %q8_noalias, %expert_table_noalias, %gate_noalias, %up_noalias, %output_noalias = buffer.assume.noalias %q8_input, %expert_table, %gate_weight, %up_weight, %output : buffer, buffer, buffer, buffer, buffer - %q8_words = buffer.view %q8_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%q4_block_count]x72xi32> - %count_view = buffer.view %expert_table_noalias[%c0_offset] : buffer -> view<[%table_expert_count]xi32> - %assignment_view = buffer.view %expert_table_noalias[%assignment_table_byte_base] : buffer -> view<[%table_expert_count]x[%bounded_token_count]xi32> - %gate_words = buffer.view %gate_noalias[%c0_offset] : buffer -> view<[%bounded_expert_count]x[%bounded_output_size]x[%q4_block_count]x36xi32> - %up_words = buffer.view %up_noalias[%c0_offset] : buffer -> view<[%bounded_expert_count]x[%bounded_output_size]x[%q4_block_count]x36xi32> - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%output_route_count]x[%bounded_output_size]xf32> - %row_ids = buffer.alloca align(16) %row_id_bytes : buffer - %gate_stage = buffer.alloca align(16) %weight_stage_bytes : buffer - %up_stage = buffer.alloca align(16) %weight_stage_bytes : buffer - %q8_stage = buffer.alloca align(16) %q8_stage_bytes : buffer - %row_ids_view = buffer.view %row_ids[%c0_offset] : buffer -> view<32xi32> - %gate_stage_words = buffer.view %gate_stage[%c0_offset] : buffer -> view<1152xi32> - %up_stage_words = buffer.view %up_stage[%c0_offset] : buffer -> view<1152xi32> - %q8_stage_words = buffer.view %q8_stage[%c0_offset] : buffer -> view<2304xi32> - %channel_base = index.mul %channel_tile, %c32 : index - %channel0 = index.rem %lane, %c32 : index - %channel = index.add %channel_base, %channel0 : index - %row_base0 = index.div %lane, %c32 : index - %row_base = index.assume %row_base0 [range(%row_base0, 0, 7)] : index - %initial_tile_base = index.mul %route_tile, %c32 : index - // Four interleaved route partitions cover the first 128 rows. Smaller - // token counts execute at most one iteration; larger expert populations - // continue in 128-row strides without increasing the launch grid. - %route_partition_step = index.constant 128 : index - %is_lane_zero = index.cmp eq, %lane, %c0 : index - %lane_expert_route_count = scf.if %is_lane_zero -> (i32) { - %loaded = view.load %count_view[%expert] : view<[%table_expert_count]xi32> -> i32 - scf.yield %loaded : i32 - } else { - scf.yield %c0_i32 : i32 - } - %expert_route_count_reduced = kernel.workgroup.reduce %lane_expert_route_count : i32 - %expert_route_count_i32 = kernel.subgroup.broadcast.first %expert_route_count_reduced : i32 - %expert_route_count0 = index.cast %expert_route_count_i32 : i32 to index - %expert_route_count = index.assume %expert_route_count0 [range(%expert_route_count0, 0, 4096)] : index - scf.for %tile_base = [%initial_tile_base to %expert_route_count step %route_partition_step] { - %tile_base_i32 = index.cast %tile_base : index to i32 - %remaining_rows_i32 = scalar.subi %expert_route_count_i32, %tile_base_i32 : i32 - %remaining_rows0 = index.cast %remaining_rows_i32 : i32 to index - %remaining_rows = index.assume %remaining_rows0 [range(%remaining_rows0, 1, 4096)] : index - %has_full_tile = index.cmp uge, %remaining_rows, %c32 : index - %tile_row_count = scf.if %has_full_tile -> (index) { - scf.yield %c32 : index - } else { - scf.yield %remaining_rows : index - } - %loads_row = index.cmp ult, %lane, %tile_row_count : index - scf.if %loads_row { - %local_row = index.assume %lane [range(%lane, 0, 31)] : index - %expert_assignment_ordinal0 = index.add %tile_base, %local_row : index - %expert_assignment_ordinal, %table_token_count = index.assume %expert_assignment_ordinal0, %bounded_token_count [lt(%expert_assignment_ordinal0, %bounded_token_count)] : index, index - %assignment_i32 = view.load %assignment_view[%expert, %expert_assignment_ordinal] : view<[%table_expert_count]x[%bounded_token_count]xi32> -> i32 - view.store %assignment_i32, %row_ids_view[%local_row] : i32, view<32xi32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %gate_acc, %up_acc = scf.for %q4_block = [%c0 to %q4_block_count step %c1](%gate_block_acc = %c0_f32_rows : vector<4xf32>, %up_block_acc = %c0_f32_rows : vector<4xf32>) -> (vector<4xf32>, vector<4xf32>) { - %bounded_q4_block, %body_q4_block_count = index.assume %q4_block, %q4_block_count [lt(%q4_block, %q4_block_count)] : index, index - scf.for %stage_word = [%lane to %gate_stage_word_count step %c256] { - %local_channel = index.div %stage_word, %c36 : index - %word_in_block0 = index.rem %stage_word, %c36 : index - %word_in_block = index.assume %word_in_block0 [range(%word_in_block0, 0, 35)] : index - %global_channel = index.add %channel_base, %local_channel : index - %valid_channel = index.cmp ult, %global_channel, %bounded_output_size : index - %gate_word, %up_word = scf.if %valid_channel -> (i32, i32) { - %bounded_global_channel, %table_output_size = index.assume %global_channel, %bounded_output_size [lt(%global_channel, %bounded_output_size)] : index, index - %gate_loaded = view.load %gate_words[%expert, %bounded_global_channel, %bounded_q4_block, %word_in_block] : view<[%bounded_expert_count]x[%bounded_output_size]x[%q4_block_count]x36xi32> -> i32 - %up_loaded = view.load %up_words[%expert, %bounded_global_channel, %bounded_q4_block, %word_in_block] : view<[%bounded_expert_count]x[%bounded_output_size]x[%q4_block_count]x36xi32> -> i32 - scf.yield %gate_loaded, %up_loaded : i32, i32 - } else { - scf.yield %c0_i32, %c0_i32 : i32, i32 - } - view.store %gate_word, %gate_stage_words[%stage_word] : i32, view<1152xi32> - view.store %up_word, %up_stage_words[%stage_word] : i32, view<1152xi32> - } - scf.for %stage_word = [%lane to %q8_stage_word_count step %c256] { - %local_row = index.div %stage_word, %c72 : index - %word_in_row0 = index.rem %stage_word, %c72 : index - %word_in_row = index.assume %word_in_row0 [range(%word_in_row0, 0, 71)] : index - %valid_row = index.cmp ult, %local_row, %tile_row_count : index - %q8_word = scf.if %valid_row -> (i32) { - %bounded_local_row = index.assume %local_row [range(%local_row, 0, 31)] : index - %assignment_i32 = view.load %row_ids_view[%bounded_local_row] : view<32xi32> -> i32 - %assignment0 = index.cast %assignment_i32 : i32 to index - %assignment = index.assume %assignment0 [range(%assignment0, 0, 4095)] : index - %bounded_assignment, %bounded_assignment_count = index.assume %assignment, %assignment_count [lt(%assignment, %assignment_count)] : index, index - %token0 = index.div %bounded_assignment, %configured_route_count : index - %token, %table_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index - %loaded = view.load %q8_words[%token, %bounded_q4_block, %word_in_row] : view<[%bounded_token_count]x[%q4_block_count]x72xi32> -> i32 - scf.yield %loaded : i32 - } else { - scf.yield %c0_i32 : i32 - } - view.store %q8_word, %q8_stage_words[%stage_word] : i32, view<2304xi32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %channel_byte_base = index.scale %channel0, %q4_block_bytes : index, offset -> offset - %after_groups_gate, %after_groups_up = scf.for %q4_group = [%c0 to %c8 step %c1](%gate_group_acc = %gate_block_acc : vector<4xf32>, %up_group_acc = %up_block_acc : vector<4xf32>) -> (vector<4xf32>, vector<4xf32>) unroll { - %row1 = index.add %row_base, %c8 : index - %row2_add = index.constant 16 : index - %row2 = index.add %row_base, %row2_add : index - %row3_add = index.constant 24 : index - %row3 = index.add %row_base, %row3_add : index - %row0_byte_base = index.scale %row_base, %q8_row_bytes : index, offset -> offset - %row1_byte_base = index.scale %row1, %q8_row_bytes : index, offset -> offset - %row2_byte_base = index.scale %row2, %q8_row_bytes : index, offset -> offset - %row3_byte_base = index.scale %row3, %q8_row_bytes : index, offset -> offset - %q8_values0, %q8_d0, %q8_s0 = func.call @ggml_q8_1_x4_block(%q8_stage, %row0_byte_base, %q4_group) : (buffer, offset, index) -> (vector<32xi8>, f32, f32) - %q8_values1, %q8_d1, %q8_s1 = func.call @ggml_q8_1_x4_block(%q8_stage, %row1_byte_base, %q4_group) : (buffer, offset, index) -> (vector<32xi8>, f32, f32) - %q8_values2, %q8_d2, %q8_s2 = func.call @ggml_q8_1_x4_block(%q8_stage, %row2_byte_base, %q4_group) : (buffer, offset, index) -> (vector<32xi8>, f32, f32) - %q8_values3, %q8_d3, %q8_s3 = func.call @ggml_q8_1_x4_block(%q8_stage, %row3_byte_base, %q4_group) : (buffer, offset, index) -> (vector<32xi8>, f32, f32) - %q8_low0 = vector.slice %q8_values0[0] : vector<32xi8> -> vector<16xi8> - %q8_low1 = vector.slice %q8_values1[0] : vector<32xi8> -> vector<16xi8> - %q8_low2 = vector.slice %q8_values2[0] : vector<32xi8> -> vector<16xi8> - %q8_low3 = vector.slice %q8_values3[0] : vector<32xi8> -> vector<16xi8> - %q8_high0 = vector.slice %q8_values0[16] : vector<32xi8> -> vector<16xi8> - %q8_high1 = vector.slice %q8_values1[16] : vector<32xi8> -> vector<16xi8> - %q8_high2 = vector.slice %q8_values2[16] : vector<32xi8> -> vector<16xi8> - %q8_high3 = vector.slice %q8_values3[16] : vector<32xi8> -> vector<16xi8> - %q8_rows_low0 = vector.constant 0 : vector<4x16xi8> - %q8_rows_low1 = vector.insert %q8_low0 into %q8_rows_low0[0] : vector<16xi8>, vector<4x16xi8> - %q8_rows_low2 = vector.insert %q8_low1 into %q8_rows_low1[1] : vector<16xi8>, vector<4x16xi8> - %q8_rows_low3 = vector.insert %q8_low2 into %q8_rows_low2[2] : vector<16xi8>, vector<4x16xi8> - %q8_rows_low = vector.insert %q8_low3 into %q8_rows_low3[3] : vector<16xi8>, vector<4x16xi8> - %q8_rows_high0 = vector.constant 0 : vector<4x16xi8> - %q8_rows_high1 = vector.insert %q8_high0 into %q8_rows_high0[0] : vector<16xi8>, vector<4x16xi8> - %q8_rows_high2 = vector.insert %q8_high1 into %q8_rows_high1[1] : vector<16xi8>, vector<4x16xi8> - %q8_rows_high3 = vector.insert %q8_high2 into %q8_rows_high2[2] : vector<16xi8>, vector<4x16xi8> - %q8_rows_high = vector.insert %q8_high3 into %q8_rows_high3[3] : vector<16xi8>, vector<4x16xi8> - %q8_d = vector.from_elements %q8_d0, %q8_d1, %q8_d2, %q8_d3 : vector<4xf32> - %q8_s = vector.from_elements %q8_s0, %q8_s1, %q8_s2, %q8_s3 : vector<4xf32> - %gate_q4_low, %gate_d_scale_low, %gate_dmin_scale_low = func.call @qwen3_moe_q4k_chunk_local(%gate_stage, %channel_byte_base, %c0, %q4_group, %c0) : (buffer, offset, index, index, index) -> (vector<16xi8>, f32, f32) - %gate_q4_high, %gate_d_scale_high, %gate_dmin_scale_high = func.call @qwen3_moe_q4k_chunk_local(%gate_stage, %channel_byte_base, %c0, %q4_group, %c1) : (buffer, offset, index, index, index) -> (vector<16xi8>, f32, f32) - %up_q4_low, %up_d_scale_low, %up_dmin_scale_low = func.call @qwen3_moe_q4k_chunk_local(%up_stage, %channel_byte_base, %c0, %q4_group, %c0) : (buffer, offset, index, index, index) -> (vector<16xi8>, f32, f32) - %up_q4_high, %up_d_scale_high, %up_dmin_scale_high = func.call @qwen3_moe_q4k_chunk_local(%up_stage, %channel_byte_base, %c0, %q4_group, %c1) : (buffer, offset, index, index, index) -> (vector<16xi8>, f32, f32) - %gate_low = func.call @qwen3_moe_q4k_q8_1_dot4_rows(%gate_q4_low, %gate_d_scale_low, %gate_dmin_scale_low, %q8_rows_low, %q8_d, %q8_s) : (vector<16xi8>, f32, f32, vector<4x16xi8>, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) - %gate_high = func.call @qwen3_moe_q4k_q8_1_dot4_rows(%gate_q4_high, %gate_d_scale_high, %gate_dmin_scale_high, %q8_rows_high, %q8_d, %q8_s) : (vector<16xi8>, f32, f32, vector<4x16xi8>, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) - %up_low = func.call @qwen3_moe_q4k_q8_1_dot4_rows(%up_q4_low, %up_d_scale_low, %up_dmin_scale_low, %q8_rows_low, %q8_d, %q8_s) : (vector<16xi8>, f32, f32, vector<4x16xi8>, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) - %up_high = func.call @qwen3_moe_q4k_q8_1_dot4_rows(%up_q4_high, %up_d_scale_high, %up_dmin_scale_high, %q8_rows_high, %q8_d, %q8_s) : (vector<16xi8>, f32, f32, vector<4x16xi8>, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) - %gate_pair = vector.addf %gate_low, %gate_high : vector<4xf32> - %up_pair = vector.addf %up_low, %up_high : vector<4xf32> - %gate_next = vector.addf %gate_group_acc, %gate_pair : vector<4xf32> - %up_next = vector.addf %up_group_acc, %up_pair : vector<4xf32> - scf.yield %gate_next, %up_next : vector<4xf32>, vector<4xf32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - scf.yield %after_groups_gate, %after_groups_up : vector<4xf32>, vector<4xf32> - } - %gate_silu = vector.siluf %gate_acc : vector<4xf32> - %result_rows = vector.mulf %gate_silu, %up_acc : vector<4xf32> - %valid_channel = index.cmp ult, %channel, %bounded_output_size : index - scf.for %row_variant = [%c0 to %c4 step %c1] unroll { - %row_add = index.mul %row_variant, %c8 : index - %local_row = index.add %row_base, %row_add : index - %valid_row = index.cmp ult, %local_row, %tile_row_count : index - %writes_output = scalar.andi %valid_channel, %valid_row : i1 - scf.if %writes_output { - %bounded_local_row = index.assume %local_row [range(%local_row, 0, 31)] : index - %assignment_i32 = view.load %row_ids_view[%bounded_local_row] : view<32xi32> -> i32 - %assignment0 = index.cast %assignment_i32 : i32 to index - %assignment = index.assume %assignment0 [range(%assignment0, 0, 4095)] : index - %bounded_assignment, %bounded_assignment_count = index.assume %assignment, %assignment_count [lt(%assignment, %assignment_count)] : index, index - %result = vector.extract %result_rows[%row_variant] : vector<4xf32> -> f32 - view.store %result, %output_view[%bounded_assignment, %channel] : f32, view<[%output_route_count]x[%bounded_output_size]xf32> - } - } - } - kernel.return -} - -check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_nonzero_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(2048) : index - %route_count = check.literal value(1) : index - %route_stride = check.literal value(1) : index - %expert_count = check.literal value(1) : index - %output_size = check.literal value(1) : index - %input = check.generate.fill value(0.00390625) : tensor<1x2048xf32> - %q8_input = check.generate.fill value(0) : tensor<1x2304xi8> - %route_ids = check.generate.fill value(0) : tensor<1xi32> - %gate_weight = check.generate.fill value(85) : tensor<8x144xi8> - %up_weight = check.generate.fill value(-86) : tensor<8x144xi8> - %output = check.generate.fill value(0.0) : tensor<1xf32> - %expected = check.generate.fill value(-9024645.0) : tensor<1xf32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<1x2048xf32>, tensor<1x2304xi8>) - kernel.launch @ggml_mul_mat_id_vector_pair_binary_publish_f32[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %output) : [index, index, index, index, index](index, index, index, index, index, tensor<1x2304xi8>, tensor<1xi32>, tensor<8x144xi8>, tensor<8x144xi8>, tensor<1xf32>) - check.expect.close actual(%output) expected(%expected) atol(16.0) rtol(9.9999999999999995e-07) nan(same) : tensor<1xf32> - check.return -} - -// Exact decode topology with compact route IDs and every physical output -// group populated. Two fused invocations prove counter reuse in addition to -// comparing the F32 and packed outputs with the ordinary composition. -check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8_differential_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(2048) : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(768) : index - %routed_row_count = check.literal value(8) : index - %input = check.generate.fill value(0.00390625) : tensor<1x2048xf32> - %q8_input = check.generate.fill value(0) : tensor<2304xi8> - %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<8xi32> - %gate_weight = check.generate.iota offset(-72) step(1) period(144) : tensor<128x768x8x144xi8> - %up_weight = check.generate.iota offset(-71) step(1) period(144) : tensor<128x768x8x144xi8> - %expected_output = check.generate.fill value(0.0) : tensor<8x768xf32> - %expected_q8 = check.generate.fill value(0) : tensor<8x864xi8> - %actual_output0 = check.generate.fill value(1.0) : tensor<8x768xf32> - %actual_q8_0 = check.generate.fill value(1) : tensor<8x864xi8> - %actual_output1 = check.generate.fill value(2.0) : tensor<8x768xf32> - %actual_q8_1 = check.generate.fill value(2) : tensor<8x864xi8> - %completion_counters = check.generate.fill value(0) : tensor<48xi32> - %expected_counters = check.generate.fill value(0) : tensor<48xi32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<1x2048xf32>, tensor<2304xi8>) - kernel.launch @ggml_mul_mat_id_vector_pair_binary_publish_f32[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %expected_output) : [index, index, index, index, index](index, index, index, index, index, tensor<2304xi8>, tensor<8xi32>, tensor<128x768x8x144xi8>, tensor<128x768x8x144xi8>, tensor<8x768xf32>) - kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %output_size](%routed_row_count, %output_size, %expected_output, %expected_q8) : [index, index](index, index, tensor<8x768xf32>, tensor<8x864xi8>) - kernel.launch @ggml_mul_mat_id_vector_pair_binary_publish_q8[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %actual_output0, %completion_counters, %actual_q8_0) : [index, index, index, index, index](index, index, index, index, index, tensor<2304xi8>, tensor<8xi32>, tensor<128x768x8x144xi8>, tensor<128x768x8x144xi8>, tensor<8x768xf32>, tensor<48xi32>, tensor<8x864xi8>) - kernel.launch @ggml_mul_mat_id_vector_pair_binary_publish_q8[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %actual_output1, %completion_counters, %actual_q8_1) : [index, index, index, index, index](index, index, index, index, index, tensor<2304xi8>, tensor<8xi32>, tensor<128x768x8x144xi8>, tensor<128x768x8x144xi8>, tensor<8x768xf32>, tensor<48xi32>, tensor<8x864xi8>) - check.expect.close actual(%actual_output0) expected(%expected_output) atol(16.0) rtol(9.9999999999999995e-07) nan(same) : tensor<8x768xf32> - check.expect.close actual(%actual_output1) expected(%expected_output) atol(16.0) rtol(9.9999999999999995e-07) nan(same) : tensor<8x768xf32> - check.expect.equal actual(%actual_q8_0) expected(%expected_q8) : tensor<8x864xi8> - check.expect.equal actual(%actual_q8_1) expected(%expected_q8) : tensor<8x864xi8> - check.expect.equal actual(%completion_counters) expected(%expected_counters) : tensor<48xi32> - check.return -} - -check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_differential_case { - %token_count = check.literal value(64) : index - %input_size = check.literal value(512) : index - %route_count = check.literal value(2) : index - %route_stride = check.literal value(4) : index - %expert_count = check.literal value(4) : index - %output_size = check.literal value(32) : index - %input = check.generate.fill value(0.00390625) : tensor<64x512xf32> - %q8_input = check.generate.fill value(0) : tensor<64x576xi8> - // The physical row retains four argsort entries while the logical top-k view - // selects experts 0 and 1. Both experts span two 32-row grouped tiles. - %route_ids = check.generate.iota offset(0) step(1) period(4) : tensor<64x4xi32> - %expert_table = check.generate.fill value(-1) : tensor<260xi32> - %gate_weight = check.generate.iota offset(-72) step(1) period(144) : tensor<4x32x2x144xi8> - %up_weight = check.generate.iota offset(-71) step(1) period(144) : tensor<4x32x2x144xi8> - %expected = check.generate.fill value(0.0) : tensor<64x2x32xf32> - %actual = check.generate.fill value(1.0) : tensor<64x2x32xf32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<64x512xf32>, tensor<64x576xi8>) - kernel.launch @ggml_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<64x4xi32>, tensor<260xi32>) - kernel.launch @ggml_mul_mat_id_vector_pair_binary_publish_f32[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %expected) : [index, index, index, index, index](index, index, index, index, index, tensor<64x576xi8>, tensor<64x4xi32>, tensor<4x32x2x144xi8>, tensor<4x32x2x144xi8>, tensor<64x2x32xf32>) - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped[%token_count, %route_count, %expert_count, %output_size](%token_count, %route_count, %expert_count, %output_size, %q8_input, %expert_table, %gate_weight, %up_weight, %actual) : [index, index, index, index](index, index, index, index, tensor<64x576xi8>, tensor<260xi32>, tensor<4x32x2x144xi8>, tensor<4x32x2x144xi8>, tensor<64x2x32xf32>) - check.expect.close actual(%actual) expected(%expected) atol(0.25) rtol(9.9999999999999995e-07) nan(same) : tensor<64x2x32xf32> - check.return -} - -check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_tail_case { - %token_count = check.literal value(37) : index - %input_size = check.literal value(512) : index - %route_count = check.literal value(3) : index - %route_stride = check.literal value(5) : index - %expert_count = check.literal value(4) : index - %output_size = check.literal value(33) : index - %input = check.generate.fill value(0.00390625) : tensor<37x512xf32> - %q8_input = check.generate.fill value(0) : tensor<37x576xi8> - // Three unique selected experts rotate through a five-entry physical row. - // This leaves the second route tile empty and the second channel tile with - // only one live output channel. - %route_ids = check.generate.iota offset(0) step(1) period(4) : tensor<37x5xi32> - %expert_table = check.generate.fill value(-1) : tensor<152xi32> - %gate_weight = check.generate.iota offset(-72) step(1) period(144) : tensor<4x33x2x144xi8> - %up_weight = check.generate.iota offset(-71) step(1) period(144) : tensor<4x33x2x144xi8> - %expected = check.generate.fill value(0.0) : tensor<37x3x33xf32> - %actual = check.generate.fill value(1.0) : tensor<37x3x33xf32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<37x512xf32>, tensor<37x576xi8>) - kernel.launch @ggml_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<37x5xi32>, tensor<152xi32>) - kernel.launch @ggml_mul_mat_id_vector_pair_binary_publish_f32[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %expected) : [index, index, index, index, index](index, index, index, index, index, tensor<37x576xi8>, tensor<37x5xi32>, tensor<4x33x2x144xi8>, tensor<4x33x2x144xi8>, tensor<37x3x33xf32>) - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped[%token_count, %route_count, %expert_count, %output_size](%token_count, %route_count, %expert_count, %output_size, %q8_input, %expert_table, %gate_weight, %up_weight, %actual) : [index, index, index, index](index, index, index, index, tensor<37x576xi8>, tensor<152xi32>, tensor<4x33x2x144xi8>, tensor<4x33x2x144xi8>, tensor<37x3x33xf32>) - check.expect.close actual(%actual) expected(%expected) atol(0.25) rtol(9.9999999999999995e-07) nan(same) : tensor<37x3x33xf32> - check.return -} - -// Forces one route partition to process a second 128-row band. This models a -// maximally skewed router while retaining a noncompact physical route stride. -check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_partition_stride_case { - %token_count = check.literal value(129) : index - %input_size = check.literal value(512) : index - %route_count = check.literal value(1) : index - %route_stride = check.literal value(3) : index - %expert_count = check.literal value(2) : index - %output_size = check.literal value(1) : index - %input = check.generate.fill value(0.00390625) : tensor<129x512xf32> - %q8_input = check.generate.fill value(0) : tensor<129x576xi8> - %route_ids = check.generate.fill value(0) : tensor<129x3xi32> - %expert_table = check.generate.fill value(-1) : tensor<260xi32> - %gate_weight = check.generate.fill value(85) : tensor<2x1x2x144xi8> - %up_weight = check.generate.fill value(-86) : tensor<2x1x2x144xi8> - %expected = check.generate.fill value(0.0) : tensor<129x1x1xf32> - %actual = check.generate.fill value(1.0) : tensor<129x1x1xf32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<129x512xf32>, tensor<129x576xi8>) - kernel.launch @ggml_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<129x3xi32>, tensor<260xi32>) - kernel.launch @ggml_mul_mat_id_vector_pair_binary_publish_f32[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %expected) : [index, index, index, index, index](index, index, index, index, index, tensor<129x576xi8>, tensor<129x3xi32>, tensor<2x1x2x144xi8>, tensor<2x1x2x144xi8>, tensor<129x1x1xf32>) - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped[%token_count, %route_count, %expert_count, %output_size](%token_count, %route_count, %expert_count, %output_size, %q8_input, %expert_table, %gate_weight, %up_weight, %actual) : [index, index, index, index](index, index, index, index, tensor<129x576xi8>, tensor<260xi32>, tensor<2x1x2x144xi8>, tensor<2x1x2x144xi8>, tensor<129x1x1xf32>) - check.expect.close actual(%actual) expected(%expected) atol(0.25) rtol(9.9999999999999995e-07) nan(same) : tensor<129x1x1xf32> - check.return -} - -check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case { - %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512]) name("token_count") : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(128) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(768) : index - %q8_input = check.generate.fill value(0) : tensor<[%token_count]x2304xi8> - %route_ids = check.generate.iota offset(0) step(1) period(127) : tensor<[%token_count]x128xi32> - %gate_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %up_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x8x768xf32> - %expected = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf32> - kernel.launch @ggml_mul_mat_id_vector_pair_binary_publish_f32[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %output) : [index, index, index, index, index](index, index, index, index, index, tensor<[%token_count]x2304xi8>, tensor<[%token_count]x128xi32>, tensor<128x768x8x144xi8>, tensor<128x768x8x144xi8>, tensor<[%token_count]x8x768xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x8x768xf32> - check.return -} - -check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8_benchmark_case { - %token_count = check.literal value(1) : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(768) : index - %q8_input = check.generate.fill value(0) : tensor<2304xi8> - %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<8xi32> - %gate_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %up_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %output = check.generate.fill value(1.0) : tensor<8x768xf32> - %completion_counters = check.generate.fill value(0) : tensor<48xi32> - %next_q8_output = check.generate.fill value(1) : tensor<8x864xi8> - %expected_output = check.generate.fill value(0.0) : tensor<8x768xf32> - %expected_q8 = check.generate.fill value(0) : tensor<8x864xi8> - %expected_counters = check.generate.fill value(0) : tensor<48xi32> - kernel.launch @ggml_mul_mat_id_vector_pair_binary_publish_q8[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %output, %completion_counters, %next_q8_output) : [index, index, index, index, index](index, index, index, index, index, tensor<2304xi8>, tensor<8xi32>, tensor<128x768x8x144xi8>, tensor<128x768x8x144xi8>, tensor<8x768xf32>, tensor<48xi32>, tensor<8x864xi8>) - check.expect.close actual(%output) expected(%expected_output) atol(0.0) rtol(0.0) nan(same) : tensor<8x768xf32> - check.expect.equal actual(%next_q8_output) expected(%expected_q8) : tensor<8x864xi8> - check.expect.equal actual(%completion_counters) expected(%expected_counters) : tensor<48xi32> - check.return -} - -check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8_composed_benchmark_case { - %token_count = check.literal value(1) : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(768) : index - %routed_row_count = check.literal value(8) : index - %q8_input = check.generate.fill value(0) : tensor<2304xi8> - %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<8xi32> - %gate_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %up_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %output = check.generate.fill value(1.0) : tensor<8x768xf32> - %next_q8_output = check.generate.fill value(1) : tensor<8x864xi8> - %expected_output = check.generate.fill value(0.0) : tensor<8x768xf32> - %expected_q8 = check.generate.fill value(0) : tensor<8x864xi8> - kernel.launch @ggml_mul_mat_id_vector_pair_binary_publish_f32[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %output) : [index, index, index, index, index](index, index, index, index, index, tensor<2304xi8>, tensor<8xi32>, tensor<128x768x8x144xi8>, tensor<128x768x8x144xi8>, tensor<8x768xf32>) - kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %output_size](%routed_row_count, %output_size, %output, %next_q8_output) : [index, index](index, index, tensor<8x768xf32>, tensor<8x864xi8>) - check.expect.close actual(%output) expected(%expected_output) atol(0.0) rtol(0.0) nan(same) : tensor<8x768xf32> - check.expect.equal actual(%next_q8_output) expected(%expected_q8) : tensor<8x864xi8> - check.return -} - -check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case { - %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512]) name("token_count") : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(128) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(768) : index - %q8_input = check.generate.fill value(0) : tensor<[%token_count]x2304xi8> - // A period of 127 rotates each physical 128-entry row by one expert while - // keeping the first eight logical routes distinct within every token. - %route_ids = check.generate.iota offset(0) step(1) period(127) : tensor<[%token_count]x128xi32> - // One packed transient buffer holds 128 counts followed by room for 512 - // assignments per expert. The production allocation uses the exact token - // count; this fixed test capacity permits one parameterized benchmark case. - %expert_table = check.generate.fill value(-1) : tensor<65664xi32> - %gate_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %up_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x8x768xf32> - %expected = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf32> - kernel.launch @ggml_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x128xi32>, tensor<65664xi32>) - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped[%token_count, %route_count, %expert_count, %output_size](%token_count, %route_count, %expert_count, %output_size, %q8_input, %expert_table, %gate_weight, %up_weight, %output) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x2304xi8>, tensor<65664xi32>, tensor<128x768x8x144xi8>, tensor<128x768x8x144xi8>, tensor<[%token_count]x8x768xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x8x768xf32> - check.return -} - -// Maximally diverse decode-batch control. Compact route storage makes the -// flattened 0..127 iota assign every M=16 route to a different expert. -check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_diverse_benchmark_case { - %token_count = check.param.choice values([1, 2, 4, 8, 16]) name("token_count") : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(768) : index - %q8_input = check.generate.fill value(0) : tensor<[%token_count]x2304xi8> - %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<[%token_count]x8xi32> - %expert_table = check.generate.fill value(-1) : tensor<65664xi32> - %gate_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %up_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x8x768xf32> - %expected = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf32> - kernel.launch @ggml_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x8xi32>, tensor<65664xi32>) - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped[%token_count, %route_count, %expert_count, %output_size](%token_count, %route_count, %expert_count, %output_size, %q8_input, %expert_table, %gate_weight, %up_weight, %output) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x2304xi8>, tensor<65664xi32>, tensor<128x768x8x144xi8>, tensor<128x768x8x144xi8>, tensor<[%token_count]x8x768xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x8x768xf32> - check.return -} - -check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_pipeline_benchmark_case { - %token_count = check.param.choice values([1, 2, 4, 8, 16, 32, 128, 512]) name("token_count") : index - %input_size = check.literal value(2048) : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(128) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(768) : index - %input = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> - %q8_input = check.generate.fill value(1) : tensor<[%token_count]x2304xi8> - %route_ids = check.generate.iota offset(0) step(1) period(127) : tensor<[%token_count]x128xi32> - %gate_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %up_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x8x768xf32> - %expected = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<[%token_count]x2048xf32>, tensor<[%token_count]x2304xi8>) - kernel.launch @ggml_mul_mat_id_vector_pair_binary_publish_f32[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %output) : [index, index, index, index, index](index, index, index, index, index, tensor<[%token_count]x2304xi8>, tensor<[%token_count]x128xi32>, tensor<128x768x8x144xi8>, tensor<128x768x8x144xi8>, tensor<[%token_count]x8x768xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x8x768xf32> - check.return -} - -check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_benchmark_case { - %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512]) name("token_count") : index - %input_size = check.literal value(2048) : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(128) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(768) : index - %input = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> - %q8_input = check.generate.fill value(1) : tensor<[%token_count]x2304xi8> - %route_ids = check.generate.iota offset(0) step(1) period(127) : tensor<[%token_count]x128xi32> - %expert_table = check.generate.fill value(-1) : tensor<65664xi32> - %gate_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %up_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x8x768xf32> - %expected = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<[%token_count]x2048xf32>, tensor<[%token_count]x2304xi8>) - kernel.launch @ggml_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x128xi32>, tensor<65664xi32>) - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped[%token_count, %route_count, %expert_count, %output_size](%token_count, %route_count, %expert_count, %output_size, %q8_input, %expert_table, %gate_weight, %up_weight, %output) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x2304xi8>, tensor<65664xi32>, tensor<128x768x8x144xi8>, tensor<128x768x8x144xi8>, tensor<[%token_count]x8x768xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x8x768xf32> - check.return -} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_nonzero_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_small - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_decode {token_count = 1} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_small_batch_2 {token_count = 2} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8_decode - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8_composed_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8_composed_decode - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_small_batch_4 {token_count = 4} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_small_batch_8 {token_count = 8} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_small_batch_16 {token_count = 16} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_prefill_17 {token_count = 17} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_prefill_32 {token_count = 32} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_prefill_63 {token_count = 63} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_prefill_128 {token_count = 128} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_prefill_129 {token_count = 129} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_prefill_512 {token_count = 512} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_prefill_17 {token_count = 17} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_decode {token_count = 1} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_small_batch_2 {token_count = 2} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_small_batch_4 {token_count = 4} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_small_batch_8 {token_count = 8} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_small_batch_16 {token_count = 16} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_diverse_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_diverse_decode {token_count = 1} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_diverse_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_diverse_small_batch_2 {token_count = 2} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_diverse_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_diverse_small_batch_4 {token_count = 4} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_diverse_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_diverse_small_batch_8 {token_count = 8} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_diverse_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_diverse_small_batch_16 {token_count = 16} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_prefill_32 {token_count = 32} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_prefill_63 {token_count = 63} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_prefill_128 {token_count = 128} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_prefill_129 {token_count = 129} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_prefill_512 {token_count = 512} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_pipeline_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_pipeline_decode {token_count = 1} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_pipeline_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_pipeline_small_batch_8 {token_count = 8} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_pipeline_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_pipeline_prefill_32 {token_count = 32} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_pipeline_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_pipeline_prefill_128 {token_count = 128} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_pipeline_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_pipeline_prefill_512 {token_count = 512} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_prefill_17 {token_count = 17} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_prefill_32 {token_count = 32} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_prefill_63 {token_count = 63} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_prefill_128 {token_count = 128} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_prefill_129 {token_count = 129} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_prefill_512 {token_count = 512} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_linear_q4k_f16_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_linear_q4k_f16_wmma.loom deleted file mode 100644 index 59324f8fd3c2..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_linear_q4k_f16_wmma.loom +++ /dev/null @@ -1,949 +0,0 @@ -// Gfx11 routed-expert projection matching llama.cpp Vulkan's matrix path. -// -// Each two-wave workgroup computes 64 output channels for 32 routed rows of one -// expert. The waves share padded FP16 operand stages, each owns 32 output -// channels and four WMMA accumulators, and scatters one transposed accumulator -// fragment at a time through a compact route map. Gate and up invoke this same -// projection independently before the separate SwiGLU epilogue. -// -// The raw weights remain [expert][output channel][K / 256][144 bytes]. -// No persistent repacking or expanded-weight allocation is required. -amdgpu.target @qwen3_moe_gfx11_wave64 {subgroup_size = 64} - -amdgpu.target @qwen3_moe_gfx11_wave32 {subgroup_size = 32} - -config.decl @qwen3_moe.routed_gate_up.input_size : %value: index where [range(%value, 512, 32768), mul(%value, 512)] - -// Top-k is a model hyperparameter and is specialized with the kernel. Keeping -// it out of the dynamic workload lets address arithmetic and assignment decode -// fold to the exact model contract. -config.decl @qwen3_moe.routed_gate_up.route_count : %value: index where [range(%value, 1, 8)] - -config.decl @qwen3_moe.routed_gate_up.expert_count : %value: index where [range(%value, 1, 512)] - -config.decl @qwen3_moe.routed_gate_up.output_size : %value: index where [range(%value, 1, 4096)] - -kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count$4: index, %input_size$5: index) launch(%token_count$6: index, %input_size$7: index, %input: buffer, %output: buffer) - -kernel.decl @ggml_moe_build_expert_table(%token_count$10: index, %route_count$11: index, %route_stride$12: index, %expert_count$13: index) launch(%token_count$14: index, %route_count$15: index, %route_stride$16: index, %expert_count$17: index, %route_ids: buffer, %expert_table: buffer) - -kernel.decl @ggml_moe_build_expert_partition_table(%token_count$20: index, %route_count$21: index, %expert_count$22: index) launch(%token_count$23: index, %route_count$24: index, %expert_count$25: index, %expert_table: buffer, %partition_table: buffer) - -func.decl @qwen3_moe_unpack_expert_partition_descriptor(%descriptor: i32) -> (index, index, index) - -kernel.decl @ggml_mul_mat_id_vector_pair_binary_publish_f32(%token_count$32: index, %route_count$33: index, %route_stride$34: index, %expert_count$35: index, %output_size$36: index) launch(%token_count$37: index, %route_count$38: index, %route_stride$39: index, %expert_count$40: index, %output_size$41: index, %q8_input: buffer, %route_ids: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) - -func.decl @qwen3_moe_q4k_scale_from_header(%scale0: i32, %scale1: i32, %scale2: i32, %q4_group: index) -> (i32, i32) - -// Acquires the packed code word shared by one adjacent Q4_K group pair. -func.def inline @qwen3_moe_q4k_wmma_load_code(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group_pair: index, %packet: index) -> (vector<1xi32>) { - %c8 = index.constant 8 : index - %block_bytes = index.constant 144 : offset - %code_offset = index.constant 16 : offset - %bounded_group_pair = index.assume %q4_group_pair [range(%q4_group_pair, 0, 3)] : index - %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index - %block_byte_add = index.scale %q4_block, %block_bytes : index, offset -> offset - %block_byte_base = index.add %row_byte_base, %block_byte_add : offset - %code_byte_base = index.add %block_byte_base, %code_offset : offset - %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<32xi32> - %q_page = index.mul %bounded_group_pair, %c8 : index - %q_word_index0 = index.add %q_page, %bounded_packet : index - %q_word_index = index.assume %q_word_index0 [range(%q_word_index0, 0, 31)] : index - %q_word = vector.load %code_view[%q_word_index] : view<32xi32> -> vector<1xi32> - func.return %q_word : vector<1xi32> -} - -// Decodes the four adjacent Q4_K values owned by one load packet from an -// already-loaded block header and packed code word. Matrix schedules choose -// the lifetime of both immutable packets. -func.def inline @qwen3_moe_q4k_wmma_vector4_from_header_code(%q4_group: index, %header_words: vector<4xi32>, %q_word: vector<1xi32>) -> (vector<4xf16>) { - %c2 = index.constant 2 : index - %c4 = index.constant 4 : index - %q4_mask = vector.constant 252645135 : vector<1xi32> - %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index - %header_halves = vector.bitcast %header_words : vector<4xi32> to vector<8xf16> - %d_f16 = vector.extract %header_halves[0] : vector<8xf16> -> f16 - %dmin_f16 = vector.extract %header_halves[1] : vector<8xf16> -> f16 - %scale0 = vector.extract %header_words[1] : vector<4xi32> -> i32 - %scale1 = vector.extract %header_words[2] : vector<4xi32> -> i32 - %scale2 = vector.extract %header_words[3] : vector<4xi32> -> i32 - %d = scalar.extf %d_f16 : f16 to f32 - %dmin = scalar.extf %dmin_f16 : f16 to f32 - %scale, %minimum = func.call @qwen3_moe_q4k_scale_from_header(%scale0, %scale1, %scale2, %bounded_group) : (i32, i32, i32, index) -> (i32, i32) - %scale_f32 = scalar.uitofp %scale : i32 to f32 - %minimum_f32 = scalar.uitofp %minimum : i32 to f32 - %d_scale = scalar.mulf %d, %scale_f32 : f32 - %minimum_scale = scalar.mulf %dmin, %minimum_f32 : f32 - %q_half = index.rem %bounded_group, %c2 : index - %q_shift_index = index.mul %q_half, %c4 : index - %q_shift_i32 = index.cast %q_shift_index : index to i32 - %q_shift = vector.splat %q_shift_i32 : vector<1xi32> - %shifted_q = vector.shrui %q_word, %q_shift : vector<1xi32> - %masked_q = vector.andi %shifted_q, %q4_mask : vector<1xi32> - %q_i8 = vector.bitcast %masked_q : vector<1xi32> to vector<4xi8> - %q_f32 = vector.uitofp %q_i8 : vector<4xi8> to vector<4xf32> - // Form adjacent FP16 lanes from fused FP32 affine expressions. AMDGPU maps - // this natural shape to packed mixlo/mixhi instructions where available. - %negative_minimum_scale = scalar.negf %minimum_scale : f32 - %q0 = vector.extract %q_f32[0] : vector<4xf32> -> f32 - %q1 = vector.extract %q_f32[1] : vector<4xf32> -> f32 - %q2 = vector.extract %q_f32[2] : vector<4xf32> -> f32 - %q3 = vector.extract %q_f32[3] : vector<4xf32> -> f32 - %value0 = scalar.fmaf %q0, %d_scale, %negative_minimum_scale : f32 - %value1 = scalar.fmaf %q1, %d_scale, %negative_minimum_scale : f32 - %value2 = scalar.fmaf %q2, %d_scale, %negative_minimum_scale : f32 - %value3 = scalar.fmaf %q3, %d_scale, %negative_minimum_scale : f32 - %half0 = scalar.fptrunc %value0 : f32 to f16 - %half1 = scalar.fptrunc %value1 : f32 to f16 - %half2 = scalar.fptrunc %value2 : f32 to f16 - %half3 = scalar.fptrunc %value3 : f32 to f16 - %result = vector.from_elements %half0, %half1, %half2, %half3 : vector<4xf16> - func.return %result : vector<4xf16> -} - -// Decodes one group when its caller has retained only the Q4_K block header. -func.def inline @qwen3_moe_q4k_wmma_vector4_from_header(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group: index, %packet: index, %header_words: vector<4xi32>) -> (vector<4xf16>) { - %c2 = index.constant 2 : index - %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index - %q4_group_pair = index.div %bounded_group, %c2 : index - %q_word = func.call @qwen3_moe_q4k_wmma_load_code(%weight, %row_byte_base, %q4_block, %q4_group_pair, %packet) : (buffer, offset, index, index, index) -> (vector<1xi32>) - %values = func.call @qwen3_moe_q4k_wmma_vector4_from_header_code(%bounded_group, %header_words, %q_word) : (index, vector<4xi32>, vector<1xi32>) -> (vector<4xf16>) - func.return %values : vector<4xf16> -} - -// Acquires one naturally aligned Q4_K block header as a single 16-byte packet. -func.def inline @qwen3_moe_q4k_wmma_load_header(%weight: buffer, %row_byte_base: offset, %q4_block: index) -> (vector<4xi32>) { - %c0 = index.constant 0 : index - %block_bytes = index.constant 144 : offset - %block_byte_add = index.scale %q4_block, %block_bytes : index, offset -> offset - %block_byte_base = index.add %row_byte_base, %block_byte_add : offset - %header_view = buffer.view %weight[%block_byte_base] : buffer -> view<4xi32> - %header_words = vector.load %header_view[%c0] : view<4xi32> -> vector<4xi32> - func.return %header_words : vector<4xi32> -} - -// Acquires one block header before decoding the selected four-value group -// packet. Matrix schedules that span several groups call the two operations -// separately so the header lifetime matches their complete block loop. -func.def inline @qwen3_moe_q4k_wmma_vector4(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group: index, %packet: index) -> (vector<4xf16>) { - %header_words = func.call @qwen3_moe_q4k_wmma_load_header(%weight, %row_byte_base, %q4_block) : (buffer, offset, index) -> (vector<4xi32>) - %values = func.call @qwen3_moe_q4k_wmma_vector4_from_header(%weight, %row_byte_base, %q4_block, %q4_group, %packet, %header_words) : (buffer, offset, index, index, index, vector<4xi32>) -> (vector<4xf16>) - func.return %values : vector<4xf16> -} - -kernel.def target(@qwen3_moe_gfx11_wave64) @qwen3_moe_routed_linear_q4k_f16_wmma(%token_count: index) { - %expert_count = config.get @qwen3_moe.routed_gate_up.expert_count : index - %output_size = config.get @qwen3_moe.routed_gate_up.output_size : index - %c1 = index.constant 1 : index - %c63 = index.constant 63 : index - %c64 = index.constant 64 : index - %c128 = index.constant 128 : index - %padded_output_size = index.add %output_size, %c63 : index - %output_tiles = index.div %padded_output_size, %c64 : index - %padded_token_count = index.add %token_count, %c63 : index - %route_tiles = index.div %padded_token_count, %c64 : index - kernel.launch.config workgroups(%output_tiles, %route_tiles, %expert_count) workgroup_size(%c128, %c1, %c1) : index -} launch(%token_count: index, %input: buffer, %expert_table: buffer, %weight: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { - %input_size = config.get @qwen3_moe.routed_gate_up.input_size : index - %route_count = config.get @qwen3_moe.routed_gate_up.route_count : index - %expert_count = config.get @qwen3_moe.routed_gate_up.expert_count : index - %output_size = config.get @qwen3_moe.routed_gate_up.output_size : index - %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index - %bounded_expert_count = index.assume %expert_count [range(%expert_count, 1, 128)] : index - %bounded_output_size = index.assume %output_size [range(%output_size, 1, 4096)] : index - %channel_tile = kernel.workgroup.id : index - %route_tile = kernel.workgroup.id : index - %expert = kernel.workgroup.id : index - %workitem = kernel.workitem.id : index - %subgroup0 = kernel.subgroup.id : index - %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 1)] : index - %lane = kernel.subgroup.lane.id : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c4 = index.constant 4 : index - %c8 = index.constant 8 : index - %c16 = index.constant 16 : index - %c32 = index.constant 32 : index - %c40 = index.constant 40 : index - %c63 = index.constant 63 : index - %c64 = index.constant 64 : index - %c256 = index.constant 256 : index - %c0_offset = index.constant 0 : offset - %c4_bytes = index.constant 4 : offset - %q4_block_bytes = index.constant 144 : offset - %weight_stage_bytes = index.constant 5120 : offset - %activation_stage_bytes = index.constant 2560 : offset - %route_stage_bytes = index.constant 128 : offset - %wave_result_stage_bytes = index.constant 512 : offset - %result_stage_bytes = index.constant 1024 : offset - %c0_i32 = scalar.constant 0 : i32 - %cn1_i32 = scalar.constant -1 : i32 - %c0_f16x4 = vector.constant 0.0 : vector<4xf16> - %zero_accumulator = vector.constant 0.0 : vector<8xf16> - %m = index.constant 16 : index - %n = index.constant 16 : index - %k = index.constant 16 : index - %assignment_count = index.mul %token_count, %bounded_route_count : index - %assignment_table_byte_base = index.scale %bounded_expert_count, %c4_bytes : index, offset -> offset - %q4_block_count = index.div %input_size, %c256 : index - %weight_row_bytes = index.scale %q4_block_count, %q4_block_bytes : index, offset -> offset - %weight_expert_bytes = index.scale %bounded_output_size, %weight_row_bytes : index, offset -> offset - %output_row_count = index.mul %token_count, %bounded_route_count : index - %input_noalias, %expert_table_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %input, %expert_table, %weight, %output : buffer, buffer, buffer, buffer - %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%token_count]x[%input_size]xf32> - %count_view = buffer.view %expert_table_noalias[%c0_offset] : buffer -> view<[%bounded_expert_count]xi32> - %assignment_view = buffer.view %expert_table_noalias[%assignment_table_byte_base] : buffer -> view<[%bounded_expert_count]x[%token_count]xi32> - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%output_row_count]x[%bounded_output_size]xf32> - %weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer - %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer - %route_stage = buffer.alloca align(16) %route_stage_bytes : buffer - %result_stage = buffer.alloca align(16) %result_stage_bytes : buffer - %weight_stage_view = buffer.view %weight_stage[%c0_offset] : buffer -> view<64x40xf16> - %activation_stage_physical_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x40xf16> - %activation_fragment_layout = encoding.layout.strided [1, %c40] : encoding - %activation_fragment_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x32xf16, %activation_fragment_layout> - %route_stage_view = buffer.view %route_stage[%c0_offset] : buffer -> view<32xi32> - %wave_result_stage_offset = index.scale %subgroup, %wave_result_stage_bytes : index, offset -> offset - %result_fragment_layout = encoding.layout.strided [1, %c16] : encoding - %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 - %initial_route_tile_base = index.mul %route_tile, %c32 : index - %padded_token_count = index.add %token_count, %c63 : index - %route_partition_count = index.div %padded_token_count, %c64 : index - %route_partition_step = index.mul %route_partition_count, %c32 : index - %bounded_expert, %table_expert_count = index.assume %expert, %bounded_expert_count [lt(%expert, %bounded_expert_count)] : index, index - %is_workitem_zero = index.cmp eq, %workitem, %c0 : index - %lane_expert_route_count = scf.if %is_workitem_zero -> (i32) { - %loaded = view.load %count_view[%bounded_expert] : view<[%bounded_expert_count]xi32> -> i32 - scf.yield %loaded : i32 - } else { - scf.yield %c0_i32 : i32 - } - %expert_route_count_reduced = kernel.workgroup.reduce %lane_expert_route_count : i32 - %expert_route_count_i32 = kernel.subgroup.broadcast.first %expert_route_count_reduced : i32 - %expert_route_count0 = index.cast %expert_route_count_i32 : i32 to index - %expert_route_count = index.assume %expert_route_count0 [range(%expert_route_count0, 0, 2048)] : index - // Route partitions are distributed across the launch grid and continue in - // uniform strides for concentrated routing. Balanced Qwen prefill gives each - // expert 32, 64, or 128 rows at 512, 1024, or 2048 tokens. Concentrated - // experts can consume the full token count without changing the launch - // geometry. - scf.for %route_tile_base = [%initial_route_tile_base to %expert_route_count step %route_partition_step] { - // The first wave snapshots the compact route map once. Every K tile then - // reuses these 32 entries while the full workgroup cooperatively fills the - // operand stages. - %loads_route = index.cmp ult, %workitem, %c32 : index - scf.if %loads_route { - %local_route = index.assume %workitem [range(%workitem, 0, 31)] : index - %assignment_ordinal = index.add %route_tile_base, %local_route : index - %valid_row = index.cmp ult, %assignment_ordinal, %expert_route_count : index - %assignment_i32 = scf.if %valid_row -> (i32) { - %bounded_assignment_ordinal, %table_token_count = index.assume %assignment_ordinal, %token_count [lt(%assignment_ordinal, %token_count)] : index, index - %loaded = view.load %assignment_view[%bounded_expert, %bounded_assignment_ordinal] : view<[%bounded_expert_count]x[%token_count]xi32> -> i32 - scf.yield %loaded : i32 - } else { - scf.yield %cn1_i32 : i32 - } - view.store %assignment_i32, %route_stage_view[%local_route] : i32, view<32xi32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %load_packet = index.rem %workitem, %c8 : index - %load_k = index.mul %load_packet, %c4 : index - %load_row0 = index.div %workitem, %c8 : index - %load_row = index.assume %load_row0 [range(%load_row0, 0, 15)] : index - %expert_byte_base = index.scale %bounded_expert, %weight_expert_bytes : index, offset -> offset - %subgroup_channel_add = index.mul %subgroup, %c32 : index - %subgroup_channel1 = index.add %subgroup_channel_add, %c16 : index - %init00 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> - %init01 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> - %init10 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> - %init11 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> - %result00, %result01, %result10, %result11 = scf.for %q4_block = [%c0 to %q4_block_count 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>) { - %block_result00, %block_result01, %block_result10, %block_result11 = scf.for %q4_group = [%c0 to %c8 step %c1](%acc00 = %block_acc00 : vector<8xf16>, %acc01 = %block_acc01 : vector<8xf16>, %acc10 = %block_acc10 : vector<8xf16>, %acc11 = %block_acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { - %block_k_base = index.mul %q4_block, %c256 : index - %group_k_add = index.mul %q4_group, %c32 : index - %k_origin = index.add %block_k_base, %group_k_add : index - scf.for %row_offset = [%c0 to %c64 step %c16] unroll { - %local_row0 = index.add %load_row, %row_offset : index - %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>) { - %channel_byte_add = index.scale %channel, %weight_row_bytes : index, offset -> offset - %row_byte_base = index.add %expert_byte_base, %channel_byte_add : offset - %decoded = func.call @qwen3_moe_q4k_wmma_vector4(%weight_noalias, %row_byte_base, %q4_block, %q4_group, %load_packet) : (buffer, offset, index, index, index) -> (vector<4xf16>) - scf.yield %decoded : vector<4xf16> - } else { - scf.yield %c0_f16x4 : vector<4xf16> - } - %is_activation_row = index.cmp ult, %local_row, %c32 : index - vector.store %weight_values, %weight_stage_view[%local_row, %load_k] : vector<4xf16>, view<64x40xf16> - scf.if %is_activation_row { - %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>) { - %assignment0 = index.cast %assignment_i32 : i32 to index - %assignment = index.assume %assignment0 [range(%assignment0, 0, 16383)] : index - %bounded_assignment, %bounded_assignment_count = index.assume %assignment, %assignment_count [lt(%assignment, %assignment_count)] : index, index - %token0 = index.div %bounded_assignment, %bounded_route_count : index - %token, %input_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index - %input_k = index.add %k_origin, %load_k : index - %loaded = vector.load %input_view[%token, %input_k] : view<[%token_count]x[%input_size]xf32> -> vector<4xf32> - %converted = vector.fptrunc %loaded : vector<4xf32> to vector<4xf16> - scf.yield %converted : vector<4xf16> - } else { - scf.yield %c0_f16x4 : vector<4xf16> - } - vector.store %activation_values, %activation_stage_physical_view[%activation_row, %load_k] : vector<4xf16>, view<32x40xf16> - } - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %next00, %next01, %next10, %next11 = scf.for %k_half = [%c0 to %c32 step %c16](%half_acc00 = %acc00 : vector<8xf16>, %half_acc01 = %acc01 : vector<8xf16>, %half_acc10 = %acc10 : vector<8xf16>, %half_acc11 = %acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) unroll { - %lhs0 = vector.fragment.load %weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> - %lhs1 = vector.fragment.load %weight_stage_view[%subgroup_channel1, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> - %rhs0 = vector.fragment.load %activation_fragment_view[%k_half, %c0] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> - %rhs1 = vector.fragment.load %activation_fragment_view[%k_half, %c16] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> - %half_next00 = vector.mma %lhs0, %rhs0, %half_acc00 : vector<16xf16>, vector<16xf16>, vector<8xf16> - %half_next01 = vector.mma %lhs0, %rhs1, %half_acc01 : vector<16xf16>, vector<16xf16>, vector<8xf16> - %half_next10 = vector.mma %lhs1, %rhs0, %half_acc10 : vector<16xf16>, vector<16xf16>, vector<8xf16> - %half_next11 = vector.mma %lhs1, %rhs1, %half_acc11 : vector<16xf16>, vector<16xf16>, vector<8xf16> - scf.yield %half_next00, %half_next01, %half_next10, %half_next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - scf.yield %next00, %next01, %next10, %next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> - } - scf.yield %block_result00, %block_result01, %block_result10, %block_result11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> - } - // WMMA produces [channel][route] fragments. Store each through a - // transposed, wave-private LDS slice so every lane can scatter four - // contiguous channels to one routed output row. This uses the full wave - // for conversion and publication instead of serializing 16 channels onto - // each of 16 lanes. Subgroup-scoped fences suffice because waves never - // access each other's result slice. - %publish_route0 = index.div %lane, %c4 : index - %publish_route = index.assume %publish_route0 [range(%publish_route0, 0, 15)] : index - %publish_packet0 = index.rem %lane, %c4 : index - %publish_packet = index.assume %publish_packet0 [range(%publish_packet0, 0, 3)] : index - %publish_channel_add = index.mul %publish_packet, %c4 : index - %local_route1 = index.add %c16, %publish_route : index - %assignment0_i32 = view.load %route_stage_view[%publish_route] : view<32xi32> -> i32 - %assignment1_i32 = view.load %route_stage_view[%local_route1] : view<32xi32> -> i32 - %assignment0_nonnegative = scalar.cmpi sge, %assignment0_i32, %c0_i32 : i32 - %assignment1_nonnegative = scalar.cmpi sge, %assignment1_i32, %c0_i32 : i32 - %safe_assignment0_i32 = scf.if %assignment0_nonnegative -> (i32) { - scf.yield %assignment0_i32 : i32 - } else { - scf.yield %c0_i32 : i32 - } - %safe_assignment1_i32 = scf.if %assignment1_nonnegative -> (i32) { - scf.yield %assignment1_i32 : i32 - } else { - scf.yield %c0_i32 : i32 - } - %safe_assignment0_0 = index.cast %safe_assignment0_i32 : i32 to index - %safe_assignment1_0 = index.cast %safe_assignment1_i32 : i32 to index - %safe_assignment0 = index.assume %safe_assignment0_0 [range(%safe_assignment0_0, 0, 16383)] : index - %safe_assignment1 = index.assume %safe_assignment1_0 [range(%safe_assignment1_0, 0, 16383)] : index - %bounded_assignment0, %bounded_assignment_count0 = index.assume %safe_assignment0, %assignment_count [lt(%safe_assignment0, %assignment_count)] : index, index - %bounded_assignment1, %bounded_assignment_count1 = index.assume %safe_assignment1, %assignment_count [lt(%safe_assignment1, %assignment_count)] : index, index - %subgroup_channel_base = index.add %channel_tile_base, %subgroup_channel_add : index - %channel0 = index.add %subgroup_channel_base, %publish_channel_add : index - %channel1_base = index.add %subgroup_channel_base, %c16 : index - %channel1 = index.add %channel1_base, %publish_channel_add : index - %valid_channel0 = index.cmp ult, %channel0, %bounded_output_size : index - %valid_channel1 = index.cmp ult, %channel1, %bounded_output_size : index - %writes00 = scalar.andi %assignment0_nonnegative, %valid_channel0 : i1 - %writes01 = scalar.andi %assignment1_nonnegative, %valid_channel0 : i1 - %writes10 = scalar.andi %assignment0_nonnegative, %valid_channel1 : i1 - %writes11 = scalar.andi %assignment1_nonnegative, %valid_channel1 : i1 - vector.fragment.store %result00, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> - kernel.barrier scope(subgroup) ordering(acq_rel) - scf.if %writes00 { - %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> - %wide = vector.extf %values : vector<4xf16> to vector<4xf32> - %mask = vector.mask.range [%channel0 to %bounded_output_size step %c1] : index -> vector<4xi1> - vector.store.mask %wide, %output_view[%bounded_assignment0, %channel0], %mask : vector<4xf32>, view<[%output_row_count]x[%bounded_output_size]xf32>, vector<4xi1> - } - kernel.barrier scope(subgroup) ordering(acq_rel) - vector.fragment.store %result01, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> - kernel.barrier scope(subgroup) ordering(acq_rel) - scf.if %writes01 { - %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> - %wide = vector.extf %values : vector<4xf16> to vector<4xf32> - %mask = vector.mask.range [%channel0 to %bounded_output_size step %c1] : index -> vector<4xi1> - vector.store.mask %wide, %output_view[%bounded_assignment1, %channel0], %mask : vector<4xf32>, view<[%output_row_count]x[%bounded_output_size]xf32>, vector<4xi1> - } - kernel.barrier scope(subgroup) ordering(acq_rel) - vector.fragment.store %result10, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> - kernel.barrier scope(subgroup) ordering(acq_rel) - scf.if %writes10 { - %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> - %wide = vector.extf %values : vector<4xf16> to vector<4xf32> - %mask = vector.mask.range [%channel1 to %bounded_output_size step %c1] : index -> vector<4xi1> - vector.store.mask %wide, %output_view[%bounded_assignment0, %channel1], %mask : vector<4xf32>, view<[%output_row_count]x[%bounded_output_size]xf32>, vector<4xi1> - } - kernel.barrier scope(subgroup) ordering(acq_rel) - vector.fragment.store %result11, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> - kernel.barrier scope(subgroup) ordering(acq_rel) - scf.if %writes11 { - %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> - %wide = vector.extf %values : vector<4xf16> to vector<4xf32> - %mask = vector.mask.range [%channel1 to %bounded_output_size step %c1] : index -> vector<4xi1> - vector.store.mask %wide, %output_view[%bounded_assignment1, %channel1], %mask : vector<4xf32>, view<[%output_row_count]x[%bounded_output_size]xf32>, vector<4xi1> - } - // Route and result stages are reused by the next concentrated-routing - // partition. All waves must finish publication before either stage changes. - kernel.barrier scope(workgroup) ordering(acq_rel) - } - kernel.return -} - -// Fuses gate and up projection through SwiGLU for one routed expert tile. -// -// Eight waves split the 64 output channels and 32 routed rows into 16x16 result -// tiles. Each wave carries one gate and one up fragment while the workgroup -// shares one routed activation tile across both contractions. The final -// fragments meet in a wave-private LDS slice and publish the activated product -// directly, avoiding both full-size projection intermediates. -// -// The SwiGLU product is rounded to FP16 at its sole publication point because -// the grouped-down contraction consumes that exact FP16 WMMA operand. Keeping -// the transient in FP16 avoids a widen-store-load-truncate round trip and -// halves its global-memory footprint. -kernel.def target(@qwen3_moe_gfx11_wave32) @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma(%token_count: index) { - %route_count = config.get @qwen3_moe.routed_gate_up.route_count : index - %expert_count = config.get @qwen3_moe.routed_gate_up.expert_count : index - %output_size = config.get @qwen3_moe.routed_gate_up.output_size : index - %c1 = index.constant 1 : index - %c31 = index.constant 31 : index - %c32 = index.constant 32 : index - %c63 = index.constant 63 : index - %c64 = index.constant 64 : index - %c256 = index.constant 256 : index - %padded_output_size = index.add %output_size, %c63 : index - %output_tiles = index.div %padded_output_size, %c64 : index - // The exact partition count lives in device memory. The rounding bound keeps - // it below assignment_partitions + expert_count, so launching the larger - // term lets every workgroup consume at most two strided descriptors. - %assignment_count = index.mul %token_count, %route_count : index - %rounded_assignment_count = index.add %assignment_count, %c31 : index - %assignment_partition_count = index.div %rounded_assignment_count, %c32 : index - %has_more_assignment_partitions = index.cmp ugt, %assignment_partition_count, %expert_count : index - %launch_partition_count = scf.select %has_more_assignment_partitions, %assignment_partition_count, %expert_count : index - kernel.launch.config workgroups(%output_tiles, %launch_partition_count, %c1) workgroup_size(%c256, %c1, %c1) : index -} launch(%token_count: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { - %input_size = config.get @qwen3_moe.routed_gate_up.input_size : index - %route_count = config.get @qwen3_moe.routed_gate_up.route_count : index - %expert_count = config.get @qwen3_moe.routed_gate_up.expert_count : index - %output_size = config.get @qwen3_moe.routed_gate_up.output_size : index - %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index - %bounded_expert_count = index.assume %expert_count [range(%expert_count, 1, 128)] : index - %bounded_output_size = index.assume %output_size [range(%output_size, 1, 4096)] : index - %channel_tile = kernel.workgroup.id : index - %partition_ordinal = kernel.workgroup.id : index - %workitem = kernel.workitem.id : index - %subgroup0 = kernel.subgroup.id : index - %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index - %lane = kernel.subgroup.lane.id : index - %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 - %c31 = index.constant 31 : index - %c32 = index.constant 32 : index - %c40 = index.constant 40 : index - %c64 = index.constant 64 : index - %c256 = index.constant 256 : index - %c0_offset = index.constant 0 : offset - %c4_bytes = index.constant 4 : offset - %q4_block_bytes = index.constant 144 : offset - %weight_stage_bytes = index.constant 5120 : offset - %activation_stage_bytes = index.constant 2560 : offset - %route_stage_bytes = index.constant 128 : offset - %wave_result_stage_bytes = index.constant 512 : offset - %result_stage_bytes = index.constant 4096 : offset - %c0_i32 = scalar.constant 0 : i32 - %cn1_i32 = scalar.constant -1 : i32 - %c0_i32x1 = vector.constant 0 : vector<1xi32> - %c0_i32x4 = vector.constant 0 : vector<4xi32> - %c0_f16x4 = vector.constant 0.0 : vector<4xf16> - %zero_accumulator = vector.constant 0.0 : vector<16xf16> - %m = index.constant 16 : index - %n = index.constant 16 : index - %k = index.constant 16 : index - %assignment_count = index.mul %token_count, %bounded_route_count : index - %rounded_assignment_count = index.add %assignment_count, %c31 : index - %assignment_partition_count = index.div %rounded_assignment_count, %c32 : index - %maximum_partition_count = index.add %assignment_partition_count, %bounded_expert_count : index - %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 - %q4_block_count = index.div %input_size, %c256 : index - %weight_row_bytes = index.scale %q4_block_count, %q4_block_bytes : index, offset -> offset - %weight_expert_bytes = index.scale %bounded_output_size, %weight_row_bytes : index, offset -> offset - %output_row_count = index.mul %token_count, %bounded_route_count : index - %input_noalias, %expert_table_noalias, %partition_table_noalias, %gate_weight_noalias, %up_weight_noalias, %output_noalias = buffer.assume.noalias %input, %expert_table, %partition_table, %gate_weight, %up_weight, %output : buffer, buffer, buffer, buffer, buffer, buffer - %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%token_count]x[%input_size]xf32> - %assignment_view = buffer.view %expert_table_noalias[%assignment_table_byte_base] : buffer -> view<[%bounded_expert_count]x[%token_count]xi32> - %partition_count_view = buffer.view %partition_table_noalias[%c0_offset] : buffer -> view<1xi32> - %partition_descriptor_view = buffer.view %partition_table_noalias[%c4_bytes] : buffer -> view<[%maximum_partition_count]xi32> - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%output_row_count]x[%bounded_output_size]xf16> - %weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer - %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer - %route_stage = buffer.alloca align(16) %route_stage_bytes : buffer - %result_stage = buffer.alloca align(16) %result_stage_bytes : buffer - %weight_stage_view = buffer.view %weight_stage[%c0_offset] : buffer -> view<64x40xf16> - %activation_stage_physical_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x40xf16> - %activation_fragment_layout = encoding.layout.strided [1, %c40] : encoding - %activation_fragment_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x32xf16, %activation_fragment_layout> - %route_stage_view = buffer.view %route_stage[%c0_offset] : buffer -> view<32xi32> - %wave_result_stage_offset = index.scale %subgroup, %wave_result_stage_bytes : index, offset -> offset - %result_fragment_layout = encoding.layout.strided [1, %c16] : encoding - %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 - %is_workitem_zero = index.cmp eq, %workitem, %c0 : index - %lane_partition_count_i32 = scf.if %is_workitem_zero -> (i32) { - %loaded = view.load %partition_count_view[%c0] : view<1xi32> -> i32 - scf.yield %loaded : i32 - } else { - scf.yield %c0_i32 : i32 - } - %partition_count_reduced = kernel.workgroup.reduce %lane_partition_count_i32 : i32 - %partition_count_i32 = kernel.subgroup.broadcast.first %partition_count_reduced : i32 - %partition_count0 = index.cast %partition_count_i32 : i32 to index - %partition_count, %partition_capacity = index.assume %partition_count0, %maximum_partition_count [lt(%partition_count0, %maximum_partition_count)] : index, index - scf.for %active_partition = [%partition_ordinal to %partition_count step %launch_partition_count] { - %descriptor_ordinal, %descriptor_count = index.assume %active_partition, %partition_count [lt(%active_partition, %partition_count)] : index, index - %table_descriptor_ordinal, %table_descriptor_capacity = index.assume %descriptor_ordinal, %maximum_partition_count [lt(%descriptor_ordinal, %maximum_partition_count)] : index, index - %lane_descriptor_i32 = scf.if %is_workitem_zero -> (i32) { - %loaded = view.load %partition_descriptor_view[%table_descriptor_ordinal] : view<[%maximum_partition_count]xi32> -> i32 - scf.yield %loaded : i32 - } else { - scf.yield %c0_i32 : i32 - } - %descriptor_reduced_i32 = kernel.workgroup.reduce %lane_descriptor_i32 : i32 - %descriptor_i32 = kernel.subgroup.broadcast.first %descriptor_reduced_i32 : i32 - %expert, %route_tile_base, %partition_row_count = func.call @qwen3_moe_unpack_expert_partition_descriptor(%descriptor_i32) : (i32) -> (index, index, index) - %bounded_expert, %table_expert_count = index.assume %expert, %bounded_expert_count [lt(%expert, %bounded_expert_count)] : index, index - %loads_route = index.cmp ult, %workitem, %c32 : index - scf.if %loads_route { - %local_route = index.assume %workitem [range(%workitem, 0, 31)] : index - %assignment_ordinal = index.add %route_tile_base, %local_route : index - %valid_row = index.cmp ult, %local_route, %partition_row_count : index - %assignment_i32 = scf.if %valid_row -> (i32) { - %bounded_assignment_ordinal, %table_token_count = index.assume %assignment_ordinal, %token_count [lt(%assignment_ordinal, %token_count)] : index, index - %loaded = view.load %assignment_view[%bounded_expert, %bounded_assignment_ordinal] : view<[%bounded_expert_count]x[%token_count]xi32> -> i32 - scf.yield %loaded : i32 - } else { - scf.yield %cn1_i32 : i32 - } - view.store %assignment_i32, %route_stage_view[%local_route] : i32, view<32xi32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %load_packet = index.rem %workitem, %c8 : index - %load_k = index.mul %load_packet, %c4 : index - %load_row0 = index.div %workitem, %c8 : index - %load_row = index.assume %load_row0 [range(%load_row0, 0, 31)] : index - %expert_byte_base = index.scale %bounded_expert, %weight_expert_bytes : index, offset -> offset - %channel_subgroup = index.rem %subgroup, %c4 : index - %route_subgroup = index.div %subgroup, %c4 : index - %subgroup_channel_add = index.mul %channel_subgroup, %c16 : index - %subgroup_route_add = index.mul %route_subgroup, %c16 : index - %init_gate = vector.fragment %zero_accumulator shape [%m, %n] : vector<16xf16> - %init_up = vector.fragment %zero_accumulator shape [%m, %n] : vector<16xf16> - %gate_result, %up_result = scf.for %q4_block = [%c0 to %q4_block_count step %c1](%block_gate = %init_gate : vector<16xf16>, %block_up = %init_up : vector<16xf16>) -> (vector<16xf16>, vector<16xf16>) { - // Two lanes cover the 64 output rows owned by one load packet. Retain - // both projection headers across the complete eight-group block. - %gate_header0, %gate_header1, %up_header0, %up_header1 = scf.for %header_row_offset = [%c0 to %c64 step %c32](%prior_gate_header0 = %c0_i32x4 : vector<4xi32>, %prior_gate_header1 = %c0_i32x4 : vector<4xi32>, %prior_up_header0 = %c0_i32x4 : vector<4xi32>, %prior_up_header1 = %c0_i32x4 : vector<4xi32>) -> (vector<4xi32>, vector<4xi32>, vector<4xi32>, vector<4xi32>) unroll { - %header_local_row0 = index.add %load_row, %header_row_offset : index - %header_local_row = index.assume %header_local_row0 [range(%header_local_row0, 0, 63)] : index - %header_channel = index.add %channel_tile_base, %header_local_row : index - %valid_header_channel = index.cmp ult, %header_channel, %bounded_output_size : index - %loaded_gate_header, %loaded_up_header = scf.if %valid_header_channel -> (vector<4xi32>, vector<4xi32>) { - %header_channel_byte_add = index.scale %header_channel, %weight_row_bytes : index, offset -> offset - %header_row_byte_base = index.add %expert_byte_base, %header_channel_byte_add : offset - %gate_header = func.call @qwen3_moe_q4k_wmma_load_header(%gate_weight_noalias, %header_row_byte_base, %q4_block) : (buffer, offset, index) -> (vector<4xi32>) - %up_header = func.call @qwen3_moe_q4k_wmma_load_header(%up_weight_noalias, %header_row_byte_base, %q4_block) : (buffer, offset, index) -> (vector<4xi32>) - scf.yield %gate_header, %up_header : vector<4xi32>, vector<4xi32> - } else { - scf.yield %c0_i32x4, %c0_i32x4 : vector<4xi32>, vector<4xi32> - } - %updates_header0 = index.cmp eq, %header_row_offset, %c0 : index - %next_gate_header0 = scf.select %updates_header0, %loaded_gate_header, %prior_gate_header0 : vector<4xi32> - %next_gate_header1 = scf.select %updates_header0, %prior_gate_header1, %loaded_gate_header : vector<4xi32> - %next_up_header0 = scf.select %updates_header0, %loaded_up_header, %prior_up_header0 : vector<4xi32> - %next_up_header1 = scf.select %updates_header0, %prior_up_header1, %loaded_up_header : vector<4xi32> - scf.yield %next_gate_header0, %next_gate_header1, %next_up_header0, %next_up_header1 : vector<4xi32>, vector<4xi32>, vector<4xi32>, vector<4xi32> - } - %next_block_gate, %next_block_up = scf.for %q4_group_pair = [%c0 to %c4 step %c1](%pair_gate_acc = %block_gate : vector<16xf16>, %pair_up_acc = %block_up : vector<16xf16>) -> (vector<16xf16>, vector<16xf16>) { - // Each code word supplies the low and high nibbles for one adjacent - // group pair. Retain both projections' words for exactly those uses. - %gate_q_word0, %gate_q_word1, %up_q_word0, %up_q_word1 = scf.for %code_row_offset = [%c0 to %c64 step %c32](%prior_gate_q_word0 = %c0_i32x1 : vector<1xi32>, %prior_gate_q_word1 = %c0_i32x1 : vector<1xi32>, %prior_up_q_word0 = %c0_i32x1 : vector<1xi32>, %prior_up_q_word1 = %c0_i32x1 : vector<1xi32>) -> (vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>) unroll { - %code_local_row0 = index.add %load_row, %code_row_offset : index - %code_local_row = index.assume %code_local_row0 [range(%code_local_row0, 0, 63)] : index - %code_channel = index.add %channel_tile_base, %code_local_row : index - %valid_code_channel = index.cmp ult, %code_channel, %bounded_output_size : index - %loaded_gate_q_word, %loaded_up_q_word = scf.if %valid_code_channel -> (vector<1xi32>, vector<1xi32>) { - %bounded_q4_group_pair = index.assume %q4_group_pair [range(%q4_group_pair, 0, 3)] : index - %code_channel_byte_add = index.scale %code_channel, %weight_row_bytes : index, offset -> offset - %code_row_byte_base = index.add %expert_byte_base, %code_channel_byte_add : offset - %gate_q_word = func.call @qwen3_moe_q4k_wmma_load_code(%gate_weight_noalias, %code_row_byte_base, %q4_block, %bounded_q4_group_pair, %load_packet) : (buffer, offset, index, index, index) -> (vector<1xi32>) - %up_q_word = func.call @qwen3_moe_q4k_wmma_load_code(%up_weight_noalias, %code_row_byte_base, %q4_block, %bounded_q4_group_pair, %load_packet) : (buffer, offset, index, index, index) -> (vector<1xi32>) - scf.yield %gate_q_word, %up_q_word : vector<1xi32>, vector<1xi32> - } else { - scf.yield %c0_i32x1, %c0_i32x1 : vector<1xi32>, vector<1xi32> - } - %updates_q_word0 = index.cmp eq, %code_row_offset, %c0 : index - %next_gate_q_word0 = scf.select %updates_q_word0, %loaded_gate_q_word, %prior_gate_q_word0 : vector<1xi32> - %next_gate_q_word1 = scf.select %updates_q_word0, %prior_gate_q_word1, %loaded_gate_q_word : vector<1xi32> - %next_up_q_word0 = scf.select %updates_q_word0, %loaded_up_q_word, %prior_up_q_word0 : vector<1xi32> - %next_up_q_word1 = scf.select %updates_q_word0, %prior_up_q_word1, %loaded_up_q_word : vector<1xi32> - scf.yield %next_gate_q_word0, %next_gate_q_word1, %next_up_q_word0, %next_up_q_word1 : vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32> - } - %next_pair_gate, %next_pair_up = scf.for %group_within_pair = [%c0 to %c2 step %c1](%group_gate = %pair_gate_acc : vector<16xf16>, %group_up = %pair_up_acc : vector<16xf16>) -> (vector<16xf16>, vector<16xf16>) { - %q4_group_base = index.mul %q4_group_pair, %c2 : index - %q4_group0 = index.add %q4_group_base, %group_within_pair : index - %q4_group = index.assume %q4_group0 [range(%q4_group0, 0, 7)] : index - %block_k_base = index.mul %q4_block, %c256 : index - %group_k_add = index.mul %q4_group, %c32 : index - %k_origin = index.add %block_k_base, %group_k_add : index - scf.for %row_offset = [%c0 to %c64 step %c32] unroll { - %local_row0 = index.add %load_row, %row_offset : index - %local_row = index.assume %local_row0 [range(%local_row0, 0, 63)] : index - %selects_first_row = index.cmp eq, %row_offset, %c0 : index - %selected_gate_header = scf.select %selects_first_row, %gate_header0, %gate_header1 : vector<4xi32> - %selected_gate_q_word = scf.select %selects_first_row, %gate_q_word0, %gate_q_word1 : vector<1xi32> - %channel = index.add %channel_tile_base, %local_row : index - %valid_channel = index.cmp ult, %channel, %bounded_output_size : index - %gate_values = scf.if %valid_channel -> (vector<4xf16>) { - %decoded = func.call @qwen3_moe_q4k_wmma_vector4_from_header_code(%q4_group, %selected_gate_header, %selected_gate_q_word) : (index, vector<4xi32>, vector<1xi32>) -> (vector<4xf16>) - scf.yield %decoded : vector<4xf16> - } else { - scf.yield %c0_f16x4 : vector<4xf16> - } - vector.store %gate_values, %weight_stage_view[%local_row, %load_k] : vector<4xf16>, view<64x40xf16> - } - %assignment_i32 = view.load %route_stage_view[%load_row] : view<32xi32> -> i32 - %valid_assignment = scalar.cmpi sge, %assignment_i32, %c0_i32 : i32 - %activation_values = scf.if %valid_assignment -> (vector<4xf16>) { - %assignment0 = index.cast %assignment_i32 : i32 to index - %assignment = index.assume %assignment0 [range(%assignment0, 0, 16383)] : index - %bounded_assignment, %bounded_assignment_count = index.assume %assignment, %assignment_count [lt(%assignment, %assignment_count)] : index, index - %token0 = index.div %bounded_assignment, %bounded_route_count : index - %token, %input_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index - %input_k = index.add %k_origin, %load_k : index - %loaded = vector.load %input_view[%token, %input_k] : view<[%token_count]x[%input_size]xf32> -> vector<4xf32> - %converted = vector.fptrunc %loaded : vector<4xf32> to vector<4xf16> - scf.yield %converted : vector<4xf16> - } else { - scf.yield %c0_f16x4 : vector<4xf16> - } - vector.store %activation_values, %activation_stage_physical_view[%load_row, %load_k] : vector<4xf16>, view<32x40xf16> - kernel.barrier scope(workgroup) ordering(acq_rel) - %gate_next = scf.for %k_half = [%c0 to %c32 step %c16](%half_gate = %group_gate : vector<16xf16>) -> (vector<16xf16>) unroll { - %lhs = vector.fragment.load %weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> - %rhs = vector.fragment.load %activation_fragment_view[%k_half, %subgroup_route_add] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> - %next = vector.mma %lhs, %rhs, %half_gate : vector<16xf16>, vector<16xf16>, vector<16xf16> - scf.yield %next : vector<16xf16> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - scf.for %row_offset = [%c0 to %c64 step %c32] unroll { - %local_row0 = index.add %load_row, %row_offset : index - %local_row = index.assume %local_row0 [range(%local_row0, 0, 63)] : index - %selects_first_row = index.cmp eq, %row_offset, %c0 : index - %selected_up_header = scf.select %selects_first_row, %up_header0, %up_header1 : vector<4xi32> - %selected_up_q_word = scf.select %selects_first_row, %up_q_word0, %up_q_word1 : vector<1xi32> - %channel = index.add %channel_tile_base, %local_row : index - %valid_channel = index.cmp ult, %channel, %bounded_output_size : index - %up_values = scf.if %valid_channel -> (vector<4xf16>) { - %decoded = func.call @qwen3_moe_q4k_wmma_vector4_from_header_code(%q4_group, %selected_up_header, %selected_up_q_word) : (index, vector<4xi32>, vector<1xi32>) -> (vector<4xf16>) - scf.yield %decoded : vector<4xf16> - } else { - scf.yield %c0_f16x4 : vector<4xf16> - } - vector.store %up_values, %weight_stage_view[%local_row, %load_k] : vector<4xf16>, view<64x40xf16> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %up_next = scf.for %k_half = [%c0 to %c32 step %c16](%half_up = %group_up : vector<16xf16>) -> (vector<16xf16>) unroll { - %lhs = vector.fragment.load %weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> - %rhs = vector.fragment.load %activation_fragment_view[%k_half, %subgroup_route_add] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> - %next = vector.mma %lhs, %rhs, %half_up : vector<16xf16>, vector<16xf16>, vector<16xf16> - scf.yield %next : vector<16xf16> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - scf.yield %gate_next, %up_next : vector<16xf16>, vector<16xf16> - } - scf.yield %next_pair_gate, %next_pair_up : vector<16xf16>, vector<16xf16> - } - scf.yield %next_block_gate, %next_block_up : vector<16xf16>, vector<16xf16> - } - %publish_route0 = index.div %lane, %c2 : index - %publish_route = index.assume %publish_route0 [range(%publish_route0, 0, 15)] : index - %publish_packet0 = index.rem %lane, %c2 : index - %publish_packet = index.assume %publish_packet0 [range(%publish_packet0, 0, 1)] : index - %publish_channel_add = index.mul %publish_packet, %c8 : index - %local_route = index.add %subgroup_route_add, %publish_route : index - %assignment_i32 = view.load %route_stage_view[%local_route] : view<32xi32> -> i32 - %assignment_nonnegative = scalar.cmpi sge, %assignment_i32, %c0_i32 : i32 - %safe_assignment_i32 = scf.if %assignment_nonnegative -> (i32) { - scf.yield %assignment_i32 : i32 - } else { - scf.yield %c0_i32 : i32 - } - %safe_assignment0 = index.cast %safe_assignment_i32 : i32 to index - %safe_assignment = index.assume %safe_assignment0 [range(%safe_assignment0, 0, 16383)] : index - %bounded_assignment, %bounded_assignment_count = index.assume %safe_assignment, %assignment_count [lt(%safe_assignment, %assignment_count)] : index, index - %subgroup_channel_base = index.add %channel_tile_base, %subgroup_channel_add : index - %channel = index.add %subgroup_channel_base, %publish_channel_add : index - %valid_channel = index.cmp ult, %channel, %bounded_output_size : index - %writes = scalar.andi %assignment_nonnegative, %valid_channel : i1 - vector.fragment.store %gate_result, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<16xf16>, view<16x16xf16, %result_fragment_layout> - kernel.barrier scope(subgroup) ordering(acq_rel) - %gate_values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<8xf16> - %gate_wide = vector.extf %gate_values : vector<8xf16> to vector<8xf32> - %activated = vector.siluf %gate_wide : vector<8xf32> - kernel.barrier scope(subgroup) ordering(acq_rel) - vector.fragment.store %up_result, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<16xf16>, view<16x16xf16, %result_fragment_layout> - kernel.barrier scope(subgroup) ordering(acq_rel) - scf.if %writes { - %up_values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<8xf16> - %up_wide = vector.extf %up_values : vector<8xf16> to vector<8xf32> - %wide_values = vector.mulf %activated, %up_wide : vector<8xf32> - %values = vector.fptrunc %wide_values : vector<8xf32> to vector<8xf16> - %mask = vector.mask.range [%channel to %bounded_output_size step %c1] : index -> vector<8xi1> - vector.store.mask %values, %output_view[%bounded_assignment, %channel], %mask : vector<8xf16>, view<[%output_row_count]x[%bounded_output_size]xf16>, vector<8xi1> - } - kernel.barrier scope(subgroup) ordering(acq_rel) - } - kernel.return -} - -// Applies the model's SwiGLU epilogue to gate and up projections already -// scattered into logical [token][route][output channel] order. -kernel.def target(@qwen3_moe_gfx11_wave64) @qwen3_moe_routed_swiglu_f32(%token_count: index) { - %c1 = index.constant 1 : index - %workgroup_size = index.constant 256 : index - %rounding = index.constant 255 : index - %output_size = config.get @qwen3_moe.routed_gate_up.output_size : index - %route_count = config.get @qwen3_moe.routed_gate_up.route_count : index - %row_count = index.mul %token_count, %route_count : index - %element_count = index.mul %row_count, %output_size : index - %rounded_count = index.add %element_count, %rounding : index - %workgroup_count = index.div %rounded_count, %workgroup_size : index - kernel.launch.config workgroups(%workgroup_count, %c1, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index -} launch(%token_count: index, %gate: buffer, %up: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { - %route_count = config.get @qwen3_moe.routed_gate_up.route_count : index - %output_size = config.get @qwen3_moe.routed_gate_up.output_size : index - %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index - %bounded_output_size = index.assume %output_size [range(%output_size, 1, 4096)] : index - %workgroup = kernel.workgroup.id : index - %lane = kernel.workitem.id : index - %workgroup_size = index.constant 256 : index - %c0_offset = index.constant 0 : offset - %row_count = index.mul %token_count, %bounded_route_count : index - %element_count = index.mul %row_count, %bounded_output_size : index - %element = index.madd %workgroup, %workgroup_size, %lane : index - %in_bounds = index.cmp ult, %element, %element_count : index - %gate_noalias, %up_noalias, %output_noalias = buffer.assume.noalias %gate, %up, %output : buffer, buffer, buffer - scf.if %in_bounds { - %bounded_element, %view_element_count = index.assume %element, %element_count [lt(%element, %element_count)] : index, index - %gate_view = buffer.view %gate_noalias[%c0_offset] : buffer -> view<[%view_element_count]xf32> - %up_view = buffer.view %up_noalias[%c0_offset] : buffer -> view<[%view_element_count]xf32> - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%view_element_count]xf32> - %gate_value = view.load %gate_view[%bounded_element] : view<[%view_element_count]xf32> -> f32 - %up_value = view.load %up_view[%bounded_element] : view<[%view_element_count]xf32> -> f32 - %activated = scalar.siluf %gate_value : f32 - %result = scalar.mulf %activated, %up_value : f32 - view.store %result, %output_view[%bounded_element] : f32, view<[%view_element_count]xf32> - } - kernel.return -} - -// Reference FP16 publication path for checking the fused gate/up provider at -// its actual model boundary. Gate and up projections remain independently -// materialized so this schedule does not share the fused kernel's matrix or -// publication implementation. -kernel.def target(@qwen3_moe_gfx11_wave64) @qwen3_moe_routed_swiglu_f16(%token_count: index) { - %c1 = index.constant 1 : index - %workgroup_size = index.constant 256 : index - %rounding = index.constant 255 : index - %output_size = config.get @qwen3_moe.routed_gate_up.output_size : index - %route_count = config.get @qwen3_moe.routed_gate_up.route_count : index - %row_count = index.mul %token_count, %route_count : index - %element_count = index.mul %row_count, %output_size : index - %rounded_count = index.add %element_count, %rounding : index - %workgroup_count = index.div %rounded_count, %workgroup_size : index - kernel.launch.config workgroups(%workgroup_count, %c1, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index -} launch(%token_count: index, %gate: buffer, %up: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { - %route_count = config.get @qwen3_moe.routed_gate_up.route_count : index - %output_size = config.get @qwen3_moe.routed_gate_up.output_size : index - %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index - %bounded_output_size = index.assume %output_size [range(%output_size, 1, 4096)] : index - %workgroup = kernel.workgroup.id : index - %lane = kernel.workitem.id : index - %workgroup_size = index.constant 256 : index - %c0_offset = index.constant 0 : offset - %row_count = index.mul %token_count, %bounded_route_count : index - %element_count = index.mul %row_count, %bounded_output_size : index - %element = index.madd %workgroup, %workgroup_size, %lane : index - %in_bounds = index.cmp ult, %element, %element_count : index - %gate_noalias, %up_noalias, %output_noalias = buffer.assume.noalias %gate, %up, %output : buffer, buffer, buffer - scf.if %in_bounds { - %bounded_element, %view_element_count = index.assume %element, %element_count [lt(%element, %element_count)] : index, index - %gate_view = buffer.view %gate_noalias[%c0_offset] : buffer -> view<[%view_element_count]xf32> - %up_view = buffer.view %up_noalias[%c0_offset] : buffer -> view<[%view_element_count]xf32> - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%view_element_count]xf16> - %gate_value = view.load %gate_view[%bounded_element] : view<[%view_element_count]xf32> -> f32 - %up_value = view.load %up_view[%bounded_element] : view<[%view_element_count]xf32> -> f32 - %activated = scalar.siluf %gate_value : f32 - %wide_result = scalar.mulf %activated, %up_value : f32 - %result = scalar.fptrunc %wide_result : f32 to f16 - view.store %result, %output_view[%bounded_element] : f16, view<[%view_element_count]xf16> - } - kernel.return -} - -// The constant input is exactly representable in Q8_1 and FP16. Comparing -// against the established integer-dot path therefore isolates Q4_K decode, -// routing, FP16 matrix accumulation, tail handling, and the SwiGLU boundary. -check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_differential_case { - %token_count = check.literal value(67) : index - %input_size = check.literal value(512) : index - %route_count = check.literal value(2) : index - %route_stride = check.literal value(4) : index - %expert_count = check.literal value(4) : index - %output_size = check.literal value(33) : index - %input = check.generate.fill value(0.00390625) : tensor<67x512xf32> - %q8_input = check.generate.fill value(0) : tensor<67x576xi8> - %route_ids = check.generate.iota offset(0) step(1) period(4) : tensor<67x4xi32> - %expert_table = check.generate.fill value(-1) : tensor<272xi32> - %partition_table = check.generate.fill value(-1) : tensor<10xi32> - %gate_weight = check.generate.fill value(34) : tensor<4x33x2x144xi8> - %up_weight = check.generate.fill value(35) : tensor<4x33x2x144xi8> - %gate_projection = check.generate.fill value(0.0) : tensor<67x2x33xf32> - %up_projection = check.generate.fill value(0.0) : tensor<67x2x33xf32> - %expected = check.generate.fill value(0.0) : tensor<67x2x33xf32> - %actual = check.generate.fill value(1.0) : tensor<67x2x33xf32> - %actual_f16 = check.generate.fill value(1.0) : tensor<67x2x33xf16> - %fused_actual = check.generate.fill value(1.0) : tensor<67x2x33xf16> - kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<67x512xf32>, tensor<67x576xi8>) - kernel.launch @ggml_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<67x4xi32>, tensor<272xi32>) - kernel.launch @ggml_moe_build_expert_partition_table[%token_count, %route_count, %expert_count](%token_count, %route_count, %expert_count, %expert_table, %partition_table) : [index, index, index](index, index, index, tensor<272xi32>, tensor<10xi32>) - kernel.launch @ggml_mul_mat_id_vector_pair_binary_publish_f32[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %expected) : [index, index, index, index, index](index, index, index, index, index, tensor<67x576xi8>, tensor<67x4xi32>, tensor<4x33x2x144xi8>, tensor<4x33x2x144xi8>, tensor<67x2x33xf32>) - kernel.launch @qwen3_moe_routed_linear_q4k_f16_wmma[%token_count](%token_count, %input, %expert_table, %gate_weight, %gate_projection) : [index](index, tensor<67x512xf32>, tensor<272xi32>, tensor<4x33x2x144xi8>, tensor<67x2x33xf32>) - kernel.launch @qwen3_moe_routed_linear_q4k_f16_wmma[%token_count](%token_count, %input, %expert_table, %up_weight, %up_projection) : [index](index, tensor<67x512xf32>, tensor<272xi32>, tensor<4x33x2x144xi8>, tensor<67x2x33xf32>) - kernel.launch @qwen3_moe_routed_swiglu_f32[%token_count](%token_count, %gate_projection, %up_projection, %actual) : [index](index, tensor<67x2x33xf32>, tensor<67x2x33xf32>, tensor<67x2x33xf32>) - kernel.launch @qwen3_moe_routed_swiglu_f16[%token_count](%token_count, %gate_projection, %up_projection, %actual_f16) : [index](index, tensor<67x2x33xf32>, tensor<67x2x33xf32>, tensor<67x2x33xf16>) - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma[%token_count](%token_count, %input, %expert_table, %partition_table, %gate_weight, %up_weight, %fused_actual) : [index](index, tensor<67x512xf32>, tensor<272xi32>, tensor<10xi32>, tensor<4x33x2x144xi8>, tensor<4x33x2x144xi8>, tensor<67x2x33xf16>) - check.expect.close actual(%actual) expected(%expected) atol(0.01) rtol(0.01) nan(same) : tensor<67x2x33xf32> - check.expect.close actual(%fused_actual) expected(%actual_f16) atol(0.01) rtol(0.01) nan(same) : tensor<67x2x33xf16> - check.return -} - -check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case { - %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512, 1024, 2048]) name("token_count") : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(768) : index - %input = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> - // The compact iota exactly matches the oracle's balanced ring: - // expert(token, route) = (8 * token + route) % 128. - %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<[%token_count]x8xi32> - // One count per expert followed by 2048 assignment slots per expert. - %expert_table = check.generate.fill value(-1) : tensor<262272xi32> - %gate_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %up_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %gate_projection = check.generate.fill value(1.0) : tensor<[%token_count]x8x768xf32> - %up_projection = check.generate.fill value(1.0) : tensor<[%token_count]x8x768xf32> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x8x768xf32> - %expected = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf32> - kernel.launch @ggml_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x8xi32>, tensor<262272xi32>) - kernel.launch @qwen3_moe_routed_linear_q4k_f16_wmma[%token_count](%token_count, %input, %expert_table, %gate_weight, %gate_projection) : [index](index, tensor<[%token_count]x2048xf32>, tensor<262272xi32>, tensor<128x768x8x144xi8>, tensor<[%token_count]x8x768xf32>) - kernel.launch @qwen3_moe_routed_linear_q4k_f16_wmma[%token_count](%token_count, %input, %expert_table, %up_weight, %up_projection) : [index](index, tensor<[%token_count]x2048xf32>, tensor<262272xi32>, tensor<128x768x8x144xi8>, tensor<[%token_count]x8x768xf32>) - kernel.launch @qwen3_moe_routed_swiglu_f32[%token_count](%token_count, %gate_projection, %up_projection, %output) : [index](index, tensor<[%token_count]x8x768xf32>, tensor<[%token_count]x8x768xf32>, tensor<[%token_count]x8x768xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x8x768xf32> - check.return -} - -check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case { - %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512, 1024, 2048]) name("token_count") : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(768) : index - %input = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> - %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<[%token_count]x8xi32> - // One count per expert followed by 2048 assignment slots per expert. - %expert_table = check.generate.fill value(-1) : tensor<262272xi32> - // One exact count followed by at most 640 packed descriptors. - %partition_table = check.generate.fill value(-1) : tensor<641xi32> - %gate_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %up_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x8x768xf16> - %expected = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf16> - kernel.launch @ggml_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x8xi32>, tensor<262272xi32>) - kernel.launch @ggml_moe_build_expert_partition_table[%token_count, %route_count, %expert_count](%token_count, %route_count, %expert_count, %expert_table, %partition_table) : [index, index, index](index, index, index, tensor<262272xi32>, tensor<641xi32>) - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma[%token_count](%token_count, %input, %expert_table, %partition_table, %gate_weight, %up_weight, %output) : [index](index, tensor<[%token_count]x2048xf32>, tensor<262272xi32>, tensor<641xi32>, tensor<128x768x8x144xi8>, tensor<128x768x8x144xi8>, tensor<[%token_count]x8x768xf16>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x8x768xf16> - check.return -} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_differential_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_differential - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_decode {token_count = 1} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_small_batch_2 {token_count = 2} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_small_batch_4 {token_count = 4} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_small_batch_8 {token_count = 8} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_small_batch_16 {token_count = 16} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_prefill_17 {token_count = 17} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_prefill_32 {token_count = 32} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_prefill_63 {token_count = 63} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_prefill_128 {token_count = 128} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_prefill_129 {token_count = 129} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_prefill_512 {token_count = 512} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_prefill_1024 {token_count = 1024} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_prefill_2048 {token_count = 2048} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_decode {token_count = 1} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_small_batch_2 {token_count = 2} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_small_batch_4 {token_count = 4} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_small_batch_8 {token_count = 8} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_small_batch_16 {token_count = 16} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_prefill_17 {token_count = 17} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_prefill_32 {token_count = 32} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_prefill_63 {token_count = 63} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_prefill_128 {token_count = 128} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_prefill_129 {token_count = 129} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_prefill_512 {token_count = 512} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_prefill_1024 {token_count = 1024} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_prefill_2048 {token_count = 2048} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/router_projection_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/router_projection_f32.loom index be653b373509..44e464e29634 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/router_projection_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/router_projection_f32.loom @@ -16,25 +16,12 @@ amdgpu.target @qwen3_moe_router_projection_gfx11_wave32 {subgroup amdgpu.target @qwen3_moe_router_projection_gfx11_wave64 {subgroup_size = 64} -amdgpu.target @qwen3_moe_router_projection_gfx1151 +target.decl @qwen3_moe_router_projection_gfx1151 config.decl @qwen3_moe.model.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] config.decl @qwen3_moe.router.expert_count : %value: index where [range(%value, 32, 512), mul(%value, 32)] -// Selects the target's preferred storage contract without exposing addressing -// policy at the projection call site. -template.def<@qwen3_moe.router_projection.storage> target(@qwen3_moe_router_projection_gfx1151) priority(20) @qwen3_moe_router_projection_descriptor_storage(%input: buffer, %weight: buffer, %output: buffer) -> (buffer, buffer, buffer) { - %input_descriptor = buffer.assume.memory_space %input : buffer - %weight_descriptor = buffer.assume.memory_space %weight : buffer - %output_descriptor = buffer.assume.memory_space %output : buffer - template.return %input_descriptor, %weight_descriptor, %output_descriptor : buffer, buffer, buffer -} - -template.def<@qwen3_moe.router_projection.storage> priority(1) @qwen3_moe_router_projection_global_storage(%input: buffer, %weight: buffer, %output: buffer) -> (buffer, buffer, buffer) { - template.return %input, %weight, %output : buffer, buffer, buffer -} - // Scalar differential oracle. It is reachable only from check cases and keeps // production validation independent of the packetized wave schedule. kernel.def target(@qwen3_moe_router_projection_gfx11_wave32) @qwen3_moe_router_projection_f32_reference(%token_count: index) { diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/router_projection_top8_fused_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/router_projection_top8_fused_f32.loom index b637d7801788..b542f9219452 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/router_projection_top8_fused_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/router_projection_top8_fused_f32.loom @@ -18,7 +18,7 @@ template.decl @qwen3_moe.router_projection.storage(%arg0: buffer, %arg1: buffer, amdgpu.target @qwen3_moe_router_fused_gfx11_wave64 {subgroup_size = 64} -amdgpu.target @qwen3_moe_router_fused_gfx1151 +target.decl @qwen3_moe_router_fused_gfx1151 config.decl @qwen3_moe.model.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] @@ -28,88 +28,7 @@ kernel.decl @qwen3_moe_router_projection_f32_four_row_wave32(%token_count$26: in kernel.decl @qwen3_moe_router_top8_f32(%token_count$31: index, %route_id_stride$32: index) launch(%token_count$33: index, %route_id_stride$34: index, %logits: buffer, %route_ids: buffer, %route_weights: buffer) -template.def<@qwen3_moe.router.fused.experts_per_wave> target(@qwen3_moe_router_fused_gfx1151) requires [#target.subgroup.size<64>] priority(20) @qwen3_moe_router_fused_two_experts_per_wave() -> (index) { - %c2 = index.constant 2 : index - template.return %c2 : index -} - -template.def<@qwen3_moe.router.fused.experts_per_wave> requires [#target.subgroup.size<64>] priority(1) @qwen3_moe_router_fused_four_experts_per_wave() -> (index) { - %c4 = index.constant 4 : index - template.return %c4 : index -} - -template.def<@qwen3_moe.router.fused.projection> device target(@qwen3_moe_router_fused_gfx1151) requires [#target.subgroup.size<64>] priority(20) @qwen3_moe_router_fused_two_expert_projection(%hidden_size: index, %expert_count: index, %token_count: index, %token: index, %expert_tile: index, %lane: index, %writes_projection: i1, %input: buffer, %weight: buffer, %logits: buffer) { - %c0_offset = index.constant 0 : offset - %c1 = index.constant 1 : index - %c2 = index.constant 2 : index - %c4 = index.constant 4 : index - %c256 = index.constant 256 : index - %c0_f32 = scalar.constant 0.0 : f32 - %expert_base = index.mul %expert_tile, %c2 : index - %expert1 = index.add %expert_base, %c1 : index - %lane_channel = index.mul %lane, %c4 : index - %input_view = buffer.view %input[%c0_offset] : buffer -> view<[%token_count]x[%hidden_size]xf32> - %weight_view = buffer.view %weight[%c0_offset] : buffer -> view<[%expert_count]x[%hidden_size]xf32> - %logits_view = buffer.view %logits[%c0_offset] : buffer -> view<[%token_count]x[%expert_count]xf32> - %c0_f32x2 = vector.splat %c0_f32 : vector<2xf32> - %lane_sums = scf.for %channel = [%lane_channel to %hidden_size step %c256](%accumulators = %c0_f32x2 : vector<2xf32>) -> (vector<2xf32>) unroll(%c4) schedule(interleaved) { - %input_values = vector.load %input_view[%token, %channel] : view<[%token_count]x[%hidden_size]xf32> -> vector<4xf32> - %weight_values0 = vector.load %weight_view[%expert_base, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> - %weight_values1 = vector.load %weight_view[%expert1, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> - %accumulator0 = vector.extract %accumulators[0] : vector<2xf32> -> f32 - %accumulator1 = vector.extract %accumulators[1] : vector<2xf32> -> f32 - %next0 = vector.dotf %input_values, %weight_values0, %accumulator0 : vector<4xf32>, vector<4xf32>, f32 - %next1 = vector.dotf %input_values, %weight_values1, %accumulator1 : vector<4xf32>, vector<4xf32>, f32 - %next_accumulators = vector.from_elements %next0, %next1 : vector<2xf32> - scf.yield %next_accumulators : vector<2xf32> - } - %sums = kernel.subgroup.reduce %lane_sums : vector<2xf32> - scf.if %writes_projection { - vector.store %sums, %logits_view[%token, %expert_base] : vector<2xf32>, view<[%token_count]x[%expert_count]xf32> - } - template.return -} - -template.def<@qwen3_moe.router.fused.projection> device requires [#target.subgroup.size<64>] priority(1) @qwen3_moe_router_fused_four_expert_projection(%hidden_size: index, %expert_count: index, %token_count: index, %token: index, %expert_tile: index, %lane: index, %writes_projection: i1, %input: buffer, %weight: buffer, %logits: buffer) { - %c0_offset = index.constant 0 : offset - %c1 = index.constant 1 : index - %c2 = index.constant 2 : index - %c3 = index.constant 3 : index - %c4 = index.constant 4 : index - %c256 = index.constant 256 : index - %c0_f32 = scalar.constant 0.0 : f32 - %expert_base = index.mul %expert_tile, %c4 : index - %expert1 = index.add %expert_base, %c1 : index - %expert2 = index.add %expert_base, %c2 : index - %expert3 = index.add %expert_base, %c3 : index - %lane_channel = index.mul %lane, %c4 : index - %input_view = buffer.view %input[%c0_offset] : buffer -> view<[%token_count]x[%hidden_size]xf32> - %weight_view = buffer.view %weight[%c0_offset] : buffer -> view<[%expert_count]x[%hidden_size]xf32> - %logits_view = buffer.view %logits[%c0_offset] : buffer -> view<[%token_count]x[%expert_count]xf32> - %c0_f32x4 = vector.splat %c0_f32 : vector<4xf32> - %lane_sums = scf.for %channel = [%lane_channel to %hidden_size step %c256](%accumulators = %c0_f32x4 : vector<4xf32>) -> (vector<4xf32>) unroll(%c4) schedule(interleaved) { - %input_values = vector.load %input_view[%token, %channel] : view<[%token_count]x[%hidden_size]xf32> -> vector<4xf32> - %weight_values0 = vector.load %weight_view[%expert_base, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> - %weight_values1 = vector.load %weight_view[%expert1, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> - %weight_values2 = vector.load %weight_view[%expert2, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> - %weight_values3 = vector.load %weight_view[%expert3, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> - %accumulator0 = vector.extract %accumulators[0] : vector<4xf32> -> f32 - %accumulator1 = vector.extract %accumulators[1] : vector<4xf32> -> f32 - %accumulator2 = vector.extract %accumulators[2] : vector<4xf32> -> f32 - %accumulator3 = vector.extract %accumulators[3] : vector<4xf32> -> f32 - %next0 = vector.dotf %input_values, %weight_values0, %accumulator0 : vector<4xf32>, vector<4xf32>, f32 - %next1 = vector.dotf %input_values, %weight_values1, %accumulator1 : vector<4xf32>, vector<4xf32>, f32 - %next2 = vector.dotf %input_values, %weight_values2, %accumulator2 : vector<4xf32>, vector<4xf32>, f32 - %next3 = vector.dotf %input_values, %weight_values3, %accumulator3 : vector<4xf32>, vector<4xf32>, f32 - %next_accumulators = vector.from_elements %next0, %next1, %next2, %next3 : vector<4xf32> - scf.yield %next_accumulators : vector<4xf32> - } - %sums = kernel.subgroup.reduce %lane_sums : vector<4xf32> - scf.if %writes_projection { - vector.store %sums, %logits_view[%token, %expert_base] : vector<4xf32>, view<[%token_count]x[%expert_count]xf32> - } - template.return -} +kernel.decl @qwen3_moe_router_top8_f32_check_one() launch(%logits: buffer, %route_ids: buffer, %route_weights: buffer) kernel.def target(@qwen3_moe_router_fused_gfx11_wave64) @qwen3_moe_router_projection_top8_fused_decode_f32(%token_count: index, %route_id_stride: index) { %expert_count = config.get @qwen3_moe.router.expert_count : index @@ -191,7 +110,7 @@ check.case public @qwen3_moe_router_projection_top8_fused_differential_case { %completion_counter = check.generate.fill value(0) : tensor<1xi32> %expected_counter = check.generate.fill value(0) : tensor<1xi32> kernel.launch @qwen3_moe_router_projection_f32_four_row_wave32[%token_count](%token_count, %input, %weight, %expected_logits) : [index](index, tensor<1x2048xf32>, tensor<128x2048xf32>, tensor<1x128xf32>) - kernel.launch @qwen3_moe_router_top8_f32[%token_count, %route_id_stride](%token_count, %route_id_stride, %expected_logits, %expected_route_ids, %expected_route_weights) : [index, index](index, index, tensor<1x128xf32>, tensor<1x8xi32>, tensor<1x8xf32>) + kernel.launch @qwen3_moe_router_top8_f32_check_one(%expected_logits, %expected_route_ids, %expected_route_weights) : (tensor<1x128xf32>, tensor<1x8xi32>, tensor<1x8xf32>) kernel.launch @qwen3_moe_router_projection_top8_fused_decode_f32[%token_count, %route_id_stride](%token_count, %route_id_stride, %input, %weight, %actual_logits0, %completion_counter, %actual_route_ids0, %actual_route_weights0) : [index, index](index, index, tensor<1x2048xf32>, tensor<128x2048xf32>, tensor<1x128xf32>, tensor<1xi32>, tensor<1x8xi32>, tensor<1x8xf32>) kernel.launch @qwen3_moe_router_projection_top8_fused_decode_f32[%token_count, %route_id_stride](%token_count, %route_id_stride, %input, %weight, %actual_logits1, %completion_counter, %actual_route_ids1, %actual_route_weights1) : [index, index](index, index, tensor<1x2048xf32>, tensor<128x2048xf32>, tensor<1x128xf32>, tensor<1xi32>, tensor<1x8xi32>, tensor<1x8xf32>) check.expect.close actual(%actual_logits0) expected(%expected_logits) atol(0.001) rtol(0.001) nan(same) : tensor<1x128xf32> @@ -230,6 +149,8 @@ check.case public @qwen3_moe_router_projection_top8_composed_benchmark_case { check.return } +check.benchmark<@qwen3_moe_router_projection_top8_fused_differential_case> @qwen3_moe_router_projection_top8_fused_differential + check.benchmark<@qwen3_moe_router_projection_top8_fused_benchmark_case> @qwen3_moe_router_projection_top8_fused_decode check.benchmark<@qwen3_moe_router_projection_top8_composed_benchmark_case> @qwen3_moe_router_projection_top8_composed_decode diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/router_top8_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/router_top8_f32.loom index 1ecccc71ba44..88e404695f0c 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/router_top8_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/router_top8_f32.loom @@ -24,98 +24,6 @@ config.decl @qwen3_moe.router.expert_count : %value: index where [range(%value, config.decl @qwen3_moe.router.route_count : %value: index where [range(%value, 1, 32)] -// Selects and normalizes one logical router row. The caller owns mapping a -// subgroup to a safe physical row and tells the helper whether that row should -// publish. Keeping row semantics independent of launch geometry lets fused -// producers consume this exact tie-breaking and normalization contract. -template.def<@qwen3_moe.router.top8.row> device requires [#target.subgroup.size<64>] @qwen3_moe_router_top8_row_f32(%valid_token: i1, %token: index, %token_count: index, %route_id_stride: index, %logits: buffer, %route_ids: buffer, %route_weights: buffer) { - %expert_count0 = config.get @qwen3_moe.router.expert_count : index - %route_count0 = config.get @qwen3_moe.router.route_count : index - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index - %bounded_route_id_stride0 = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index - %expert_count, %route_count, %bounded_route_id_stride = index.assume %expert_count0, %route_count0, %bounded_route_id_stride0 [range(%expert_count0, 32, 512), mul(%expert_count0, 32), range(%route_count0, 1, 32), le(%route_count0, %expert_count0), le(%route_count0, %bounded_route_id_stride0)] : index, index, index - %safe_token, %launch_token_count = index.assume %token, %bounded_token_count [lt(%token, %bounded_token_count)] : index, index - %lane = kernel.subgroup.lane.id : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %wave_size = target.subgroup.size : index - %negative_large = scalar.constant -3.4028234663852886e+38 : f32 - %largest_i32 = scalar.constant 2147483647 : i32 - %c0_offset = index.constant 0 : offset - %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 - %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> - %route_ids_view = buffer.view %route_ids_noalias[%c0_offset] : buffer -> view<[%route_id_storage_count]xi32> - %route_weights_view = buffer.view %route_weights_noalias[%c0_offset] : buffer -> view<[%route_weight_storage_count]xf32> - %initial_logits = vector.load %logits_view[%safe_token, %lane_expert_base] : view<[%launch_token_count]x[%expert_count]xf32> -> vector<[%experts_per_lane]xf32> - %remaining_final, %selected_logit = scf.for %route = [%c0 to %route_count step %c1](%remaining_logits = %initial_logits : vector<[%experts_per_lane]xf32>, %lane_selected_logit = %negative_large : f32) -> (vector<[%experts_per_lane]xf32>, f32) { - %local_value, %local_id = scf.for %slot = [%c0 to %experts_per_lane step %c1](%best_value = %negative_large : f32, %best_id = %largest_i32 : i32) -> (f32, i32) unroll { - %candidate_value = vector.extract %remaining_logits[%slot] : vector<[%experts_per_lane]xf32> -> f32 - %candidate_expert = index.add %lane_expert_base, %slot : index - %candidate_id = index.cast %candidate_expert : index to i32 - %is_greater = scalar.cmpf ogt, %candidate_value, %best_value : f32 - %is_equal = scalar.cmpf oeq, %candidate_value, %best_value : f32 - %is_lower_id = scalar.cmpi ult, %candidate_id, %best_id : i32 - %is_lower_tie = scalar.andi %is_equal, %is_lower_id : i1 - %is_better = scalar.ori %is_greater, %is_lower_tie : i1 - %next_value, %next_id = scf.if %is_better -> (f32, i32) { - scf.yield %candidate_value, %candidate_id : f32, i32 - } else { - scf.yield %best_value, %best_id : f32, i32 - } - scf.yield %next_value, %next_id : f32, i32 - } - %winner_value = kernel.subgroup.reduce %local_value : f32 - %matches_winner = scalar.cmpf oeq, %local_value, %winner_value : f32 - %winner_lane_mask = kernel.subgroup.vote.ballot %matches_winner : i1 -> i64 - %nonzero_winner_lane_mask = scalar.assume %winner_lane_mask [ne(%winner_lane_mask, 0)] : i64 - %winner_lane_i64 = scalar.cttzi %nonzero_winner_lane_mask : i64 - %winner_lane0 = index.cast %winner_lane_i64 : i64 to index - %winner_lane = index.assume %winner_lane0 [range(%winner_lane0, 0, 63)] : index - %local_expert0 = index.cast %local_id : i32 to index - %local_expert = index.assume %local_expert0 [range(%local_expert0, 0, 511)] : index - %winner_slot = index.rem %local_expert, %experts_per_lane : index - %owns_winner = index.cmp eq, %lane, %winner_lane : index - %next_remaining_logits = scf.if %owns_winner -> (vector<[%experts_per_lane]xf32>) { - %removed = vector.insert %negative_large into %remaining_logits[%winner_slot] : f32, vector<[%experts_per_lane]xf32> - scf.yield %removed : vector<[%experts_per_lane]xf32> - } else { - scf.yield %remaining_logits : vector<[%experts_per_lane]xf32> - } - %publishes_route_id = scalar.andi %valid_token, %owns_winner : i1 - scf.if %publishes_route_id { - %route_id_token_base = index.mul %safe_token, %bounded_route_id_stride : index - %route_id_index = index.add %route_id_token_base, %route : index - view.store %local_id, %route_ids_view[%route_id_index] : i32, view<[%route_id_storage_count]xi32> - } - %lane_publishes_route = index.cmp eq, %lane, %route : index - %publishes_route = scalar.andi %valid_token, %lane_publishes_route : i1 - %next_lane_selected_logit = scf.if %publishes_route -> (f32) { - scf.yield %winner_value : f32 - } else { - scf.yield %lane_selected_logit : f32 - } - scf.yield %next_remaining_logits, %next_lane_selected_logit : vector<[%experts_per_lane]xf32>, f32 - } - %selected_max = kernel.subgroup.reduce %selected_logit : f32 - %selected_delta = scalar.subf %selected_logit, %selected_max : f32 - %unnormalized_weight = scalar.expf %selected_delta : f32 - %selected_sum = kernel.subgroup.reduce %unnormalized_weight : f32 - %route_weight = scalar.divf %unnormalized_weight, %selected_sum : f32 - %lane_publishes_weight = index.cmp ult, %lane, %route_count : index - %publishes_weight = scalar.andi %valid_token, %lane_publishes_weight : i1 - scf.if %publishes_weight { - %route_weight_token_base = index.mul %safe_token, %route_count : index - %route_weight_index = index.add %route_weight_token_base, %lane : index - view.store %route_weight, %route_weights_view[%route_weight_index] : f32, view<[%route_weight_storage_count]xf32> - } - template.return -} - kernel.def target(@qwen3_moe_router_gfx11_wave64) @qwen3_moe_router_top8_f32(%token_count: index, %route_id_stride: index) { %c1 = index.constant 1 : index %c3 = index.constant 3 : index @@ -162,6 +70,54 @@ kernel.def @qwen3_moe_router_top8_wide_stride_reference() { kernel.return } +// Fixed-shape wrappers keep the correctness cases independent of dynamic +// subrange proof without changing the production kernel's lowering. +kernel.def target(@qwen3_moe_router_gfx11_wave64) @qwen3_moe_router_top8_f32_check_one() { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%c1, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%logits: buffer, %route_ids: buffer, %route_weights: buffer) { + %subgroup = kernel.subgroup.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c8 = index.constant 8 : index + %valid_token = index.cmp ult, %subgroup, %c1 : index + %safe_token = scf.select %valid_token, %subgroup, %c0 : index + template.apply<@qwen3_moe.router.top8.row>(%valid_token, %safe_token, %c1, %c8, %logits, %route_ids, %route_weights) : (i1, index, index, index, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@qwen3_moe_router_gfx11_wave64) @qwen3_moe_router_top8_f32_check_two_compact() { + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%c1, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%logits: buffer, %route_ids: buffer, %route_weights: buffer) { + %subgroup = kernel.subgroup.id : index + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c8 = index.constant 8 : index + %valid_token = index.cmp ult, %subgroup, %c2 : index + %safe_token = scf.select %valid_token, %subgroup, %c0 : index + template.apply<@qwen3_moe.router.top8.row>(%valid_token, %safe_token, %c2, %c8, %logits, %route_ids, %route_weights) : (i1, index, index, index, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@qwen3_moe_router_gfx11_wave64) @qwen3_moe_router_top8_f32_check_two_wide() { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%c1, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%logits: buffer, %route_ids: buffer, %route_weights: buffer) { + %subgroup = kernel.subgroup.id : index + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c16 = index.constant 16 : index + %valid_token = index.cmp ult, %subgroup, %c2 : index + %safe_token = scf.select %valid_token, %subgroup, %c0 : index + template.apply<@qwen3_moe.router.top8.row>(%valid_token, %safe_token, %c2, %c16, %logits, %route_ids, %route_weights) : (i1, index, index, index, buffer, buffer, buffer) + kernel.return +} + // Period-eight logits place sixteen equal maxima in different lane-local // slots. Repeated selection must remove each winner and retain the lower-ID // half of that tie set. Two rows leave two inactive waves in the four-wave @@ -179,9 +135,9 @@ check.case public @qwen3_moe_router_top8_f32_repeated_maxima_case { %expected_ids = check.generate.iota offset(7) step(8) period(8) : tensor<2x8xi32> %expected_wide_ids = check.generate.fill value(-1) : tensor<2x16xi32> %expected_weights = check.generate.fill value(0.125) : tensor<2x8xf32> - kernel.launch @qwen3_moe_router_top8_f32[%token_count, %route_id_stride](%token_count, %route_id_stride, %logits, %route_ids, %route_weights) : [index, index](index, index, tensor<2x128xf32>, tensor<2x8xi32>, tensor<2x8xf32>) + kernel.launch @qwen3_moe_router_top8_f32_check_two_compact(%logits, %route_ids, %route_weights) : (tensor<2x128xf32>, tensor<2x8xi32>, tensor<2x8xf32>) kernel.launch @qwen3_moe_router_top8_wide_stride_reference(%expected_wide_ids) : (tensor<2x16xi32>) - kernel.launch @qwen3_moe_router_top8_f32[%token_count, %wide_route_id_stride](%token_count, %wide_route_id_stride, %logits, %wide_route_ids, %wide_route_weights) : [index, index](index, index, tensor<2x128xf32>, tensor<2x16xi32>, tensor<2x8xf32>) + kernel.launch @qwen3_moe_router_top8_f32_check_two_wide(%logits, %wide_route_ids, %wide_route_weights) : (tensor<2x128xf32>, tensor<2x16xi32>, tensor<2x8xf32>) check.expect.equal actual(%route_ids) expected(%expected_ids) : tensor<2x8xi32> check.expect.close actual(%route_weights) expected(%expected_weights) atol(1.0000000000000001e-05) rtol(1.0000000000000001e-05) nan(same) : tensor<2x8xf32> check.expect.equal actual(%wide_route_ids) expected(%expected_wide_ids) : tensor<2x16xi32> diff --git a/ggml/src/ggml-hrx/tools/benchmarks/generate-model-benchmarks.py b/ggml/src/ggml-hrx/tools/benchmarks/generate-model-benchmarks.py index 5ab37b10fa33..a3c1ca06f630 100755 --- a/ggml/src/ggml-hrx/tools/benchmarks/generate-model-benchmarks.py +++ b/ggml/src/ggml-hrx/tools/benchmarks/generate-model-benchmarks.py @@ -715,11 +715,7 @@ def render_case(symbol: str, command: dict[str, Any], export: dict[str, Any]) -> elif kernel == "ggml_mul_mat_add_f32_f32_decode_wave64": case = case_mul_mat_add_decode(symbol, command, kernel) elif kernel in ( - "ggml_mul_mat_tiled_input_f32_bias_publish_f32", - "ggml_mul_mat_tiled_input_f32_residual_publish_f32", "ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", - "ggml_mul_mat_skinny_input_f32_bias_publish_f32", - "ggml_mul_mat_skinny_input_f32_residual_publish_f32", "ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", ): case = case_mul_mat_postops(symbol, command, kernel) diff --git a/ggml/src/ggml-hrx/tools/compile_qwen_kernel_corpus.py b/ggml/src/ggml-hrx/tools/compile_qwen_kernel_corpus.py index d92136fd3ca2..e7104c635cb0 100755 --- a/ggml/src/ggml-hrx/tools/compile_qwen_kernel_corpus.py +++ b/ggml/src/ggml-hrx/tools/compile_qwen_kernel_corpus.py @@ -19,7 +19,16 @@ ATTR_RE = re.compile(r"(?P[A-Za-z0-9_.]+)\s*=\s*(?P-?[0-9]+)") CASE_RE = re.compile(r"check\.case(?:\s+public)?\s+@(?P[A-Za-z0-9_]+)\s*\{") LITERAL_RE = re.compile(r"%(?P[A-Za-z0-9_]+)\s*=\s*check\.literal\s+value\((?P-?[0-9]+)\)\s*:\s*index") -FUNC_CALL_RE = re.compile(r"func\.call\s+@(?P[A-Za-z0-9_]+)\s*\(") +KERNEL_LAUNCH_RE = re.compile( + r"kernel\.launch\s+@(?P[A-Za-z0-9_]+)" + r"(?:\[(?P[^]]*)\])?\s*\(" +) +SSA_VALUE_RE = re.compile(r"%(?P[A-Za-z0-9_]+)") +TARGET_DEF_RE = re.compile(r"amdgpu\.target<(?P[^>]+)>\s+@(?P[A-Za-z0-9_]+)") +KERNEL_TARGET_RE = re.compile( + r"kernel\.def[^\n{]*target\(@(?P[A-Za-z0-9_]+)\)[^\n{]*@(?P[A-Za-z0-9_]+)\s*\(" +) +KERNEL_DEF_RE = re.compile(r"kernel\.def[^\n{]*@(?P[A-Za-z0-9_]+)\s*\(") def run(command: list[str], *, env: dict[str, str] | None = None) -> subprocess.CompletedProcess[str]: @@ -66,6 +75,28 @@ def main() -> int: linked_dir.mkdir(exist_ok=True) linked_paths: dict[str, pathlib.Path] = {} + def module_inputs(name: str) -> list[pathlib.Path]: + paths: list[pathlib.Path] = [] + seen: set[pathlib.Path] = set() + + def add(path: pathlib.Path) -> None: + if path not in seen: + seen.add(path) + paths.append(path) + + def visit(module_name: str) -> None: + recipe = modules[module_name] + for source in recipe["srcs"]: + add(args.corpus_dir / source) + for library in recipe["libraries"]: + if library.startswith(":"): + visit(library[1:]) + else: + add(args.corpus_dir / library) + + visit(name) + return paths + def link_module(name: str) -> pathlib.Path: if name in linked_paths: return linked_paths[name] @@ -81,20 +112,45 @@ def link_module(name: str) -> pathlib.Path: return output exports = {item["symbol"]: item for item in manifest["exports"]} + kernel_defs_by_source: dict[pathlib.Path, set[str]] = {} + + def kernel_defs(path: pathlib.Path) -> set[str]: + if path not in kernel_defs_by_source: + text = path.read_text(encoding="utf-8") + kernel_defs_by_source[path] = {match.group("name") for match in KERNEL_DEF_RE.finditer(text)} + return kernel_defs_by_source[path] - def source_for_root(root: str) -> pathlib.Path: - source_name = exports[root]["source"] - direct_modules = [name for name, recipe in modules.items() if source_name in recipe["srcs"]] - if direct_modules: - return link_module(direct_modules[0]) - library_modules = [name for name, recipe in modules.items() if source_name in recipe["libraries"]] - if library_modules: - return link_module(library_modules[0]) - return args.corpus_dir / source_name + def internal_inputs(root: str, inputs: list[pathlib.Path]) -> list[pathlib.Path]: + owners = {path for path in inputs if root in kernel_defs(path)} + if not owners: + return inputs + return [path for path in inputs if path in owners or not kernel_defs(path)] + + def export_inputs(root: str) -> list[pathlib.Path]: + recipe = exports[root]["compile_recipe"] + return [ + args.corpus_dir / source + for source in recipe["primary_sources"] + recipe["library_sources"] + ] + + def link_root(root: str, configs: tuple[str, ...] | list[str], inputs: list[pathlib.Path], output: pathlib.Path) -> None: + output.parent.mkdir(parents=True, exist_ok=True) + target = "amdgpu:gfx11-generic" if args.target == "gfx1100" else "amdgpu:gfx1151" + command = [ + str(args.loom_link), + "--mode=link", + f"--root=@{root}", + f"--target={target}", + "--require-resolved-config", + "--to=text", + f"--output={output}", + ] + command.extend(f"--config={config}" for config in configs) + command.extend(str(path) for path in inputs) + require_run(command) results: list[dict[str, object]] = [] compiled_keys: set[tuple[str, tuple[int, ...], tuple[str, ...]]] = set() - root_sources: dict[str, pathlib.Path] = {} planned_invocation_count = 0 resolved_invocation_count = 0 for case in manifest["plan_cases"]: @@ -115,13 +171,20 @@ def source_for_root(root: str) -> pathlib.Path: if not plan_rows: raise RuntimeError(f"BUILD recipe {case['name']} produced no planner rows") source_text = source.read_text(encoding="utf-8") + target_selectors = { + match.group("name"): match.group("selector") for match in TARGET_DEF_RE.finditer(source_text) + } + kernel_selectors = { + match.group("name"): target_selectors.get(match.group("target"), "") + for match in KERNEL_TARGET_RE.finditer(source_text) + } benchmark_defs = { match.group("name"): (match.group("case"), {item.group("name"): int(item.group("value")) for item in ATTR_RE.finditer(match.group("attrs") or "")}) for match in BENCHMARK_RE.finditer(source.read_text(encoding="utf-8")) } case_literals: dict[str, dict[str, int]] = {} - case_calls: dict[str, list[str]] = {} + case_invocations: dict[str, list[tuple[str, tuple[str, ...]]]] = {} for match in CASE_RE.finditer(source_text): depth = 1 cursor = match.end() @@ -133,43 +196,63 @@ def source_for_root(root: str) -> pathlib.Path: case_literals[match.group("name")] = { item.group("name"): int(item.group("value")) for item in LITERAL_RE.finditer(body) } - case_calls[match.group("name")] = [item.group("name") for item in FUNC_CALL_RE.finditer(body)] + case_invocations[match.group("name")] = [ + ( + item.group("name"), + tuple(value.group("name") for value in SSA_VALUE_RE.finditer(item.group("workload") or "")), + ) + for item in KERNEL_LAUNCH_RE.finditer(body) + ] configs = [item.removeprefix("--config=") for item in case["args"] if item.startswith("--config=")] selects_benchmark = any(item.startswith("--benchmark=") for item in case["args"]) owned_sources = set(modules[case["link_module"]]["srcs"]) if "link_module" in case else {case["source"]} + case_inputs = module_inputs(case["link_module"]) if "link_module" in case else [args.corpus_dir / case["source"]] for row in plan_rows: benchmark_case, attrs = benchmark_defs.get(row["benchmark"], (row["case"], {})) concrete_values = dict(case_literals.get(benchmark_case, {})) concrete_values.update(attrs) - expected_count = int(row.get("actual_invocation_count", 1)) + expected_count = int(row.get("kernel_launch_count", row.get("actual_invocation_count", 1))) planned_invocation_count += expected_count - if row.get("actual_entry"): - roots = [row["actual_entry"]] - else: - roots = [root for root in case_calls.get(benchmark_case, []) if root in exports] - if len(roots) != expected_count: + invocations = case_invocations.get(benchmark_case, []) + actual_entry = row.get("kernel_entry", row.get("actual_entry")) + if actual_entry: + invocations = [item for item in invocations if item[0] == actual_entry] + if not invocations and actual_entry in exports: + invocations = [ + ( + actual_entry, + tuple(item["name"] for item in exports[actual_entry]["workload_parameters"]), + ) + ] + if len(invocations) != expected_count: raise RuntimeError( f"BUILD recipe {case['name']} planner reports {expected_count} invocations for " - f"{benchmark_case}, but source resolves {len(roots)} exported calls: {roots}" + f"{benchmark_case}, but source resolves {len(invocations)} kernel launches: " + f"{[item[0] for item in invocations]}" ) - resolved_invocation_count += len(roots) - for root in roots: - if root not in exports: - raise RuntimeError(f"BUILD recipe {case['name']} selected unmanifested kernel {root}") - if not selects_benchmark and exports[root]["source"] not in owned_sources: + resolved_invocation_count += len(invocations) + for root, parameter_names in invocations: + if not selects_benchmark and root in exports and exports[root]["source"] not in owned_sources: continue - root_sources.setdefault(root, source) - parameter_names = [item["name"] for item in exports[root]["workload_parameters"]] missing = [name for name in parameter_names if name not in concrete_values] if missing: raise RuntimeError(f"{case['name']} recipe omits concrete {missing} for {root}") workload = tuple(concrete_values[name] for name in parameter_names) - key = (root, workload, tuple(configs)) + key = (root, workload, tuple(sorted(configs))) if key in compiled_keys: continue compiled_keys.add(key) compile_dir = case_dir / f"{root}-{len(compiled_keys):03d}" - command = [str(args.compiler), "--target", args.target, "--source", str(source), + selector = kernel_selectors.get(root, "") + if selector and selector not in ("gfx11-generic", args.target): + results.append({"case": case["name"], "root": root, "workload": workload, + "configs": configs, "artifact_dir": str(compile_dir), + "status": "skipped-incompatible-target", "target_selector": selector}) + continue + linked_root = compile_dir / "linked.loom" + inputs = export_inputs(root) if root in exports else internal_inputs(root, case_inputs) + link_root(root, configs, inputs, linked_root) + command = [str(args.compiler), "--target", args.target, "--source", str(linked_root), "--root", root, "--output", str(compile_dir)] for config in configs: command.extend(("--config", config)) @@ -189,44 +272,65 @@ def source_for_root(root: str) -> pathlib.Path: schedule = json.loads(schedule_path.read_text(encoding="utf-8")) schedule_dir = args.output_dir / "schedules" / schedule_path.parent.name schedule_dir.mkdir(parents=True, exist_ok=True) - for invocation in schedule.get("invocations", []): - for dispatch in invocation.get("dispatches", []): - specialization = dispatch["kernel"] - if specialization.get("execution") != "native": - raise RuntimeError(f"schedule {schedule_path} contains non-native dispatch {specialization['variant']}") - root = specialization["variant"] - if root not in exports: - raise RuntimeError(f"schedule {schedule_path} selects unmanifested kernel {root}") - schedule_requirement_count += 1 - parameters = specialization.get("parameters", {}) - parameter_names = [item["name"] for item in exports[root]["workload_parameters"]] - missing = [name for name in parameter_names if name not in parameters] - if missing: - raise RuntimeError(f"schedule {schedule_path} omits concrete {missing} for {root}") - workload = tuple(int(parameters[name]) for name in parameter_names) - configs = tuple(f"{name}={value}" for name, value in - sorted(specialization.get("compile_parameters", {}).items())) - key = (root, workload, configs) - if key in compiled_keys: + specializations: list[dict[str, object]] = [] + if schedule.get("schema") == "ggml-hrx-command-program-v1": + for command in schedule.get("commands", []): + if command.get("kind") != "Kernel": continue - schedule_unique_requirement_count += 1 - compiled_keys.add(key) - source = root_sources[root] if root in root_sources else source_for_root(root) - compile_dir = schedule_dir / f"{root}-{schedule_unique_requirement_count:03d}" - command = [str(args.compiler), "--target", args.target, "--source", str(source), - "--root", root, "--output", str(compile_dir)] - for config in configs: - command.extend(("--config", config)) - for value in workload: - command.extend(("--workload", str(value))) - compile_result = run(command, env=os.environ.copy()) - (schedule_dir / f"{root}-{schedule_unique_requirement_count:03d}.stdout.txt").write_text( - compile_result.stdout, encoding="utf-8") - (schedule_dir / f"{root}-{schedule_unique_requirement_count:03d}.stderr.txt").write_text( - compile_result.stderr, encoding="utf-8") - results.append({"case": f"schedule:{schedule_path}", "root": root, "workload": workload, - "configs": configs, "artifact_dir": str(compile_dir), - "status": "ok" if compile_result.returncode == 0 else "failed"}) + qualified_root = str(command.get("kernel", "")) + if not qualified_root.startswith("qwen3_moe:"): + continue + specializations.append({ + "variant": qualified_root.removeprefix("qwen3_moe:"), + "parameters": command.get("integer_parameters", {}), + "compile_parameters": command.get("compile_parameters", {}), + }) + else: + for invocation in schedule.get("invocations", []): + for dispatch in invocation.get("dispatches", []): + specialization = dispatch["kernel"] + if specialization.get("execution") != "native": + raise RuntimeError( + f"schedule {schedule_path} contains non-native dispatch {specialization['variant']}" + ) + specializations.append(specialization) + if not specializations: + raise RuntimeError(f"schedule {schedule_path} contains no Qwen MoE kernel requirements") + for specialization in specializations: + root = str(specialization["variant"]) + if root not in exports: + raise RuntimeError(f"schedule {schedule_path} selects unmanifested Qwen MoE kernel {root}") + schedule_requirement_count += 1 + parameters = specialization.get("parameters", {}) + parameter_names = [item["name"] for item in exports[root]["workload_parameters"]] + missing = [name for name in parameter_names if name not in parameters] + if missing: + raise RuntimeError(f"schedule {schedule_path} omits concrete {missing} for {root}") + workload = tuple(int(parameters[name]) for name in parameter_names) + compile_parameters = specialization.get("compile_parameters", {}) + configs = tuple(f"{name}={value}" for name, value in sorted(compile_parameters.items())) + key = (root, workload, configs) + if key in compiled_keys: + continue + schedule_unique_requirement_count += 1 + compiled_keys.add(key) + compile_dir = schedule_dir / f"{root}-{schedule_unique_requirement_count:03d}" + linked_root = compile_dir / "linked.loom" + link_root(root, configs, export_inputs(root), linked_root) + command = [str(args.compiler), "--target", args.target, "--source", str(linked_root), + "--root", root, "--output", str(compile_dir)] + for config in configs: + command.extend(("--config", config)) + for value in workload: + command.extend(("--workload", str(value))) + compile_result = run(command, env=os.environ.copy()) + (schedule_dir / f"{root}-{schedule_unique_requirement_count:03d}.stdout.txt").write_text( + compile_result.stdout, encoding="utf-8") + (schedule_dir / f"{root}-{schedule_unique_requirement_count:03d}.stderr.txt").write_text( + compile_result.stderr, encoding="utf-8") + results.append({"case": f"schedule:{schedule_path}", "root": root, "workload": workload, + "configs": configs, "artifact_dir": str(compile_dir), + "status": "ok" if compile_result.returncode == 0 else "failed"}) summary = { "schema": "ggml-hrx-qwen-compile-report-v1", @@ -241,7 +345,10 @@ def source_for_root(root: str) -> pathlib.Path: "schedule_requirement_count": schedule_requirement_count, "schedule_unique_requirement_count": schedule_unique_requirement_count, "compile_count": len(results), - "failed_count": sum(item["status"] != "ok" for item in results), + "skipped_incompatible_target_count": sum( + item["status"] == "skipped-incompatible-target" for item in results + ), + "failed_count": sum(item["status"] == "failed" for item in results), "results": results, } (args.output_dir / "summary.json").write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8") diff --git a/tests/hrx-backend-test.cpp b/tests/hrx-backend-test.cpp index 529ed4fdf0ae..db113e6e0970 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -9,6 +9,7 @@ #include "dispatch_registration/dispatch-registry.h" #include "ggml-alloc.h" #include "ggml-backend.h" +#include "ggml-cpu.h" #include "ggml-hrx.h" #include "ggml-impl.h" #include "ggml.h" @@ -70,10 +71,16 @@ static bool contains_value_id(const std::vector & ids, ggml: } static bool command_program_verifies(const ggml::hrx::CommandProgram & program) { - return ggml::hrx::verify_command_program(program, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151").valid(); + const ggml::hrx::VerificationResult result = + ggml::hrx::verify_command_program(program, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + if (!result.valid()) { + for (const std::string & error : result.status.errors()) { + std::fprintf(stderr, "%s\n", error.c_str()); + } + } + return result.valid(); } - static bool async_jit_expected_from_environment() { const char * value = std::getenv("GGML_HRX_ASYNC_JIT"); return value == nullptr || @@ -183,6 +190,36 @@ static std::string kernel_name_for_id(uint64_t kernel_id) { return ggml::hrx::kernel_definition_name(*resolved.definition); } +template static void require_scheduled_command_program(ggml_cgraph * graph, Check && check) { + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + + const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( + imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(command_program_verifies(commands)); + check(imported.graph, scheduler.plan(), commands); +} + +static std::vector scheduled_kernel_names(ggml_cgraph * graph) { + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + + std::vector result; + for (const ggml::hrx::Dispatch & dispatch : scheduler.plan().dispatches) { + result.push_back(kernel_name_for_id(dispatch.kernel.kernel_id)); + } + return result; +} + static void require_compile_parameter(const ggml::hrx::Dispatch & dispatch, const char * name, const std::string & value) { @@ -690,13 +727,12 @@ static void run_dispatch_registry_checks() { "common.mul_mat_add.skinny_f32_f32_decode", ggml::hrx::DispatchMatchKind::Fused)); REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.mul_mat_unary.tiled_f32_f32", - ggml::hrx::DispatchMatchKind::Fused)); + "common.mul_mat_unary.tiled_f32_f32", ggml::hrx::DispatchMatchKind::Fused)); REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.mul_mat_unary.skinny_f32_f32", - ggml::hrx::DispatchMatchKind::Fused)); + "common.mul_mat_unary.skinny_f32_f32", ggml::hrx::DispatchMatchKind::Fused)); REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.mul_mat_postops.tiled_f32_f32", ggml::hrx::DispatchMatchKind::Fused)); + "common.mul_mat_postops.tiled_f32_f32", + ggml::hrx::DispatchMatchKind::Fused)); REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), "common.mul_mat_postops.skinny_f32_f32", ggml::hrx::DispatchMatchKind::Fused)); @@ -717,9 +753,12 @@ static void run_dispatch_registry_checks() { "qwen.attention_postprocess_f32_f16")); REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_GET_ROWS), "common.get_rows.f32")); REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_GET_ROWS), "common.get_rows.f32_next")); + REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_GET_ROWS), "common.get_rows_scale.f32")); REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_GET_ROWS), "common.gather_add_f32")); REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_ROPE), "common.rope_set_rows.f32", ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_ROPE), "common.rope_concat.f32", + ggml::hrx::DispatchMatchKind::Fused)); REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_ROPE), "common.rope.f32", ggml::hrx::DispatchMatchKind::SingleOp)); REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_SET_ROWS), "common.set_rows")); @@ -789,13 +828,12 @@ static void run_dispatch_registry_checks() { "common.mul_mat_add.skinny_f32_f32_decode", ggml::hrx::DispatchMatchKind::Fused)); REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.mul_mat_unary.tiled_f32_f32", - ggml::hrx::DispatchMatchKind::Fused)); + "common.mul_mat_unary.tiled_f32_f32", ggml::hrx::DispatchMatchKind::Fused)); REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.mul_mat_unary.skinny_f32_f32", - ggml::hrx::DispatchMatchKind::Fused)); + "common.mul_mat_unary.skinny_f32_f32", ggml::hrx::DispatchMatchKind::Fused)); REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.mul_mat_postops.tiled_f32_f32", ggml::hrx::DispatchMatchKind::Fused)); + "common.mul_mat_postops.tiled_f32_f32", + ggml::hrx::DispatchMatchKind::Fused)); REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), "common.mul_mat_postops.skinny_f32_f32", ggml::hrx::DispatchMatchKind::Fused)); @@ -1231,8 +1269,7 @@ static void run_graph_import_checks() { transient_flow.commands[1].bindings[0].origin = ggml::hrx::CommandBindingOrigin::Transient; transient_flow.commands[1].bindings[0].value = transient_value; transient_flow.commands[1].bindings[0].length = command.bindings[2].length; - verification = - ggml::hrx::verify_command_program(transient_flow, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + verification = ggml::hrx::verify_command_program(transient_flow, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(verification.valid()); ggml::hrx::CommandProgram read_before_write = copy_command_program_shape(transient_flow); @@ -1242,8 +1279,7 @@ static void run_graph_import_checks() { read_before_write.commands[1].ordinal = 1; read_before_write.commands[1].dependencies.clear(); read_before_write.commands[1].dependencies.push_back(0); - verification = - ggml::hrx::verify_command_program(read_before_write, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + verification = ggml::hrx::verify_command_program(read_before_write, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(!verification.valid()); REQUIRE(status_contains(verification.status, "reads transient value 100000 before write by command 1")); @@ -1487,6 +1523,172 @@ static void run_scale_f32_dispatch_checks() { ggml_free(ctx); } +static void require_scale_add_dispatch(ggml_context * ctx, + ggml_tensor * input, + ggml_tensor * residual, + int64_t expected_element_count, + int64_t expected_input_stride1, + int64_t expected_residual_stride1, + ggml::hrx::BinaryKind binary_kind = ggml::hrx::BinaryKind::Add, + bool scaled_lhs = true) { + ggml_tensor * scaled = ggml_scale_bias(ctx, input, 0.177800179f, 0.25f); + ggml_tensor * output = nullptr; + switch (binary_kind) { + case ggml::hrx::BinaryKind::Add: + output = ggml_add(ctx, scaled, residual); + break; + case ggml::hrx::BinaryKind::Sub: + output = scaled_lhs ? ggml_sub(ctx, scaled, residual) : ggml_sub(ctx, residual, scaled); + break; + case ggml::hrx::BinaryKind::Mul: + output = ggml_mul(ctx, scaled, residual); + break; + case ggml::hrx::BinaryKind::Div: + output = scaled_lhs ? ggml_div(ctx, scaled, residual) : ggml_div(ctx, residual, scaled); + break; + default: + REQUIRE(false); + } + REQUIRE(scaled != nullptr); + REQUIRE(output != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + + require_scheduled_command_program(graph, [&](const ggml::hrx::Graph &, const ggml::hrx::CommandPlan & plan, + const ggml::hrx::CommandProgram & commands) { + REQUIRE(plan.dispatches.size() == 1); + const ggml::hrx::Dispatch & dispatch = plan.dispatches.front(); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_scale_add_f32"); + REQUIRE(dispatch.kernel.integer_parameters.at("element_count") == expected_element_count); + require_compile_parameter(dispatch, "ggml.scale_add_f32.scale", expected_config_value(0.177800179f)); + require_compile_parameter(dispatch, "ggml.scale_add_f32.bias", expected_config_value(0.25f)); + require_compile_parameter(dispatch, "ggml.scale_add_f32.binary_op", + std::to_string(ggml::hrx::binary_kind_config_value(binary_kind))); + require_compile_parameter(dispatch, "ggml.scale_add_f32.scaled_lhs", scaled_lhs ? "1" : "0"); + require_compile_parameter(dispatch, "ggml.scale_add_f32.input_stride1", std::to_string(expected_input_stride1)); + require_compile_parameter(dispatch, "ggml.scale_add_f32.residual_stride1", + std::to_string(expected_residual_stride1)); + REQUIRE(dispatch.bindings.size() == 3); + REQUIRE(commands.commands.size() == 1); + REQUIRE(commands.commands.front().bindings[0].name == "input"); + REQUIRE(commands.commands.front().bindings[1].name == "residual"); + REQUIRE(commands.commands.front().bindings[2].name == "output"); + }); +} + +static void run_scale_add_f32_dispatch_checks() { + ggml_init_params params = {}; + params.mem_size = 2 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * packed_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2560, 64); + ggml_tensor * packed_residual = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2560, 64); + REQUIRE(packed_input != nullptr); + REQUIRE(packed_residual != nullptr); + require_scale_add_dispatch(ctx, packed_input, packed_residual, 163840, 2560, 2560); + require_scale_add_dispatch(ctx, packed_input, packed_residual, 163840, 2560, 2560, ggml::hrx::BinaryKind::Sub, + false); + require_scale_add_dispatch(ctx, packed_input, packed_residual, 163840, 2560, 2560, ggml::hrx::BinaryKind::Mul, + true); + require_scale_add_dispatch(ctx, packed_input, packed_residual, 163840, 2560, 2560, ggml::hrx::BinaryKind::Div, + true); + + ggml_tensor * input_storage = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 320, 3); + ggml_tensor * residual_storage = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 384, 3); + ggml_tensor * strided_input = ggml_view_2d(ctx, input_storage, 257, 3, input_storage->nb[1], 0); + ggml_tensor * strided_residual = ggml_view_2d(ctx, residual_storage, 257, 3, residual_storage->nb[1], 0); + REQUIRE(strided_input != nullptr); + REQUIRE(strided_residual != nullptr); + require_scale_add_dispatch(ctx, strided_input, strided_residual, 771, 320, 384); + + ggml_tensor * scaled = ggml_scale(ctx, packed_input, 0.5f); + ggml_tensor * alias_output = ggml_add_inplace(ctx, scaled, packed_residual); + REQUIRE(scaled != nullptr); + REQUIRE(alias_output != nullptr); + REQUIRE(!can_schedule_tensor(ctx, alias_output)); + + ggml_tensor * alias_scaled = ggml_scale_inplace(ctx, packed_input, 0.5f); + ggml_tensor * output_from_alias = ggml_add(ctx, alias_scaled, packed_residual); + REQUIRE(alias_scaled != nullptr); + REQUIRE(output_from_alias != nullptr); + ggml_cgraph * alias_graph = ggml_new_graph(ctx); + REQUIRE(alias_graph != nullptr); + ggml_build_forward_expand(alias_graph, output_from_alias); + ggml::hrx::GraphImportResult alias_imported = ggml::hrx::import_ggml_graph(*alias_graph); + REQUIRE(alias_imported.valid()); + ggml::hrx::DispatchScheduler alias_scheduler; + REQUIRE(alias_scheduler.schedule_graph(alias_imported.graph, test_dispatch_target())); + REQUIRE(alias_scheduler.plan().dispatches.size() == 2); + REQUIRE(kernel_name_for_id(alias_scheduler.plan().dispatches[0].kernel.kernel_id) == + "loom_libs:ggml_scale_bias_f32"); + REQUIRE(kernel_name_for_id(alias_scheduler.plan().dispatches[1].kernel.kernel_id) == "loom_libs:ggml_binary_f32"); + + ggml_free(ctx); +} + +static ggml_tensor * make_test_binary(ggml_context * ctx, + ggml_tensor * lhs, + ggml_tensor * rhs, + ggml::hrx::BinaryKind kind) { + switch (kind) { + case ggml::hrx::BinaryKind::Add: + return ggml_add(ctx, lhs, rhs); + case ggml::hrx::BinaryKind::Sub: + return ggml_sub(ctx, lhs, rhs); + case ggml::hrx::BinaryKind::Mul: + return ggml_mul(ctx, lhs, rhs); + case ggml::hrx::BinaryKind::Div: + return ggml_div(ctx, lhs, rhs); + default: + REQUIRE(false); + return nullptr; + } +} + +static void run_rmsnorm_two_binary_dispatch_checks() { + const struct { + ggml::hrx::BinaryKind normalized_op; + ggml::hrx::BinaryKind output_op; + bool binary_lhs; + } cases[] = { + { ggml::hrx::BinaryKind::Add, ggml::hrx::BinaryKind::Mul, true }, + { ggml::hrx::BinaryKind::Sub, ggml::hrx::BinaryKind::Sub, true }, + { ggml::hrx::BinaryKind::Mul, ggml::hrx::BinaryKind::Div, true }, + { ggml::hrx::BinaryKind::Div, ggml::hrx::BinaryKind::Sub, false }, + }; + + for (const auto & test : cases) { + ggml_init_params params = {}; + params.mem_size = 512 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 128, 4); + ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 128); + ggml_tensor * residual = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 128, 4); + ggml_tensor * rms = ggml_rms_norm(ctx, input, 1.0e-6f); + ggml_tensor * scaled = make_test_binary(ctx, rms, weight, test.normalized_op); + ggml_tensor * output = test.binary_lhs ? make_test_binary(ctx, scaled, residual, test.output_op) : + make_test_binary(ctx, residual, scaled, test.output_op); + REQUIRE(output != nullptr); + + const ggml::hrx::Dispatch dispatch = schedule_single_dispatch_for_tensor(ctx, output); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rmsnorm_mul_add_f32"); + REQUIRE(dispatch.bindings.size() == 4); + require_compile_parameter(dispatch, "ggml.rmsnorm_mul_add_f32.normalized_op", + std::to_string(ggml::hrx::binary_kind_config_value(test.normalized_op))); + require_compile_parameter(dispatch, "ggml.rmsnorm_mul_add_f32.output_op", + std::to_string(ggml::hrx::binary_kind_config_value(test.output_op))); + require_compile_parameter(dispatch, "ggml.rmsnorm_mul_add_f32.binary_lhs", test.binary_lhs ? "1" : "0"); + ggml_free(ctx); + } +} + static void run_cont_f32_dispatch_checks() { ggml_init_params params = {}; params.mem_size = 512 * 1024; @@ -1947,7 +2149,7 @@ static void run_graph_index_checks() { ggml_free(ctx); - params.mem_size = 256 * 1024; + params.mem_size = 2 * 1024 * 1024; params.no_alloc = true; ctx = ggml_init(params); REQUIRE(ctx != nullptr); @@ -2027,12 +2229,80 @@ static void run_graph_index_checks() { REQUIRE(command_program_verifies(commands)); } + { + constexpr int64_t hidden_size = 256; + constexpr int64_t token_count = 64; + constexpr int64_t input_stride = 288; + const size_t storage_elements = static_cast(token_count - 1) * input_stride + hidden_size; + ggml_tensor * storage = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, storage_elements); + ggml_tensor * view = ggml_view_2d(ctx, storage, hidden_size, token_count, input_stride * sizeof(float), 0); + ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size); + ggml_tensor * norm = ggml_rms_norm(ctx, view, 0.00001f); + ggml_tensor * scaled = ggml_mul(ctx, norm, weight); + REQUIRE(storage != nullptr); + REQUIRE(view != nullptr); + REQUIRE(weight != nullptr); + REQUIRE(norm != nullptr); + REQUIRE(scaled != nullptr); + + ggml_cgraph * strided_graph = ggml_new_graph(ctx); + REQUIRE(strided_graph != nullptr); + ggml_build_forward_expand(strided_graph, scaled); + ggml::hrx::GraphImportResult strided_imported = ggml::hrx::import_ggml_graph(*strided_graph); + REQUIRE(strided_imported.valid()); + REQUIRE(scheduler.schedule_graph(strided_imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 1); + const ggml::hrx::Dispatch & strided_dispatch = scheduler.plan().dispatches[0]; + REQUIRE(kernel_name_for_id(strided_dispatch.kernel.kernel_id) == "loom_libs:ggml_rmsnorm_binary_strided_f32"); + require_compile_parameter(strided_dispatch, "ggml.rmsnorm_binary_f32.hidden_size", "256"); + require_compile_parameter(strided_dispatch, "ggml.rmsnorm_binary_f32.input_stride", "288"); + require_compile_parameter(strided_dispatch, "ggml.rmsnorm_binary_f32.op", "2"); + REQUIRE(strided_dispatch.bindings.size() == 3); + REQUIRE(strided_dispatch.bindings[0].length == + static_cast(token_count - 1) * view->nb[1] + hidden_size * sizeof(float)); + + ggml_tensor * full_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, token_count); + ggml_tensor * rejected_norm = ggml_rms_norm(ctx, view, 0.00001f); + ggml_tensor * rejected_scaled = ggml_mul(ctx, rejected_norm, full_weight); + REQUIRE(full_weight != nullptr); + REQUIRE(rejected_scaled != nullptr); + ggml_cgraph * rejected_graph = ggml_new_graph(ctx); + REQUIRE(rejected_graph != nullptr); + ggml_build_forward_expand(rejected_graph, rejected_scaled); + ggml::hrx::GraphImportResult rejected_imported = ggml::hrx::import_ggml_graph(*rejected_graph); + REQUIRE(rejected_imported.valid()); + REQUIRE(scheduler.schedule_graph(rejected_imported.graph, test_dispatch_target())); + for (const ggml::hrx::Dispatch & rejected_dispatch : scheduler.plan().dispatches) { + REQUIRE(kernel_name_for_id(rejected_dispatch.kernel.kernel_id) != + "loom_libs:ggml_rmsnorm_binary_strided_f32"); + } + + ggml_tensor * second_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size); + ggml_tensor * shared_norm = ggml_rms_norm(ctx, view, 0.00001f); + ggml_tensor * scaled_a = ggml_mul(ctx, shared_norm, weight); + ggml_tensor * scaled_b = ggml_add(ctx, shared_norm, second_weight); + ggml_tensor * joined = ggml_add(ctx, scaled_a, scaled_b); + REQUIRE(second_weight != nullptr); + REQUIRE(joined != nullptr); + ggml_cgraph * fanout_graph = ggml_new_graph(ctx); + REQUIRE(fanout_graph != nullptr); + ggml_build_forward_expand(fanout_graph, joined); + ggml::hrx::GraphImportResult fanout_imported = ggml::hrx::import_ggml_graph(*fanout_graph); + REQUIRE(fanout_imported.valid()); + REQUIRE(scheduler.schedule_graph(fanout_imported.graph, test_dispatch_target())); + for (const ggml::hrx::Dispatch & fanout_dispatch : scheduler.plan().dispatches) { + REQUIRE(kernel_name_for_id(fanout_dispatch.kernel.kernel_id) != + "loom_libs:ggml_rmsnorm_binary_strided_f32"); + } + } + ggml_free(ctx); } static void run_rope_set_rows_dispatch_checks() { ggml_init_params params = {}; - params.mem_size = 1024 * 1024; + params.mem_size = 2 * 1024 * 1024; params.no_alloc = true; ggml_context * ctx = ggml_init(params); REQUIRE(ctx != nullptr); @@ -2043,6 +2313,65 @@ static void run_rope_set_rows_dispatch_checks() { constexpr int64_t cache_rows = 8; constexpr int64_t hidden_size = head_size * head_count; + const auto check_rope_concat = [&](int64_t concat_head_count, int64_t concat_token_count, int64_t n_dims) { + constexpr int64_t prefix_size = 64; + constexpr int64_t rope_size = 32; + constexpr int64_t source_size = prefix_size + rope_size; + const size_t source_elements = static_cast(source_size) * concat_head_count * concat_token_count; + ggml_tensor * storage = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, source_elements); + ggml_tensor * prefix = + ggml_view_3d(ctx, storage, prefix_size, concat_head_count, concat_token_count, source_size * sizeof(float), + source_size * concat_head_count * sizeof(float), 0); + ggml_tensor * rope_input = + ggml_view_3d(ctx, storage, rope_size, concat_head_count, concat_token_count, source_size * sizeof(float), + source_size * concat_head_count * sizeof(float), prefix_size * sizeof(float)); + ggml_tensor * positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, concat_token_count); + ggml_tensor * freqs = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_dims / 2); + ggml_tensor * rope = ggml_rope_ext(ctx, rope_input, positions, freqs, n_dims, GGML_ROPE_TYPE_NEOX, 0, 10000.0f, + 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + ggml_tensor * output = ggml_concat(ctx, prefix, rope, 0); + REQUIRE(storage != nullptr); + REQUIRE(prefix != nullptr); + REQUIRE(rope_input != nullptr); + REQUIRE(positions != nullptr); + REQUIRE(freqs != nullptr); + REQUIRE(rope != nullptr); + 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; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + 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_rope_concat_f32"); + REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == concat_token_count); + require_compile_parameter(dispatch, "ggml.rope_concat_f32.prefix_size", std::to_string(prefix_size)); + require_compile_parameter(dispatch, "ggml.rope_concat_f32.rope_size", std::to_string(rope_size)); + require_compile_parameter(dispatch, "ggml.rope_concat_f32.n_dims", std::to_string(n_dims)); + require_compile_parameter(dispatch, "ggml.rope_concat_f32.head_count", std::to_string(concat_head_count)); + require_compile_parameter(dispatch, "ggml.rope_concat_f32.token_capacity", std::to_string(concat_token_count)); + require_compile_parameter(dispatch, "ggml.rope_concat_f32.input_stride1", std::to_string(source_size)); + require_compile_parameter(dispatch, "ggml.rope_concat_f32.input_stride2", + std::to_string(source_size * concat_head_count)); + require_compile_parameter(dispatch, "ggml.rope_concat_f32.mode", "2"); + REQUIRE(dispatch.bindings.size() == 5); + REQUIRE(dispatch.bindings[1].length == source_elements * sizeof(float)); + + const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( + imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(command_program_verifies(commands)); + }; + + check_rope_concat(40, 64, 32); + check_rope_concat(3, 2, 16); + { ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, head_size, head_count, token_count); ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); @@ -2246,6 +2575,135 @@ static void run_rope_set_rows_dispatch_checks() { REQUIRE(command_program_verifies(commands)); } + { + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, head_size, head_count, token_count); + ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); + ggml_tensor * raw_scales = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, 1, token_count); + ggml_tensor * scales = ggml_scale(ctx, raw_scales, 0.5f); + ggml_tensor * rope = ggml_rope(ctx, input, pos, head_size, GGML_ROPE_TYPE_NORMAL); + ggml_tensor * output = ggml_mul(ctx, rope, scales); + REQUIRE(input != nullptr); + REQUIRE(pos != nullptr); + REQUIRE(raw_scales != nullptr); + REQUIRE(scales != nullptr); + REQUIRE(rope != nullptr); + 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()); + const size_t rope_index = producer_index_for_tensor(imported.graph, rope); + const size_t scale_index = producer_index_for_tensor(imported.graph, scales); + std::vector covered_nodes(imported.graph.nodes().size(), false); + ggml::hrx::CommandPlan plan; + ggml::hrx::DispatchMatch match; + + REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, rope_index, match)); + REQUIRE(match.dispatches.size() == 1); + REQUIRE(kernel_name_for_id(match.dispatches.front().kernel.kernel_id) == "loom_libs:ggml_rope_f32"); + + covered_nodes[scale_index] = true; + match = {}; + REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, rope_index, match)); + REQUIRE(match.dispatches.size() == 1); + REQUIRE(kernel_name_for_id(match.dispatches.front().kernel.kernel_id) == "loom_libs:ggml_rope_token_scale_f32"); + } + + { + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, head_size, head_count, token_count); + ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); + ggml_tensor * rope = ggml_rope(ctx, input, pos, head_size, GGML_ROPE_TYPE_NEOX); + ggml_tensor * output = ggml_scale(ctx, rope, 0.125f); + REQUIRE(input != nullptr); + REQUIRE(pos != nullptr); + REQUIRE(rope != nullptr); + 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; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + 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_rope_scale_f32"); + require_compile_parameter(dispatch, "ggml.rope_scale_f32.scale", "0.125"); + REQUIRE(dispatch.bindings.size() == 5); + + const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( + imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(command_program_verifies(commands)); + } + + { + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, head_size, head_count, token_count); + ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); + ggml_tensor * scales = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, 1, token_count); + ggml_tensor * rope = ggml_rope(ctx, input, pos, head_size, GGML_ROPE_TYPE_NORMAL); + ggml_tensor * output = ggml_mul(ctx, rope, scales); + REQUIRE(input != nullptr); + REQUIRE(pos != nullptr); + REQUIRE(scales != nullptr); + REQUIRE(rope != nullptr); + 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; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + 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_rope_token_scale_f32"); + REQUIRE(dispatch.bindings.size() == 6); + + const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( + imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(command_program_verifies(commands)); + } + + { + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, head_size, head_count, token_count); + ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); + ggml_tensor * rope = ggml_rope(ctx, input, pos, head_size, GGML_ROPE_TYPE_NEOX); + ggml_tensor * output = ggml_scale(ctx, rope, 0.125f); + ggml_tensor * side = ggml_scale(ctx, rope, 0.25f); + REQUIRE(input != nullptr); + REQUIRE(pos != nullptr); + REQUIRE(rope != nullptr); + REQUIRE(output != nullptr); + REQUIRE(side != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + ggml_build_forward_expand(graph, side); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 3); + REQUIRE(kernel_name_for_id(scheduler.plan().dispatches.front().kernel.kernel_id) == "loom_libs:ggml_rope_f32"); + } + { ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, hidden_size, cache_rows); ggml_tensor * rows = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, token_count); @@ -2652,8 +3110,7 @@ static void schedule_single_matmul_command(ggml_context * ctx, const bool q8_capable_kernel = expected_kernel == "loom_libs:ggml_mul_mat_f32_f32_wmma" || expected_kernel == "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32"; const bool q8_input = !vector_kernel && expected_token_count <= 5 && expected_output_size % 64 == 0 && - (weight_type == GGML_TYPE_Q4_K || weight_type == GGML_TYPE_Q6_K) && - q8_capable_kernel; + (weight_type == GGML_TYPE_Q4_K || weight_type == GGML_TYPE_Q6_K) && q8_capable_kernel; const size_t command_count = q8_input ? 2 : 1; REQUIRE(scheduler.plan().dispatches.size() == command_count); if (q8_input) { @@ -2743,7 +3200,6 @@ static void schedule_single_matmul_command(ggml_context * ctx, } } - static void schedule_fused_matmul_unary_command(ggml_context * ctx, ggml_tensor * output, const char * expected_kernel_name, @@ -2827,8 +3283,8 @@ static void schedule_fused_matmul_swiglu_command( REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); REQUIRE(scheduler.plan().valid()); const bool low_token = expected_token_count <= 5 && expected_input_size % 256 == 0; - const bool q8_input = low_token && expected_gate_type == GGML_TYPE_Q4_K && - expected_up_type == GGML_TYPE_Q4_K && expected_output_size % 64 == 0; + const bool q8_input = low_token && expected_gate_type == GGML_TYPE_Q4_K && expected_up_type == GGML_TYPE_Q4_K && + expected_output_size % 64 == 0; const size_t command_count = q8_input ? 2 : 1; REQUIRE(scheduler.plan().dispatches.size() == command_count); if (q8_input) { @@ -2999,8 +3455,8 @@ static void schedule_packed_matmul_glu_command(ggml_context * ctx, require_compile_parameter(dispatch, "ggml.mul_mat_swiglu.op", std::to_string(ggml::hrx::binary_kind_config_value(expected_binary_op))); - const size_t half_offset = ggml_row_size(expected_weight_type, expected_input_size) * - static_cast(expected_output_size); + const size_t half_offset = + ggml_row_size(expected_weight_type, expected_input_size) * static_cast(expected_output_size); REQUIRE(dispatch.bindings[1].offset == (expected_swapped ? half_offset : 0)); REQUIRE(dispatch.bindings[2].offset == (expected_swapped ? 0 : half_offset)); REQUIRE(dispatch.bindings[3].value == imported.graph.nodes()[1].output); @@ -3307,8 +3763,7 @@ static void schedule_qwen_terminal_q6k_vector_command() { ggml_tensor * logits = ggml_mul_mat(ctx, vocab_weight, input); REQUIRE(logits != nullptr); - schedule_single_matmul_command( - ctx, logits, "loom_libs:ggml_mul_mat_vector_q6_f32_f32", 1, 2048, 151936); + schedule_single_matmul_command(ctx, logits, "loom_libs:ggml_mul_mat_vector_q6_f32_f32", 1, 2048, 151936); ggml_free(ctx); } @@ -3392,8 +3847,8 @@ static void schedule_get_rows_q8_1_alternate_command(ggml_type embedding_weight_ REQUIRE(commands.commands[0].bindings[3].name == "next_output"); REQUIRE(commands.commands[0].bindings[3].origin == ggml::hrx::CommandBindingOrigin::Transient); REQUIRE(commands.commands[1].bindings.size() == 3); - REQUIRE(commands.commands[1].bindings[0].name == (vocab_kernel_name == "qwen3_moe:ggml_linear_q6k_q8_1_x4" ? - "q8_input" : "input")); + REQUIRE(commands.commands[1].bindings[0].name == + (vocab_kernel_name == "qwen3_moe:ggml_linear_q6k_q8_1_x4" ? "q8_input" : "input")); ggml_free(ctx); } @@ -3709,8 +4164,7 @@ static void run_get_rows_rmsnorm_binary_dispatch_checks() { REQUIRE(plan.dispatches.size() == 1); REQUIRE(plan.transients.size() == 2); const ggml::hrx::Dispatch & dispatch = plan.dispatches.front(); - REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == - "loom_libs:ggml_get_rows_rmsnorm_binary_q8_1_x4_f16"); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_get_rows_rmsnorm_binary_q8_1_x4_f16"); REQUIRE(dispatch.bindings.size() == 7); require_compile_parameter(dispatch, "ggml.get_rows_f32.weight_format", "6"); require_compile_parameter(dispatch, "ggml.get_rows_rmsnorm.rms_epsilon", "9.99999975e-06"); @@ -3735,6 +4189,131 @@ static void run_get_rows_rmsnorm_binary_dispatch_checks() { ggml_free(ctx); } +static void run_get_rows_scale_dispatch_checks() { + ggml_init_params params = {}; + params.mem_size = 4 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2560, 64); + ggml_tensor * ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 5); + REQUIRE(weight != nullptr); + REQUIRE(ids != nullptr); + ggml_tensor * rows = ggml_get_rows(ctx, weight, ids); + ggml_tensor * output = ggml_scale(ctx, rows, 0.177800179f); + REQUIRE(rows != nullptr); + REQUIRE(output != nullptr); + ggml_set_output(rows); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + require_scheduled_command_program( + graph, [&](const ggml::hrx::Graph & imported_graph, const ggml::hrx::CommandPlan & plan, + const ggml::hrx::CommandProgram & commands) { + REQUIRE(imported_graph.nodes().size() == 2); + REQUIRE(plan.dispatches.size() == 1); + const ggml::hrx::Dispatch & dispatch = plan.dispatches.front(); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_get_rows_scale_f32"); + REQUIRE(dispatch.bindings.size() == 4); + require_compile_parameter(dispatch, "ggml.get_rows_f32.weight_format", "32"); + require_compile_parameter(dispatch, "ggml.get_rows_scale_f32.scale", expected_config_value(0.177800179f)); + require_compile_parameter(dispatch, "ggml.get_rows_scale_f32.bias", expected_config_value(0.0f)); + REQUIRE(commands.commands.size() == 1); + REQUIRE(commands.commands.front().bindings[0].name == "token_ids"); + REQUIRE(commands.commands.front().bindings[1].name == "weight"); + REQUIRE(commands.commands.front().bindings[2].name == "raw_output"); + REQUIRE(commands.commands.front().bindings[3].name == "output"); + }); + + ggml_tensor * q5_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q5_1, 640, 262144); + ggml_tensor * q5_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 64); + REQUIRE(q5_weight != nullptr); + REQUIRE(q5_ids != nullptr); + ggml_tensor * q5_rows = ggml_get_rows(ctx, q5_weight, q5_ids); + ggml_tensor * q5_output = ggml_scale(ctx, q5_rows, 25.2982216f); + REQUIRE(q5_rows != nullptr); + REQUIRE(q5_output != nullptr); + ggml_set_output(q5_rows); + + ggml_cgraph * q5_graph = ggml_new_graph(ctx); + REQUIRE(q5_graph != nullptr); + ggml_build_forward_expand(q5_graph, q5_output); + require_scheduled_command_program( + q5_graph, [&](const ggml::hrx::Graph & imported_graph, const ggml::hrx::CommandPlan & plan, + const ggml::hrx::CommandProgram & commands) { + REQUIRE(imported_graph.nodes().size() == 2); + REQUIRE(plan.dispatches.size() == 1); + const ggml::hrx::Dispatch & dispatch = plan.dispatches.front(); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_get_rows_scale_f32"); + REQUIRE(dispatch.bindings.size() == 4); + require_compile_parameter(dispatch, "ggml.get_rows_f32.weight_format", "51"); + require_compile_parameter(dispatch, "ggml.get_rows_scale_f32.scale", expected_config_value(25.2982216f)); + require_compile_parameter(dispatch, "ggml.get_rows_scale_f32.bias", expected_config_value(0.0f)); + REQUIRE(commands.commands.size() == 1); + }); + + ggml_tensor * q5_decode_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); + REQUIRE(q5_decode_ids != nullptr); + ggml_tensor * q5_decode_rows = ggml_get_rows(ctx, q5_weight, q5_decode_ids); + ggml_tensor * q5_decode_output = ggml_scale(ctx, q5_decode_rows, 25.2982216f); + REQUIRE(q5_decode_rows != nullptr); + REQUIRE(q5_decode_output != nullptr); + ggml_cgraph * q5_decode_graph = ggml_new_graph(ctx); + REQUIRE(q5_decode_graph != nullptr); + ggml_build_forward_expand(q5_decode_graph, q5_decode_output); + const std::vector q5_decode_kernels = scheduled_kernel_names(q5_decode_graph); + REQUIRE(std::find(q5_decode_kernels.begin(), q5_decode_kernels.end(), "loom_libs:ggml_get_rows_scale_f32") == + q5_decode_kernels.end()); + REQUIRE(std::find(q5_decode_kernels.begin(), q5_decode_kernels.end(), "loom_libs:ggml_get_rows_f32") != + q5_decode_kernels.end()); + REQUIRE(std::find(q5_decode_kernels.begin(), q5_decode_kernels.end(), "loom_libs:ggml_scale_f32") != + q5_decode_kernels.end()); + + ggml_tensor * q4_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_0, 640, 64); + ggml_tensor * q4_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 5); + REQUIRE(q4_weight != nullptr); + REQUIRE(q4_ids != nullptr); + ggml_tensor * q4_rows = ggml_get_rows(ctx, q4_weight, q4_ids); + ggml_tensor * q4_output = ggml_scale(ctx, q4_rows, 0.5f); + REQUIRE(q4_rows != nullptr); + REQUIRE(q4_output != nullptr); + ggml_cgraph * q4_graph = ggml_new_graph(ctx); + REQUIRE(q4_graph != nullptr); + ggml_build_forward_expand(q4_graph, q4_output); + const std::vector q4_kernels = scheduled_kernel_names(q4_graph); + REQUIRE(std::find(q4_kernels.begin(), q4_kernels.end(), "loom_libs:ggml_get_rows_scale_f32") == q4_kernels.end()); + REQUIRE(std::find(q4_kernels.begin(), q4_kernels.end(), "loom_libs:ggml_get_rows_f32") != q4_kernels.end()); + + ggml_tensor * side_use = ggml_add(ctx, rows, rows); + REQUIRE(side_use != nullptr); + ggml_cgraph * fallback_graph = ggml_new_graph(ctx); + REQUIRE(fallback_graph != nullptr); + ggml_build_forward_expand(fallback_graph, output); + ggml_build_forward_expand(fallback_graph, side_use); + const std::vector fallback_kernels = scheduled_kernel_names(fallback_graph); + REQUIRE(std::find(fallback_kernels.begin(), fallback_kernels.end(), "loom_libs:ggml_get_rows_scale_f32") == + fallback_kernels.end()); + REQUIRE(std::find(fallback_kernels.begin(), fallback_kernels.end(), "loom_libs:ggml_get_rows_f32") != + fallback_kernels.end()); + + ggml_tensor * inplace_rows = ggml_get_rows(ctx, weight, ids); + ggml_tensor * inplace_output = ggml_scale_inplace(ctx, inplace_rows, 0.5f); + REQUIRE(inplace_rows != nullptr); + REQUIRE(inplace_output != nullptr); + ggml_cgraph * inplace_graph = ggml_new_graph(ctx); + REQUIRE(inplace_graph != nullptr); + ggml_build_forward_expand(inplace_graph, inplace_output); + const std::vector inplace_kernels = scheduled_kernel_names(inplace_graph); + REQUIRE(std::find(inplace_kernels.begin(), inplace_kernels.end(), "loom_libs:ggml_get_rows_scale_f32") == + inplace_kernels.end()); + REQUIRE(std::find(inplace_kernels.begin(), inplace_kernels.end(), "loom_libs:ggml_get_rows_f32") != + inplace_kernels.end()); + + ggml_free(ctx); +} + static ggml::hrx::Graph build_manual_gather_add_graph(ggml_context * ctx, int64_t hidden_size, int64_t source_token_count, @@ -3777,11 +4356,198 @@ static ggml::hrx::Graph build_manual_gather_add_graph(ggml_context * ctx, const ggml::hrx::ValueId output_value = graph.values().get_or_add_tensor_value(output, ggml::hrx::ValueKind::External); - graph.add_node(GGML_OP_GET_ROWS, selected0_value, { attention_value, row_ids0_value }); - graph.add_node(GGML_OP_GET_ROWS, selected1_value, { residual_value, row_ids1_value }); - graph.add_node(GGML_OP_ADD, output_value, { selected0_value, selected1_value }); - REQUIRE(graph.build_index().success()); - return graph; + graph.add_node(GGML_OP_GET_ROWS, selected0_value, { attention_value, row_ids0_value }); + graph.add_node(GGML_OP_GET_ROWS, selected1_value, { residual_value, row_ids1_value }); + graph.add_node(GGML_OP_ADD, output_value, { selected0_value, selected1_value }); + REQUIRE(graph.build_index().success()); + return graph; +} + +struct GatherAddRmsNormGraph { + ggml_tensor * attention = nullptr; + ggml_tensor * residual = nullptr; + ggml_tensor * row_ids = nullptr; + ggml_tensor * weight = nullptr; + ggml_tensor * selected = nullptr; + ggml_tensor * raw_output = nullptr; + ggml_tensor * normalized_output = nullptr; + ggml_cgraph * graph = nullptr; +}; + +static GatherAddRmsNormGraph build_gather_add_rmsnorm_graph( + ggml_context * ctx, + int64_t hidden_size, + int64_t source_token_count, + int64_t output_token_count, + bool produced_row_ids = false, + bool produced_weight = false, + ggml::hrx::BinaryKind gather_op = ggml::hrx::BinaryKind::Add, + bool first_is_lhs = true) { + GatherAddRmsNormGraph result; + result.attention = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, source_token_count); + result.residual = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, source_token_count); + ggml_tensor * raw_row_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, output_token_count); + ggml_tensor * raw_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size); + result.row_ids = produced_row_ids ? ggml_dup(ctx, raw_row_ids) : raw_row_ids; + result.weight = produced_weight ? ggml_scale(ctx, raw_weight, 0.5f) : raw_weight; + ggml_tensor * selected_attention = ggml_get_rows(ctx, result.attention, result.row_ids); + ggml_tensor * selected_residual = ggml_get_rows(ctx, result.residual, result.row_ids); + result.selected = selected_attention; + switch (gather_op) { + case ggml::hrx::BinaryKind::Add: + result.raw_output = ggml_add(ctx, selected_attention, selected_residual); + break; + case ggml::hrx::BinaryKind::Sub: + result.raw_output = first_is_lhs ? ggml_sub(ctx, selected_attention, selected_residual) : + ggml_sub(ctx, selected_residual, selected_attention); + break; + case ggml::hrx::BinaryKind::Mul: + result.raw_output = ggml_mul(ctx, selected_attention, selected_residual); + break; + case ggml::hrx::BinaryKind::Div: + result.raw_output = first_is_lhs ? ggml_div(ctx, selected_attention, selected_residual) : + ggml_div(ctx, selected_residual, selected_attention); + break; + default: + REQUIRE(false); + } + ggml_tensor * rms = ggml_rms_norm(ctx, result.raw_output, 0.000001f); + result.normalized_output = ggml_mul(ctx, rms, result.weight); + REQUIRE(result.attention != nullptr); + REQUIRE(result.residual != nullptr); + REQUIRE(raw_row_ids != nullptr); + REQUIRE(raw_weight != nullptr); + REQUIRE(result.row_ids != nullptr); + REQUIRE(result.weight != nullptr); + REQUIRE(selected_attention != nullptr); + REQUIRE(selected_residual != nullptr); + REQUIRE(result.raw_output != nullptr); + REQUIRE(rms != nullptr); + REQUIRE(result.normalized_output != nullptr); + + ggml_set_output(result.raw_output); + result.graph = ggml_new_graph(ctx); + REQUIRE(result.graph != nullptr); + ggml_build_forward_expand(result.graph, result.normalized_output); + return result; +} + +static ggml::hrx::DispatchRegistry gather_add_test_registry() { + ggml::hrx::DispatchRegistryBuilder builder; + ggml::hrx::register_gather_add_dispatch(builder); + return builder.build(); +} + +static bool match_gather_add_at(const ggml::hrx::DispatchRegistry & registry, + const ggml::hrx::Graph & graph, + size_t root_index, + const std::vector & covered_nodes, + ggml::hrx::DispatchMatch & match) { + ggml::hrx::CommandPlan plan; + const ggml::hrx::DispatchMatchContext context = { + graph, &graph.nodes()[root_index], + root_index, covered_nodes, + plan, ggml::hrx::ValueId(static_cast(graph.values().size())), + }; + return registry.match(context, match); +} + +static void require_gather_add_rmsnorm_route(ggml_context * ctx, + ggml::hrx::BinaryKind gather_op = ggml::hrx::BinaryKind::Add, + bool first_is_lhs = true) { + GatherAddRmsNormGraph tensors = + build_gather_add_rmsnorm_graph(ctx, 128, 4, 5, false, false, gather_op, first_is_lhs); + require_scheduled_command_program(tensors.graph, [&](const ggml::hrx::Graph &, const ggml::hrx::CommandPlan & plan, + const ggml::hrx::CommandProgram & commands) { + REQUIRE(plan.dispatches.size() == 1); + const ggml::hrx::Dispatch & dispatch = plan.dispatches.front(); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_gather_add_rmsnorm_binary_f32"); + REQUIRE(dispatch.bindings.size() == 6); + require_compile_parameter(dispatch, "ggml.gather_add_rmsnorm_binary_f32.hidden_size", "128"); + require_compile_parameter(dispatch, "ggml.gather_add_rmsnorm_binary_f32.rms_epsilon", "9.99999997e-07"); + require_compile_parameter(dispatch, "ggml.gather_add_rmsnorm_binary_f32.gather_op", + std::to_string(ggml::hrx::binary_kind_config_value(gather_op))); + REQUIRE(commands.commands.size() == 1); + REQUIRE(commands.commands.front().bindings.size() == 6); + REQUIRE(commands.commands.front().bindings[0].name == "attention"); + REQUIRE(commands.commands.front().bindings[1].name == "residual"); + REQUIRE(commands.commands.front().bindings[2].name == "output_ids"); + REQUIRE(commands.commands.front().bindings[3].name == "raw_output"); + REQUIRE(commands.commands.front().bindings[4].name == "weight"); + REQUIRE(commands.commands.front().bindings[5].name == "normalized_output"); + }); +} + +static void require_gather_add_rmsnorm_availability(ggml_context * ctx) { + GatherAddRmsNormGraph tensors = build_gather_add_rmsnorm_graph(ctx, 128, 4, 5, true, true); + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*tensors.graph); + REQUIRE(imported.valid()); + const ggml::hrx::DispatchRegistry registry = gather_add_test_registry(); + const size_t root_index = producer_index_for_tensor(imported.graph, tensors.selected); + const size_t row_index = producer_index_for_tensor(imported.graph, tensors.row_ids); + const size_t weight_index = producer_index_for_tensor(imported.graph, tensors.weight); + std::vector covered_nodes(imported.graph.nodes().size(), false); + ggml::hrx::DispatchMatch match; + + REQUIRE(!match_gather_add_at(registry, imported.graph, root_index, covered_nodes, match)); + + covered_nodes[row_index] = true; + match = {}; + REQUIRE(match_gather_add_at(registry, imported.graph, root_index, covered_nodes, match)); + REQUIRE(match.dispatches.size() == 1); + REQUIRE(kernel_name_for_id(match.dispatches.front().kernel.kernel_id) == "hrx:ggml_gather_add_f32"); + + covered_nodes[weight_index] = true; + match = {}; + REQUIRE(match_gather_add_at(registry, imported.graph, root_index, covered_nodes, match)); + REQUIRE(match.dispatches.size() == 1); + REQUIRE(kernel_name_for_id(match.dispatches.front().kernel.kernel_id) == + "loom_libs:ggml_gather_add_rmsnorm_binary_f32"); +} + +static void require_gather_add_rmsnorm_row_ids_noalias(ggml_context * ctx) { + GatherAddRmsNormGraph tensors = build_gather_add_rmsnorm_graph(ctx, 128, 4, 5); + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*tensors.graph); + REQUIRE(imported.valid()); + const ggml::hrx::Value * row_ids = imported.graph.values().find_tensor(tensors.row_ids); + const ggml::hrx::Value * attention = imported.graph.values().find_tensor(tensors.attention); + REQUIRE(row_ids != nullptr); + REQUIRE(attention != nullptr); + + ggml::hrx::Graph aliased_graph; + for (const ggml::hrx::ValueStorage & storage : imported.graph.values().storages()) { + REQUIRE(aliased_graph.values().add_snapshot_storage(storage).success()); + } + for (ggml::hrx::Value value : imported.graph.values().values()) { + if (value.id == row_ids->id) { + value.storage = attention->storage; + value.storage_root = attention->storage_root; + value.alias_source = attention->id; + value.storage_offset = attention->storage_offset; + value.storage_byte_count = attention->storage_byte_count; + } + REQUIRE(aliased_graph.values().add_snapshot_value(std::move(value)).success()); + } + for (const ggml::hrx::GraphNode & node : imported.graph.nodes()) { + ggml::hrx::GraphNode & copied = aliased_graph.add_node(node.op, node.output, node.inputs); + copied.params = node.params; + } + REQUIRE(aliased_graph.build_index().success()); + + const ggml::hrx::DispatchRegistry registry = gather_add_test_registry(); + size_t root_index = aliased_graph.nodes().size(); + for (size_t i = 0; i < aliased_graph.nodes().size(); ++i) { + if (aliased_graph.nodes()[i].op == GGML_OP_GET_ROWS) { + root_index = i; + break; + } + } + REQUIRE(root_index < aliased_graph.nodes().size()); + std::vector covered_nodes(aliased_graph.nodes().size(), false); + ggml::hrx::DispatchMatch match; + REQUIRE(match_gather_add_at(registry, aliased_graph, root_index, covered_nodes, match)); + REQUIRE(match.dispatches.size() == 1); + REQUIRE(kernel_name_for_id(match.dispatches.front().kernel.kernel_id) == "hrx:ggml_gather_add_f32"); } static bool manual_gather_add_graph_is_supported(ggml_context * ctx, @@ -3875,24 +4641,18 @@ static void require_gather_add_rejects_unavailable_sources(ggml_context * ctx) { ggml::hrx::DispatchMatch match; std::vector covered_nodes(graph.nodes().size(), false); ggml::hrx::DispatchMatchContext context = { - graph, - &graph.nodes()[get_rows_index], - get_rows_index, - covered_nodes, - plan, - ggml::hrx::ValueId(static_cast(graph.values().size())), + graph, &graph.nodes()[get_rows_index], + get_rows_index, covered_nodes, + plan, ggml::hrx::ValueId(static_cast(graph.values().size())), }; REQUIRE(!gather_add_registry.match(context, match)); covered_nodes[producer_index] = true; const ggml::hrx::DispatchMatchContext available_context = { - graph, - &graph.nodes()[get_rows_index], - get_rows_index, - covered_nodes, - plan, - ggml::hrx::ValueId(static_cast(graph.values().size())), + graph, &graph.nodes()[get_rows_index], + get_rows_index, covered_nodes, + plan, ggml::hrx::ValueId(static_cast(graph.values().size())), }; match = {}; REQUIRE(gather_add_registry.match(available_context, match)); @@ -3905,7 +4665,12 @@ static void schedule_gather_add_command(ggml::hrx::Graph & graph, int64_t expected_source_token_count, int64_t expected_output_token_count) { ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(graph, test_dispatch_target())); + if (!scheduler.schedule_graph(graph, test_dispatch_target())) { + for (const std::string & error : scheduler.plan().status.errors()) { + std::fprintf(stderr, "%s\n", error.c_str()); + } + REQUIRE(false); + } REQUIRE(scheduler.plan().valid()); REQUIRE(scheduler.plan().dispatches.size() == 1); @@ -3914,6 +4679,7 @@ static void schedule_gather_add_command(ggml::hrx::Graph & graph, REQUIRE(dispatch.kernel.integer_parameters.at("hidden_size") == expected_hidden_size); REQUIRE(dispatch.kernel.integer_parameters.at("source_token_count") == expected_source_token_count); REQUIRE(dispatch.kernel.integer_parameters.at("output_token_count") == expected_output_token_count); + require_compile_parameter(dispatch, "ggml.gather_add_f32.binary_op", "0"); REQUIRE(dispatch.bindings.size() == 4); const ggml::hrx::CommandProgram commands = @@ -3949,6 +4715,12 @@ static void run_gather_add_dispatch_checks() { REQUIRE(!manual_gather_add_graph_is_supported(ctx, 96, 13, 1, true)); REQUIRE(!partial_gather_add_graph_is_supported(ctx)); require_gather_add_rejects_unavailable_sources(ctx); + require_gather_add_rmsnorm_route(ctx); + require_gather_add_rmsnorm_route(ctx, ggml::hrx::BinaryKind::Sub, false); + require_gather_add_rmsnorm_route(ctx, ggml::hrx::BinaryKind::Mul); + require_gather_add_rmsnorm_route(ctx, ggml::hrx::BinaryKind::Div); + require_gather_add_rmsnorm_availability(ctx); + require_gather_add_rmsnorm_row_ids_noalias(ctx); ggml_free(ctx); } @@ -4007,7 +4779,9 @@ static void schedule_flash_attention_command(ggml_context * ctx, REQUIRE(copy.bindings.size() == 2); REQUIRE(copy.bindings[1].value == dispatch.bindings[2].value); REQUIRE(copy.bindings[0].value != copy.bindings[1].value); - REQUIRE(copy.bindings[1].length == static_cast(key_value_token_count * key_value_head_count * value_head_size) * sizeof(ggml_fp16_t)); + REQUIRE(copy.bindings[1].length == + static_cast(key_value_token_count * key_value_head_count * value_head_size) * + sizeof(ggml_fp16_t)); require_compile_parameter(dispatch, "ggml.flash_attention.value_layout", "1"); } else { REQUIRE(dispatch.kernel.compile_parameters.count("ggml.flash_attention.value_layout") == 0); @@ -4116,15 +4890,16 @@ static void run_qwen_flash_attention_dispatch_checks() { ggml_context * ctx = ggml_init(params); REQUIRE(ctx != nullptr); - for (const auto & shape : { std::array{ 256, 512, 64, 1 }, + for (const auto & shape : { + std::array{ 256, 512, 64, 1 }, std::array{ 512, 512, 256, 1 }, std::array{ 512, 2048, 320, 1 }, std::array{ 255, 512, 256, 0 }, - std::array{ 512, 513, 256, 0 } }) { + std::array{ 512, 513, 256, 0 } + }) { const float scale = 1.0f / std::sqrt(static_cast(shape[2])); - ggml_tensor * output = build_qwen_flash_attention_graph(ctx, shape[0], shape[1], 4, 2, - GGML_TYPE_F32, GGML_TYPE_F16, true, false, - shape[2], scale); + ggml_tensor * output = build_qwen_flash_attention_graph(ctx, shape[0], shape[1], 4, 2, GGML_TYPE_F32, + GGML_TYPE_F16, true, false, shape[2], scale); schedule_flash_attention_command(ctx, output, "loom_libs:ggml_flash_attention_f32_f16_wmma", "ggml.flash_attention.", shape[0], shape[1], 4, 2, shape[2], shape[2], scale, shape[3] != 0); @@ -4688,6 +5463,7 @@ static void run_packed_f16_producer_consumer_checks() { { GGML_TYPE_Q4_K, 1024, 4096, 4096, false, false }, { GGML_TYPE_Q6_K, 512, 4096, 2048, true, false, true }, }; + for (const auto & test : cases) { ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 256, test.tokens); ggml_tensor * gate_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 256, test.hidden); @@ -4710,12 +5486,14 @@ static void run_packed_f16_producer_consumer_checks() { REQUIRE(plan.valid()); const auto * value = imported.graph.values().find_tensor(glu); REQUIRE(value != nullptr); - const auto * generated = plan.metadata.find_generated_resource(value->id, ggml::hrx::GeneratedResourceRole::F16K16Major); + const auto * generated = + plan.metadata.find_generated_resource(value->id, ggml::hrx::GeneratedResourceRole::F16K16Major); const auto * alternate = plan.metadata.find_alternate_value(value->id, GGML_TYPE_F16, value->byte_count / 2); REQUIRE((generated != nullptr) == (test.packed || test.packed_copy)); REQUIRE((alternate != nullptr) != test.packed); const auto producer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32"; + return kernel_name_for_id(dispatch.kernel.kernel_id) == + "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32"; }); REQUIRE(producer != plan.dispatches.end()); require_compile_parameter(*producer, "ggml.mul_mat_swiglu.f16_output_layout", test.packed ? "1" : "0"); @@ -4737,7 +5515,8 @@ static void run_packed_f16_producer_consumer_checks() { REQUIRE(std::next(copy) != plan.dispatches.end()); REQUIRE(std::next(copy)->bindings.front().value == generated->generated_value); } - const auto commands = ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + const auto commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); REQUIRE(command_program_verifies(commands)); } @@ -4813,14 +5592,10 @@ static void run_tiled_matmul_alternate_publish_checks() { }; const Case cases[] = { - { GGML_TYPE_Q6_K, - "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32_f16_alternate", - "loom_libs:ggml_mul_mat_q6_k_f16_wmma_prefill_wave32", - GGML_TYPE_F16 }, - { GGML_TYPE_Q5_K, - "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32", - "loom_libs:ggml_mul_mat_q5_k_iq4_xs_q8_1_x4_wmma_token256", - GGML_TYPE_Q8_1 }, + { GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32_f16_alternate", + "loom_libs:ggml_mul_mat_q6_k_f16_wmma_prefill_wave32", GGML_TYPE_F16 }, + { GGML_TYPE_Q5_K, "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32", + "loom_libs:ggml_mul_mat_q5_k_iq4_xs_q8_1_x4_wmma_token256", GGML_TYPE_Q8_1 }, }; for (const Case & test : cases) { @@ -4864,7 +5639,8 @@ static void run_tiled_matmul_alternate_publish_checks() { REQUIRE(producer->bindings.size() == 4); REQUIRE(producer->bindings[3].value == alternate->alternate_value); } else { - const auto quantize = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { + const auto quantize = + std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { return kernel_name_for_id(dispatch.kernel.kernel_id) == "qwen3_moe:ggml_quantize_q8_1_x4_f32" && dispatch.bindings.front().value == projection_value->id; }); @@ -4878,7 +5654,8 @@ static void run_tiled_matmul_alternate_publish_checks() { REQUIRE(consumer != plan.dispatches.end()); REQUIRE(consumer->bindings.front().value == alternate->alternate_value); - const auto commands = ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + const auto commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); REQUIRE(command_program_verifies(commands)); ggml_free(ctx); @@ -4914,9 +5691,9 @@ static void run_tiled_pair_matmul_postops_dispatch_checks() { ggml_tensor * output = ggml_add(ctx, build_glu(), bias); REQUIRE(output != nullptr); schedule_fused_matmul_swiglu_postops_command( - ctx, output, "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_bias_publish_f32", 4, + ctx, output, "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_bias_residual_publish_f32", 4, { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT, GGML_OP_GLU, GGML_OP_ADD }, - { "input", "gate_weight", "up_weight", "bias", "output" }); + { "input", "gate_weight", "up_weight", "bias", "residual_input", "residual_output" }); } { @@ -4925,9 +5702,9 @@ static void run_tiled_pair_matmul_postops_dispatch_checks() { ggml_tensor * output = ggml_add(ctx, build_glu(), residual); REQUIRE(output != nullptr); schedule_fused_matmul_swiglu_postops_command( - ctx, output, "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_residual_publish_f32", 4, + ctx, output, "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_bias_residual_publish_f32", 4, { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT, GGML_OP_GLU, GGML_OP_ADD }, - { "input", "gate_weight", "up_weight", "residual_input", "residual_output" }); + { "input", "gate_weight", "up_weight", "bias", "residual_input", "residual_output" }); } { @@ -4980,6 +5757,7 @@ static void run_swiglu_q8_output_checks() { { 5, 5120, 17408, GGML_TYPE_Q4_K, 1, false, false, false, true, false, false }, { 5, 5120, 17408, GGML_TYPE_Q4_K, 1, false, false, false, false, true, false }, }; + for (const auto & test : cases) { ggml_init_params params = {}; params.mem_size = 4 * 1024 * 1024; @@ -4988,7 +5766,8 @@ static void run_swiglu_q8_output_checks() { REQUIRE(ctx != nullptr); ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, test.inputs, test.tokens); ggml_tensor * gate_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, test.inputs, test.outputs); - ggml_tensor * up_weight = test.weight_alias ? ggml_view_tensor(ctx, gate_weight) : + ggml_tensor * up_weight = test.weight_alias ? + ggml_view_tensor(ctx, gate_weight) : ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, test.inputs, test.outputs); ggml_tensor * gate = ggml_mul_mat(ctx, gate_weight, input); ggml_tensor * up = ggml_mul_mat(ctx, up_weight, input); @@ -5012,7 +5791,8 @@ static void run_swiglu_q8_output_checks() { const bool scheduled = scheduler.schedule_graph(imported.graph, test_dispatch_target()); const auto & plan = scheduler.plan(); const auto producer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot_q8_output"; + return kernel_name_for_id(dispatch.kernel.kernel_id) == + "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot_q8_output"; }); REQUIRE((producer != plan.dispatches.end()) == test.packed); if (test.outputs > 32768) { @@ -5044,13 +5824,15 @@ static void run_swiglu_q8_output_checks() { if (side != nullptr) { const auto * side_value = imported.graph.values().find_tensor(side); REQUIRE(side_value != nullptr); - const auto consumer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { + const auto consumer = + std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { return !dispatch.bindings.empty() && dispatch.bindings.back().value == side_value->id; }); REQUIRE(consumer != plan.dispatches.end()); REQUIRE(consumer->bindings.front().value == output->id); } - const auto commands = ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + const auto commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); REQUIRE(command_program_verifies(commands)); ggml_free(ctx); @@ -5079,10 +5861,12 @@ static void run_k16_major_preparation_reuse_checks() { { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K, false, true }, { GGML_TYPE_Q6_K, GGML_TYPE_Q6_K, false, true }, }; + for (const auto & test : cases) { ggml_tensor * first = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 4096, 512); ggml_tensor * second = test.reshape ? ggml_reshape_2d(ctx, first, 2048, 1024) : - test.shared ? first : ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 4096, 512); + test.shared ? first : + ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 4096, 512); ggml_tensor * inputs[] = { first, second }; ggml_tensor * outputs[2]; ggml_cgraph * graph = ggml_new_graph(ctx); @@ -5098,7 +5882,8 @@ static void run_k16_major_preparation_reuse_checks() { REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); const auto & plan = scheduler.plan(); REQUIRE(plan.valid()); - const auto copy_count = std::count_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { + const auto copy_count = + std::count_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_copy_f16_k16_major"; }); REQUIRE(copy_count == (test.shared ? 1 : 2)); @@ -5108,18 +5893,20 @@ static void run_k16_major_preparation_reuse_checks() { const auto * output = imported.graph.values().find_tensor(outputs[i]); REQUIRE(input != nullptr); REQUIRE(output != nullptr); - const auto * generated = plan.metadata.find_generated_resource(input->id, ggml::hrx::GeneratedResourceRole::F16K16Major); + const auto * generated = + plan.metadata.find_generated_resource(input->id, ggml::hrx::GeneratedResourceRole::F16K16Major); REQUIRE(generated != nullptr); REQUIRE(generated->byte_count == input->byte_count / 2); packed[i] = generated->generated_value; - const auto consumer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { - return dispatch.bindings.back().value == output->id; - }); + const auto consumer = + std::find_if(plan.dispatches.begin(), plan.dispatches.end(), + [&](const auto & dispatch) { return dispatch.bindings.back().value == output->id; }); REQUIRE(consumer != plan.dispatches.end()); REQUIRE(consumer->bindings.front().value == packed[i]); } REQUIRE((packed[0] == packed[1]) == test.shared); - const auto commands = ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + const auto commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); REQUIRE(command_program_verifies(commands)); } @@ -5154,6 +5941,7 @@ static void run_rmsnorm_binary_k16_output_checks() { { GGML_TYPE_Q4_K, 5120, 512, 1, false, false, true, true, false }, { GGML_TYPE_Q4_K, 5120, 512, 1, false, false, false, false, false }, }; + for (const auto & test : cases) { ggml_init_params params = {}; params.mem_size = 4 * 1024 * 1024; @@ -5162,10 +5950,11 @@ static void run_rmsnorm_binary_k16_output_checks() { REQUIRE(ctx != nullptr); ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, test.hidden, test.tokens); ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, test.hidden); - ggml_tensor * rms = test.inplace ? ggml_rms_norm_inplace(ctx, input, 1.0e-6f) : - ggml_rms_norm(ctx, input, 1.0e-6f); + ggml_tensor * rms = + test.inplace ? ggml_rms_norm_inplace(ctx, input, 1.0e-6f) : ggml_rms_norm(ctx, input, 1.0e-6f); ggml_tensor * norm = test.inplace ? ggml_mul_inplace(ctx, rms, weight) : - test.multiply ? ggml_mul(ctx, rms, weight) : ggml_add(ctx, rms, weight); + test.multiply ? ggml_mul(ctx, rms, weight) : + ggml_add(ctx, rms, weight); if (test.observed) { ggml_set_output(norm); } @@ -5191,7 +5980,8 @@ static void run_rmsnorm_binary_k16_output_checks() { const auto * output = imported.graph.values().find_tensor(norm); REQUIRE(output != nullptr); if (test.packed) { - const auto * generated = plan.metadata.find_generated_resource(output->id, ggml::hrx::GeneratedResourceRole::F16K16Major); + const auto * generated = + plan.metadata.find_generated_resource(output->id, ggml::hrx::GeneratedResourceRole::F16K16Major); REQUIRE(generated != nullptr); REQUIRE(generated->byte_count == output->byte_count / 2); REQUIRE(producer->bindings.size() == 4); @@ -5209,13 +5999,15 @@ static void run_rmsnorm_binary_k16_output_checks() { if (side != nullptr) { const auto * side_value = imported.graph.values().find_tensor(side); REQUIRE(side_value != nullptr); - const auto consumer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { + const auto consumer = + std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { return !dispatch.bindings.empty() && dispatch.bindings.back().value == side_value->id; }); REQUIRE(consumer != plan.dispatches.end()); REQUIRE(consumer->bindings.front().value == output->id); } - const auto commands = ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + const auto commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); REQUIRE(command_program_verifies(commands)); ggml_free(ctx); @@ -5347,6 +6139,7 @@ static void run_rmsnorm_gate_packed_output_checks() { { GGML_TYPE_Q4_K, 128, 4096, 256, false, false, 0 }, { GGML_TYPE_Q4_K, 256, 256, 512, false, true, 0 }, }; + for (const auto & test : cases) { const int64_t heads = test.channels / test.hidden; ggml_tensor * input = heads == 1 ? ggml_new_tensor_2d(ctx, GGML_TYPE_F32, test.hidden, test.tokens) : @@ -5380,7 +6173,8 @@ static void run_rmsnorm_gate_packed_output_checks() { const auto * alternate = plan.metadata.find_alternate_value(output->id, GGML_TYPE_F16, output->byte_count / 2); REQUIRE((alternate == nullptr) == test.packed); if (test.packed) { - const auto * generated = plan.metadata.find_generated_resource(matmul_input->id, ggml::hrx::GeneratedResourceRole::F16K16Major); + const auto * generated = + plan.metadata.find_generated_resource(matmul_input->id, ggml::hrx::GeneratedResourceRole::F16K16Major); REQUIRE(generated != nullptr); REQUIRE(producer->bindings.back().value == generated->generated_value); REQUIRE(std::next(producer) != plan.dispatches.end()); @@ -5389,7 +6183,8 @@ static void run_rmsnorm_gate_packed_output_checks() { REQUIRE(std::count_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_copy_f16_k16_major"; }) == test.copies); - const auto commands = ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + const auto commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); REQUIRE(command_program_verifies(commands)); } @@ -5402,6 +6197,7 @@ static void run_rmsnorm_gate_q8_output_checks() { params.no_alloc = true; ggml_context * ctx = ggml_init(params); REQUIRE(ctx != nullptr); + const struct { ggml_type type; int64_t hidden; @@ -5428,6 +6224,7 @@ static void run_rmsnorm_gate_q8_output_checks() { { GGML_TYPE_Q4_K, 128, 6144, 5, 48, false, false, false }, { GGML_TYPE_Q6_K, 128, 6144, 5, 256, false, true, false }, }; + for (const auto & test : cases) { const int64_t heads = test.channels / test.hidden; ggml_tensor * input = heads == 1 ? ggml_new_tensor_2d(ctx, GGML_TYPE_F32, test.hidden, test.tokens) : @@ -5466,12 +6263,12 @@ static void run_rmsnorm_gate_q8_output_checks() { return kernel_name_for_id(dispatch.kernel.kernel_id) == "qwen3_moe:ggml_quantize_q8_1_x4_f32"; }) == 0); REQUIRE(std::any_of(std::next(producer), plan.dispatches.end(), [&](const auto & dispatch) { - return std::any_of(dispatch.bindings.begin(), dispatch.bindings.end(), [&](const auto & binding) { - return binding.value == alternate->alternate_value; - }); + return std::any_of(dispatch.bindings.begin(), dispatch.bindings.end(), + [&](const auto & binding) { return binding.value == alternate->alternate_value; }); })); } - const auto commands = ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + const auto commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); REQUIRE(command_program_verifies(commands)); } @@ -5484,6 +6281,7 @@ static void run_quantized_conv4_dispatch_checks() { params.no_alloc = true; ggml_context * ctx = ggml_init(params); REQUIRE(ctx != nullptr); + const struct { ggml_type type; int64_t inputs; @@ -5508,6 +6306,7 @@ static void run_quantized_conv4_dispatch_checks() { { GGML_TYPE_Q4_K, 5120, 10240, 512, true, false, false, true }, { GGML_TYPE_Q4_K, 5120, 10240, 512, false, true, false, true }, }; + for (const auto & test : cases) { ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, test.inputs, test.tokens); ggml_tensor * weight = ggml_new_tensor_2d(ctx, test.type, test.inputs, test.channels); @@ -5524,8 +6323,8 @@ static void run_quantized_conv4_dispatch_checks() { } ggml_tensor * filter = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 4, test.channels); ggml_tensor * window = ggml_concat(ctx, state, ggml_transpose(ctx, x), 0); - ggml_tensor * tail = ggml_view_3d(ctx, window, 3, test.channels, 1, - window->nb[1], window->nb[2], test.tokens * sizeof(float)); + ggml_tensor * tail = + ggml_view_3d(ctx, window, 3, test.channels, 1, window->nb[1], window->nb[2], test.tokens * sizeof(float)); ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 3 * test.channels, 1); ggml_tensor * target = ggml_view_2d(ctx, cache, 3 * test.channels, 1, cache->nb[1], 0); ggml_tensor * output = ggml_silu(ctx, ggml_ssm_conv(ctx, window, filter)); @@ -5551,11 +6350,13 @@ static void run_quantized_conv4_dispatch_checks() { REQUIRE(fused->bindings.size() == (test.gathered_state ? 5 : 6)); require_compile_parameter(*fused, "ggml.mul_mat.weight_format", test.type == GGML_TYPE_Q4_K ? "4" : "6"); if (test.gathered_state) { - const auto finish = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { + const auto finish = + std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:llm_ssm_conv_dconv4_silu_prefill_finish_f32"; }); - const auto gather = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { + const auto gather = + std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { return kernel_name_for_id(dispatch.kernel.kernel_id).find("get_rows") != std::string::npos; }); REQUIRE(finish != plan.dispatches.end()); @@ -5569,20 +6370,156 @@ static void run_quantized_conv4_dispatch_checks() { return transient.name == "llm.ssm_conv.window_snapshot"; })); } - const auto commands = ggml::hrx::build_command_program( - imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + const auto commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); REQUIRE(command_program_verifies(commands)); } ggml_free(ctx); } +static void run_generic_ssm_conv_binary_dispatch_checks() { + const struct { + ggml::hrx::BinaryKind kind; + bool conv_is_lhs; + } cases[] = { + { ggml::hrx::BinaryKind::Add, true }, + { ggml::hrx::BinaryKind::Sub, true }, + { ggml::hrx::BinaryKind::Sub, false }, + { ggml::hrx::BinaryKind::Mul, true }, + { ggml::hrx::BinaryKind::Div, true }, + { ggml::hrx::BinaryKind::Div, false }, + }; + + for (const auto & test : cases) { + ggml_init_params params = {}; + params.mem_size = 4 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + constexpr int64_t d_conv = 4; + constexpr int64_t d_inner = 64; + constexpr int64_t n_t = 8; + ggml_tensor * window = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, d_conv + n_t - 1, d_inner); + ggml_tensor * filter = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, d_conv, d_inner); + ggml_tensor * operand = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, d_inner, n_t); + ggml_tensor * conv = ggml_ssm_conv(ctx, window, filter); + ggml_tensor * output = nullptr; + switch (test.kind) { + case ggml::hrx::BinaryKind::Add: + output = ggml_add(ctx, conv, operand); + break; + case ggml::hrx::BinaryKind::Sub: + output = test.conv_is_lhs ? ggml_sub(ctx, conv, operand) : ggml_sub(ctx, operand, conv); + break; + case ggml::hrx::BinaryKind::Mul: + output = ggml_mul(ctx, conv, operand); + break; + case ggml::hrx::BinaryKind::Div: + output = test.conv_is_lhs ? ggml_div(ctx, conv, operand) : ggml_div(ctx, operand, conv); + break; + default: + REQUIRE(false); + } + + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, output); + auto imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + const auto & plan = scheduler.plan(); + REQUIRE(plan.valid()); + REQUIRE(plan.dispatches.size() == 1); + const auto & dispatch = plan.dispatches.front(); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:llm_ssm_conv_binary_f32"); + REQUIRE(dispatch.bindings.size() == 4); + require_compile_parameter(dispatch, "llm.ssm_conv.generic.binary_op", + std::to_string(ggml::hrx::binary_kind_config_value(test.kind))); + require_compile_parameter(dispatch, "llm.ssm_conv.generic.binary_lhs", test.conv_is_lhs ? "1" : "0"); + REQUIRE(dispatch.bindings[2].value == imported.graph.values().find_tensor(operand)->id); + REQUIRE(dispatch.bindings[3].value == imported.graph.values().find_tensor(output)->id); + + const auto commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(command_program_verifies(commands)); + ggml_free(ctx); + } +} + +static void run_lfm_dconv3_dispatch_checks() { + const struct { + int64_t tokens; + int64_t hidden; + int64_t state_rows; + int64_t filter_rows; + bool disabled; + bool fused; + } cases[] = { + {1, 2048, 2, 3, false, true}, + {64, 2048, 2, 3, false, true}, + {2, 2048, 2, 3, false, false}, + {64, 2016, 2, 3, false, false}, + {64, 2048, 3, 3, false, false}, + {64, 2048, 2, 4, false, false}, + {64, 2048, 2, 3, true, false}, + }; + + for (const auto & test : cases) { + ggml_init_params params = {}; + params.mem_size = 4 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + ggml_tensor * x = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, test.hidden, test.tokens, 1); + ggml_tensor * state = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, test.state_rows, test.hidden, 1); + ggml_tensor * filter = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, test.filter_rows, test.hidden); + ggml_tensor * window = ggml_concat(ctx, state, ggml_transpose(ctx, x), 0); + ggml_tensor * tail = ggml_view_3d(ctx, window, test.state_rows, test.hidden, 1, window->nb[1], window->nb[2], + test.tokens * sizeof(float)); + ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, test.state_rows * test.hidden, 1); + ggml_tensor * target = ggml_view_2d(ctx, cache, test.state_rows * test.hidden, 1, cache->nb[1], 0); + ggml_tensor * output = ggml_ssm_conv(ctx, window, filter); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, ggml_cpy(ctx, tail, target)); + ggml_build_forward_expand(graph, output); + if (test.disabled) { + REQUIRE(setenv("GGML_HRX_DISABLE_LFM_DCONV3_FUSION", "1", 1) == 0); + } + auto imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + const auto & plan = scheduler.plan(); + REQUIRE(plan.valid()); + const auto fused = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { + return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:llm_ssm_conv_lfm_dconv3_state_f32"; + }); + REQUIRE((fused != plan.dispatches.end()) == test.fused); + if (test.fused) { + REQUIRE(fused->bindings.size() == 5); + require_compile_parameter(*fused, "llm.ssm_conv.lfm_dconv3.n_t", std::to_string(test.tokens)); + } + const auto commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(command_program_verifies(commands)); + if (test.disabled) { + REQUIRE(unsetenv("GGML_HRX_DISABLE_LFM_DCONV3_FUSION") == 0); + } + ggml_free(ctx); + } +} + static void run_gdn_native_projection_pair_dispatch_checks() { ggml_init_params params = {}; params.mem_size = 8 * 1024 * 1024; params.no_alloc = true; ggml_context * ctx = ggml_init(params); REQUIRE(ctx != nullptr); + const struct { int64_t tokens; int64_t inputs; @@ -5608,6 +6545,7 @@ static void run_gdn_native_projection_pair_dispatch_checks() { {1, 5120, 48, false, false, true, GGML_TYPE_Q4_K, false}, {1, 5120, 48, false, false, false, GGML_TYPE_Q6_K, false}, }; + for (const auto & test : cases) { const int64_t width = 128; const int64_t qheads = test.heads / 4; @@ -5625,14 +6563,15 @@ static void run_gdn_native_projection_pair_dispatch_checks() { ggml_tensor * beta_raw = ggml_mul_mat(ctx, second_weight, second_input); ggml_tensor * bias = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, test.heads); ggml_tensor * scale = ggml_dup_tensor(ctx, bias); - ggml_tensor * gate = ggml_mul(ctx, ggml_softplus(ctx, ggml_add(ctx, - ggml_reshape_3d(ctx, alpha, test.heads, test.tokens, 1), bias)), scale); + ggml_tensor * gate = ggml_mul( + ctx, ggml_softplus(ctx, ggml_add(ctx, ggml_reshape_3d(ctx, alpha, test.heads, test.tokens, 1), bias)), + scale); gate = ggml_reshape_4d(ctx, gate, 1, test.heads, test.tokens, 1); ggml_tensor * beta = ggml_sigmoid(ctx, ggml_reshape_4d(ctx, beta_raw, 1, test.heads, test.tokens, 1)); ggml_tensor * raw = ggml_scale(ctx, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden, test.tokens), 0.5f); const auto view = [&](int64_t heads, size_t offset) { - return ggml_view_4d(ctx, raw, width, heads, test.tokens, 1, width * sizeof(float), - hidden * sizeof(float), hidden * test.tokens * sizeof(float), offset); + return ggml_view_4d(ctx, raw, width, heads, test.tokens, 1, width * sizeof(float), hidden * sizeof(float), + hidden * test.tokens * sizeof(float), offset); }; ggml_tensor * q = ggml_l2_norm(ctx, view(qheads, 0), 1.e-6f); ggml_tensor * k = ggml_l2_norm(ctx, view(qheads, width * qheads * sizeof(float)), 1.e-6f); @@ -5644,10 +6583,10 @@ static void run_gdn_native_projection_pair_dispatch_checks() { ggml_tensor * gdn = ggml_gated_delta_net(ctx, q, k, v, gate, beta, state, 5); const int64_t count = std::min(test.tokens, 5); const size_t attention_bytes = width * test.heads * test.tokens * sizeof(float); - ggml_tensor * attention = ggml_view_4d(ctx, gdn, width, test.heads, test.tokens, 1, - width * sizeof(float), width * test.heads * sizeof(float), attention_bytes, 0); - ggml_tensor * snapshots = ggml_view_3d(ctx, gdn, state_elements, 1, count, - state_bytes, state_bytes, attention_bytes); + ggml_tensor * attention = ggml_view_4d(ctx, gdn, width, test.heads, test.tokens, 1, width * sizeof(float), + width * test.heads * sizeof(float), attention_bytes, 0); + ggml_tensor * snapshots = + ggml_view_3d(ctx, gdn, state_elements, 1, count, state_bytes, state_bytes, attention_bytes); ggml_tensor * target = ggml_view_3d(ctx, cache, state_elements, 1, count, state_bytes, 4 * state_bytes, 0); ggml_cgraph * graph = ggml_new_graph(ctx); ggml_build_forward_expand(graph, raw); @@ -5678,8 +6617,8 @@ static void run_gdn_native_projection_pair_dispatch_checks() { REQUIRE(fused->bindings[3].value == imported.graph.values().find_tensor(alpha)->id); REQUIRE(fused->bindings[4].value == imported.graph.values().find_tensor(beta_raw)->id); } - const auto commands = ggml::hrx::build_command_program(imported.graph, plan, - ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + const auto commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); REQUIRE(command_program_verifies(commands)); } @@ -5692,6 +6631,7 @@ static void run_gdn_selected_snapshot_dispatch_checks() { params.no_alloc = true; ggml_context * ctx = ggml_init(params); REQUIRE(ctx != nullptr); + const struct { int64_t tokens; int64_t sequences; @@ -5784,7 +6724,8 @@ static void run_gdn_selected_snapshot_dispatch_checks() { return kernel_name_for_id(dispatch.kernel.kernel_id)=="loom_libs:ggml_get_rows_f32"; })); } - const auto commands=ggml::hrx::build_command_program(imported.graph,plan,ggml::hrx::get_qwen_kernel_corpus(),"gfx1151"); + const auto commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); REQUIRE(command_program_verifies(commands)); } @@ -5813,6 +6754,7 @@ static void run_gdn_rmsnorm_gate_dispatch_checks() { { GGML_TYPE_Q4_K, 512, false, true, false, false }, { GGML_TYPE_Q4_K, 512, false, false, true, false }, }; + for (const auto & test : cases) { const int64_t width = 128; const int64_t heads = 2; @@ -5825,17 +6767,15 @@ static void run_gdn_rmsnorm_gate_dispatch_checks() { ggml_tensor * beta = ggml_dup_tensor(ctx, g); ggml_tensor * state = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, width, width, heads, 1); ggml_tensor * gdn = ggml_gated_delta_net(ctx, q, k, v, g, beta, state, 1); - ggml_tensor * attention = ggml_view_4d(ctx, gdn, width, heads, test.tokens, 1, - width * sizeof(float), channels * sizeof(float), - channels * test.tokens * sizeof(float), 0); + ggml_tensor * attention = ggml_view_4d(ctx, gdn, width, heads, test.tokens, 1, width * sizeof(float), + channels * sizeof(float), channels * test.tokens * sizeof(float), 0); if (test.attention_output) { ggml_set_output(attention); } if (test.gdn_output) { ggml_set_output(gdn); } - ggml_tensor * new_state = ggml_view_2d(ctx, gdn, state_elements, 1, - state_elements * sizeof(float), + ggml_tensor * new_state = ggml_view_2d(ctx, gdn, state_elements, 1, state_elements * sizeof(float), channels * test.tokens * sizeof(float)); ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, state_elements, 1); ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, width); @@ -5868,13 +6808,13 @@ static void run_gdn_rmsnorm_gate_dispatch_checks() { REQUIRE(fused->bindings.size() == 11); const auto * projected = imported.graph.values().find_tensor(input); REQUIRE(projected != nullptr); - const auto * packed = plan.metadata.find_generated_resource( - projected->id, ggml::hrx::GeneratedResourceRole::F16K16Major); + const auto * packed = + plan.metadata.find_generated_resource(projected->id, ggml::hrx::GeneratedResourceRole::F16K16Major); REQUIRE(packed != nullptr && fused->bindings.back().value == packed->generated_value); require_compile_parameter(*fused, "ggml.rmsnorm_gate_f32.rms_epsilon", "1.99999999e-06"); } - const auto commands = ggml::hrx::build_command_program( - imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + const auto commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); REQUIRE(command_program_verifies(commands)); } @@ -5895,7 +6835,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(input != nullptr); ggml_tensor * output = ggml_mul_mat(ctx, weight, input); REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 128); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, + 128); } { @@ -5915,7 +6856,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(input != nullptr); ggml_tensor * output = ggml_mul_mat(ctx, weight, input); REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 2, 2048, 128); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 2, 2048, + 128); } { @@ -5925,7 +6867,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(input != nullptr); ggml_tensor * output = ggml_mul_mat(ctx, weight, input); REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 3, 2048, 128); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 3, 2048, + 128); } { @@ -5935,7 +6878,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(input != nullptr); ggml_tensor * output = ggml_mul_mat(ctx, weight, input); REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 5, 2048, 128); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 5, 2048, + 128); } { @@ -5975,7 +6919,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(input != nullptr); ggml_tensor * output = ggml_mul_mat(ctx, weight, input); REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 128); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, + 128); } { @@ -5995,7 +6940,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(input != nullptr); ggml_tensor * output = ggml_mul_mat(ctx, weight, input); REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 128); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, + 128); } { @@ -6015,7 +6961,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(input != nullptr); ggml_tensor * output = ggml_mul_mat(ctx, weight, input); REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 128); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, + 128); } { @@ -6035,7 +6982,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(input != nullptr); ggml_tensor * output = ggml_mul_mat(ctx, weight, input); REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 128); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, + 128); } { @@ -6055,7 +7003,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(input != nullptr); ggml_tensor * output = ggml_mul_mat(ctx, weight, input); REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 128); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, + 128); } { @@ -6075,7 +7024,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(input != nullptr); ggml_tensor * output = ggml_mul_mat(ctx, weight, input); REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 256); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, + 256); } { @@ -6100,8 +7050,9 @@ static void run_qwen_matmul_dispatch_checks() { ggml_tensor * output = ggml_add(ctx, projection, bias); REQUIRE(output != nullptr); schedule_fused_matmul_postops_command( - ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_bias_publish_f32", GGML_TYPE_F16, 33, 2048, 256, 2, - { GGML_OP_MUL_MAT, GGML_OP_ADD }, { "input", "weight", "bias", "output" }, false); + ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", GGML_TYPE_F16, 33, 2048, + 256, 2, { GGML_OP_MUL_MAT, GGML_OP_ADD }, + { "input", "weight", "bias", "residual_input", "residual_output" }, false); } { @@ -6116,8 +7067,9 @@ static void run_qwen_matmul_dispatch_checks() { ggml_tensor * output = ggml_add(ctx, projection, bias); REQUIRE(output != nullptr); schedule_fused_matmul_postops_command( - ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_bias_publish_f32", GGML_TYPE_BF16, 33, 2048, 256, 2, - { GGML_OP_MUL_MAT, GGML_OP_ADD }, { "input", "weight", "bias", "output" }, false); + ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", GGML_TYPE_BF16, 33, 2048, + 256, 2, { GGML_OP_MUL_MAT, GGML_OP_ADD }, + { "input", "weight", "bias", "residual_input", "residual_output" }, false); } { @@ -6132,8 +7084,9 @@ static void run_qwen_matmul_dispatch_checks() { ggml_tensor * output = ggml_add(ctx, projection, residual); REQUIRE(output != nullptr); schedule_fused_matmul_postops_command( - ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_residual_publish_f32", GGML_TYPE_F16, 33, 2048, 256, - 2, { GGML_OP_MUL_MAT, GGML_OP_ADD }, { "input", "weight", "residual_input", "residual_output" }, false); + ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", GGML_TYPE_F16, 33, 2048, + 256, 2, { GGML_OP_MUL_MAT, GGML_OP_ADD }, + { "input", "weight", "bias", "residual_input", "residual_output" }, false); } { @@ -6152,8 +7105,8 @@ static void run_qwen_matmul_dispatch_checks() { ggml_tensor * output = ggml_add(ctx, biased, residual); REQUIRE(output != nullptr); schedule_fused_matmul_postops_command( - ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", GGML_TYPE_F16, 33, 2048, 256, 3, - { GGML_OP_MUL_MAT, GGML_OP_ADD, GGML_OP_ADD }, + ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", GGML_TYPE_F16, 33, 2048, + 256, 3, { GGML_OP_MUL_MAT, GGML_OP_ADD, GGML_OP_ADD }, { "input", "weight", "bias", "residual_input", "residual_output" }, false); } @@ -6175,9 +7128,9 @@ static void run_qwen_matmul_dispatch_checks() { ggml_tensor * output = ggml_mul(ctx, rms, norm_weight); REQUIRE(output != nullptr); schedule_fused_matmul_postops_command( - ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_residual_publish_f32", GGML_TYPE_F16, 33, 2048, 256, - 4, { GGML_OP_MUL_MAT, GGML_OP_ADD, GGML_OP_RMS_NORM, GGML_OP_MUL }, - { "input", "weight", "residual_input", "residual_output" }, true); + ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", GGML_TYPE_F16, 33, 2048, + 256, 4, { GGML_OP_MUL_MAT, GGML_OP_ADD, GGML_OP_RMS_NORM, GGML_OP_MUL }, + { "input", "weight", "bias", "residual_input", "residual_output" }, true); } { @@ -6202,8 +7155,8 @@ static void run_qwen_matmul_dispatch_checks() { ggml_tensor * output = ggml_mul(ctx, rms, norm_weight); REQUIRE(output != nullptr); schedule_fused_matmul_postops_command( - ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", GGML_TYPE_F16, 33, 2048, 256, 5, - { GGML_OP_MUL_MAT, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_RMS_NORM, GGML_OP_MUL }, + ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", GGML_TYPE_F16, 33, 2048, + 256, 5, { GGML_OP_MUL_MAT, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_RMS_NORM, GGML_OP_MUL }, { "input", "weight", "bias", "residual_input", "residual_output" }, true); } @@ -6302,8 +7255,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(packed != nullptr); ggml_tensor * output = ggml_glu(ctx, packed, GGML_GLU_OP_SWIGLU, false); REQUIRE(output != nullptr); - schedule_packed_matmul_glu_command(ctx, output, GGML_TYPE_Q4_K, 2, 256, 128, - ggml::hrx::BinaryKind::SwiGLU, false); + schedule_packed_matmul_glu_command(ctx, output, GGML_TYPE_Q4_K, 2, 256, 128, ggml::hrx::BinaryKind::SwiGLU, + false); } { @@ -6315,8 +7268,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(packed != nullptr); ggml_tensor * output = ggml_glu(ctx, packed, GGML_GLU_OP_SWIGLU, true); REQUIRE(output != nullptr); - schedule_packed_matmul_glu_command(ctx, output, GGML_TYPE_Q4_K, 2, 256, 128, - ggml::hrx::BinaryKind::SwiGLU, true); + schedule_packed_matmul_glu_command(ctx, output, GGML_TYPE_Q4_K, 2, 256, 128, ggml::hrx::BinaryKind::SwiGLU, + true); } { @@ -6518,7 +7471,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(q4_input != nullptr); ggml_tensor * q4_output = ggml_mul_mat(ctx, q4_weight, q4_input); REQUIRE(q4_output != nullptr); - schedule_single_matmul_command(ctx, q4_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 128); + schedule_single_matmul_command(ctx, q4_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, + 128); ggml_tensor * q4_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 128); ggml_tensor * q4_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -6526,8 +7480,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(q4_decode_input != nullptr); ggml_tensor * q4_decode_output = ggml_mul_mat(ctx, q4_decode_weight, q4_decode_input); REQUIRE(q4_decode_output != nullptr); - schedule_single_matmul_command(ctx, q4_decode_output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 1, 2048, - 128); + schedule_single_matmul_command(ctx, q4_decode_output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 1, 2048, 128); for (ggml_type legacy_type : { GGML_TYPE_Q1_0, GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, GGML_TYPE_Q5_0, GGML_TYPE_Q5_1 }) { @@ -6537,7 +7490,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(legacy_input != nullptr); ggml_tensor * legacy_output = ggml_mul_mat(ctx, legacy_weight, legacy_input); REQUIRE(legacy_output != nullptr); - schedule_single_matmul_command(ctx, legacy_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 128); + schedule_single_matmul_command(ctx, legacy_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, + 2048, 128); ggml_tensor * legacy_decode_weight = ggml_new_tensor_2d(ctx, legacy_type, 2048, 128); ggml_tensor * legacy_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -6555,7 +7509,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(q3_input != nullptr); ggml_tensor * q3_output = ggml_mul_mat(ctx, q3_weight, q3_input); REQUIRE(q3_output != nullptr); - schedule_single_matmul_command(ctx, q3_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 128); + schedule_single_matmul_command(ctx, q3_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, + 128); ggml_tensor * q3_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q3_K, 2048, 128); ggml_tensor * q3_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -6572,7 +7527,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(iq2_s_input != nullptr); ggml_tensor * iq2_s_output = ggml_mul_mat(ctx, iq2_s_weight, iq2_s_input); REQUIRE(iq2_s_output != nullptr); - schedule_single_matmul_command(ctx, iq2_s_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 2, 2048, 640); + schedule_single_matmul_command(ctx, iq2_s_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 2, + 2048, 640); ggml_tensor * iq2_s_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ2_S, 2048, 640); ggml_tensor * iq2_s_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -6589,7 +7545,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(iq3_s_input != nullptr); ggml_tensor * iq3_s_output = ggml_mul_mat(ctx, iq3_s_weight, iq3_s_input); REQUIRE(iq3_s_output != nullptr); - schedule_single_matmul_command(ctx, iq3_s_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 128); + schedule_single_matmul_command(ctx, iq3_s_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, + 2048, 128); ggml_tensor * iq3_s_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ3_S, 2048, 128); ggml_tensor * iq3_s_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -6606,7 +7563,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(iq4_nl_input != nullptr); ggml_tensor * iq4_nl_output = ggml_mul_mat(ctx, iq4_nl_weight, iq4_nl_input); REQUIRE(iq4_nl_output != nullptr); - schedule_single_matmul_command(ctx, iq4_nl_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 128); + schedule_single_matmul_command(ctx, iq4_nl_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, + 2048, 128); ggml_tensor * iq4_nl_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ4_NL, 2048, 128); ggml_tensor * iq4_nl_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -6642,8 +7600,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(q5_0_small_input != nullptr); ggml_tensor * q5_0_small_output = ggml_mul_mat(ctx, q5_0_small_weight, q5_0_small_input); REQUIRE(q5_0_small_output != nullptr); - schedule_single_matmul_command(ctx, q5_0_small_output, "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32", - 2, 640, 256); + schedule_single_matmul_command(ctx, q5_0_small_output, "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32", 2, + 640, 256); ggml_tensor * q5_0_small_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q5_0, 640, 256); ggml_tensor * q5_0_small_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 640, 1); @@ -6660,7 +7618,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(q6_input != nullptr); ggml_tensor * q6_output = ggml_mul_mat(ctx, q6_weight, q6_input); REQUIRE(q6_output != nullptr); - schedule_single_matmul_command(ctx, q6_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 2, 2048, 128); + schedule_single_matmul_command(ctx, q6_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 2, 2048, + 128); ggml_tensor * q6_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, 2048, 128); ggml_tensor * q6_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -6686,7 +7645,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(iq4_xs_input != nullptr); ggml_tensor * iq4_xs_output = ggml_mul_mat(ctx, iq4_xs_weight, iq4_xs_input); REQUIRE(iq4_xs_output != nullptr); - schedule_single_matmul_command(ctx, iq4_xs_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 128); + schedule_single_matmul_command(ctx, iq4_xs_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, + 2048, 128); ggml_tensor * iq4_xs_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ4_XS, 2048, 128); ggml_tensor * iq4_xs_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -6694,7 +7654,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(iq4_xs_decode_input != nullptr); ggml_tensor * iq4_xs_decode_output = ggml_mul_mat(ctx, iq4_xs_decode_weight, iq4_xs_decode_input); REQUIRE(iq4_xs_decode_output != nullptr); - schedule_single_matmul_command(ctx, iq4_xs_decode_output, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 2048, 128); + schedule_single_matmul_command(ctx, iq4_xs_decode_output, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, + 2048, 128); ggml_tensor * q8_0_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q8_0, 2048, 128); ggml_tensor * q8_0_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); @@ -6702,7 +7663,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(q8_0_input != nullptr); ggml_tensor * q8_0_output = ggml_mul_mat(ctx, q8_0_weight, q8_0_input); REQUIRE(q8_0_output != nullptr); - schedule_single_matmul_command(ctx, q8_0_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 128); + schedule_single_matmul_command(ctx, q8_0_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, + 128); ggml_tensor * q8_0_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q8_0, 2048, 128); ggml_tensor * q8_0_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -6737,7 +7699,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(q8_1_input != nullptr); ggml_tensor * q8_1_output = ggml_mul_mat(ctx, q8_1_weight, q8_1_input); REQUIRE(q8_1_output != nullptr); - schedule_single_matmul_command(ctx, q8_1_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 128); + schedule_single_matmul_command(ctx, q8_1_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, + 128); ggml_tensor * q8_1_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q8_1, 2048, 128); ggml_tensor * q8_1_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -6754,7 +7717,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(f16_input != nullptr); ggml_tensor * f16_output = ggml_mul_mat(ctx, f16_weight, f16_input); REQUIRE(f16_output != nullptr); - schedule_single_matmul_command(ctx, f16_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 128); + schedule_single_matmul_command(ctx, f16_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, + 128); ggml_tensor * f16_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 2048, 128); ggml_tensor * f16_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -6771,7 +7735,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(f32_input != nullptr); ggml_tensor * f32_output = ggml_mul_mat(ctx, f32_weight, f32_input); REQUIRE(f32_output != nullptr); - schedule_single_matmul_command(ctx, f32_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 128); + schedule_single_matmul_command(ctx, f32_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, + 128); ggml_tensor * f32_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 128); ggml_tensor * f32_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -6825,7 +7790,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(input != nullptr); ggml_tensor * output = ggml_mul_mat(ctx, weight, input); REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 128); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, + 128); } { @@ -6900,7 +7866,8 @@ static void run_quantized_value_projection_dispatch_checks() { } require_compile_parameter(dispatch, "llm.attention_qkv.cache_output_format", cache_type == GGML_TYPE_F16 ? "16" : "32"); - const auto commands = ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + const auto commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); REQUIRE(command_program_verifies(commands)); } @@ -6908,6 +7875,125 @@ static void run_quantized_value_projection_dispatch_checks() { ggml_free(ctx); } +static void run_prefill_value_projection_cache_dispatch_checks() { + static constexpr const char * kDisableFusionEnv = "GGML_HRX_DISABLE_PREFILL_V_CACHE_FUSION"; + EnvironmentVariableGuard disable_fusion(kDisableFusionEnv); + disable_fusion.unset(); + + auto check_case = [&](ggml_type weight_type, int64_t input_size, int64_t output_size, int64_t token_count, + ggml_type cache_type, int64_t cache_row_count, bool publish_f16_alternate, + bool set_rows_available, bool disable, bool expect_fusion, bool observe_projection = false, + bool alias_weight = false) { + ggml_init_params params = {}; + params.mem_size = 4 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * weight_storage = ggml_new_tensor_2d(ctx, weight_type, input_size, output_size); + REQUIRE(weight_storage != nullptr); + ggml_tensor * weight = + alias_weight ? ggml_view_2d(ctx, weight_storage, input_size, output_size, weight_storage->nb[1], 0) : + weight_storage; + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, input_size, token_count); + ggml_tensor * cache = ggml_new_tensor_2d(ctx, cache_type, output_size, cache_row_count); + ggml_tensor * indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, token_count); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + REQUIRE(cache != nullptr); + REQUIRE(indices != nullptr); + ggml_tensor * projection = ggml_mul_mat(ctx, weight, input); + ggml_tensor * value = ggml_reshape_3d(ctx, projection, 128, output_size / 128, token_count); + ggml_tensor * rows = ggml_reshape_2d(ctx, value, output_size, token_count); + ggml_tensor * output = ggml_set_rows(ctx, cache, rows, indices); + REQUIRE(projection != nullptr); + REQUIRE(value != nullptr); + REQUIRE(rows != nullptr); + REQUIRE(output != nullptr); + if (observe_projection) { + ggml_set_output(projection); + } + + 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::CommandPlan plan; + const ggml::hrx::Value * input_value = imported.graph.values().find_tensor(input); + REQUIRE(input_value != nullptr); + if (publish_f16_alternate) { + const size_t bytes = static_cast(input_size * token_count) * sizeof(ggml_fp16_t); + const ggml::hrx::ValueId alternate(static_cast(imported.graph.values().size())); + plan.transients.push_back({ alternate, "test.prefill_v_cache.f16", bytes, 256 }); + plan.constant_initializations.push_back( + { alternate, "test.prefill_v_cache.f16", 0, std::vector(bytes) }); + ggml::hrx::Status status; + REQUIRE(plan.metadata.append_alternate_value( + { input_value->id, alternate, GGML_TYPE_F16, bytes, "test.prefill_v_cache.f16" }, status)); + } + + std::vector covered_nodes(imported.graph.nodes().size(), false); + if (!set_rows_available) { + covered_nodes[producer_index_for_tensor(imported.graph, output)] = true; + } + if (disable) { + disable_fusion.set("1"); + } else { + disable_fusion.unset(); + } + + ggml::hrx::DispatchMatch match; + const size_t projection_index = producer_index_for_tensor(imported.graph, projection); + REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, projection_index, match)); + REQUIRE(!match.dispatches.empty()); + const ggml::hrx::Dispatch & dispatch = match.dispatches.back(); + const bool fused = + kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:llm_attention_v_matmul_set_rows_tiled_f32_f32"; + REQUIRE(fused == expect_fusion); + if (expect_fusion) { + REQUIRE(match.dispatches.size() == 1); + REQUIRE(match.covered_nodes.size() == 4); + REQUIRE(dispatch.bindings.size() == 4); + REQUIRE(dispatch.bindings[0].value == plan.transients.front().value); + REQUIRE(dispatch.bindings[0].length == plan.transients.front().size); + require_compile_parameter(dispatch, "ggml.mul_mat.activation_format", std::to_string(GGML_TYPE_F16)); + require_compile_parameter(dispatch, "llm.attention_qkv.weight_format", "6"); + require_compile_parameter(dispatch, "llm.attention_qkv.input_size", "1536"); + require_compile_parameter(dispatch, "llm.attention_qkv.output_size", "256"); + require_compile_parameter(dispatch, "llm.attention_qkv.cache_row_count", "512"); + require_compile_parameter(dispatch, "llm.attention_qkv.cache_output_format", "16"); + + append_match_to_plan(plan, match, covered_nodes, &imported.graph); + const ggml::hrx::CommandProgram commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(command_program_verifies(commands)); + REQUIRE(commands.commands.size() == 1); + REQUIRE(commands.commands.front().bindings.size() == 4); + REQUIRE(commands.commands.front().bindings[0].name == "input"); + REQUIRE(commands.commands.front().bindings[1].name == "weight"); + REQUIRE(commands.commands.front().bindings[2].name == "indices"); + REQUIRE(commands.commands.front().bindings[3].name == "cache"); + } + ggml_free(ctx); + }; + + check_case(GGML_TYPE_Q6_K, 1536, 256, 64, GGML_TYPE_F16, 512, true, true, false, true); + check_case(GGML_TYPE_Q6_K, 1536, 256, 64, GGML_TYPE_F16, 512, true, true, true, false); + check_case(GGML_TYPE_Q6_K, 1536, 256, 64, GGML_TYPE_F16, 512, false, true, false, false); + check_case(GGML_TYPE_Q4_K, 1536, 256, 64, GGML_TYPE_F16, 512, true, true, false, false); + check_case(GGML_TYPE_Q6_K, 2048, 256, 64, GGML_TYPE_F16, 512, true, true, false, false); + check_case(GGML_TYPE_Q6_K, 1536, 512, 64, GGML_TYPE_F16, 512, true, true, false, false); + check_case(GGML_TYPE_Q6_K, 1536, 256, 32, GGML_TYPE_F16, 512, true, true, false, false); + check_case(GGML_TYPE_Q6_K, 1536, 256, 64, GGML_TYPE_F32, 512, true, true, false, false); + check_case(GGML_TYPE_Q6_K, 1536, 256, 64, GGML_TYPE_F16, 1024, true, true, false, false); + check_case(GGML_TYPE_Q6_K, 1536, 256, 64, GGML_TYPE_F16, 512, true, false, false, false); + check_case(GGML_TYPE_Q6_K, 1536, 256, 64, GGML_TYPE_F16, 512, true, true, false, false, true); + check_case(GGML_TYPE_Q6_K, 1536, 256, 64, GGML_TYPE_F16, 512, true, true, false, false, false, true); +} + static void run_llama_attention_matmul_dispatch_checks() { ggml_init_params params = {}; params.mem_size = 16 * 1024 * 1024; @@ -6941,8 +8027,9 @@ static void run_llama_attention_matmul_dispatch_checks() { REQUIRE(query != nullptr); REQUIRE(output != nullptr); schedule_fused_llama_attention_matmul_command( - ctx, output, "loom_libs:llm_attention_q_matmul_rope_tiled_f32_f32", weight_type, token_count, input_size, - output_size, head_size, head_count, 0, 0, 3, { GGML_OP_MUL_MAT, GGML_OP_RESHAPE, GGML_OP_ROPE }, + ctx, output, "loom_libs:llm_attention_q_matmul_rope_tiled_f32_f32", weight_type, token_count, + input_size, output_size, head_size, head_count, 0, 0, 3, + { GGML_OP_MUL_MAT, GGML_OP_RESHAPE, GGML_OP_ROPE }, { "input", "weight", "positions", "theta", "freq_factors", "output" }); } @@ -8000,8 +9087,7 @@ static void run_common_mul_mat_id_postops_dispatch_checks() { { const CommonMulMatIdPostOpsTensors tensors = build_common_mul_mat_id_postops_graph(ctx, true, true, true); require_common_mul_mat_id_postops_match( - ctx, tensors, common_mul_mat_id_postops_f32_kernel_name(tensors.projection->ne[2], true), true, true, - true); + ctx, tensors, common_mul_mat_id_postops_f32_kernel_name(tensors.projection->ne[2], true), true, true, true); } ggml_free(ctx); @@ -9029,8 +10115,8 @@ static void run_zero_output_scheduler_elision_checks() { REQUIRE(ids_value != nullptr); ggml::hrx::Command referenced; referenced.bindings.push_back({ "ids", ids_value->id }); - for (auto * list : { &lookup.program->commands().initialization_commands, - &lookup.program->commands().commands }) { + for (auto * list : + { &lookup.program->commands().initialization_commands, &lookup.program->commands().commands }) { list->push_back(referenced); const auto match = lookup.program->match_current_graph(*graph); REQUIRE(match.valid()); @@ -9041,8 +10127,8 @@ static void run_zero_output_scheduler_elision_checks() { const auto unreferenced = lookup.program->match_current_graph(*graph); REQUIRE(unreferenced.valid()); REQUIRE(unreferenced.external_bindings.size() == 1); - for (const auto origin : { ggml::hrx::CommandBindingOrigin::Transient, - ggml::hrx::CommandBindingOrigin::ProgramConstant }) { + for (const auto origin : + { ggml::hrx::CommandBindingOrigin::Transient, ggml::hrx::CommandBindingOrigin::ProgramConstant }) { ggml::hrx::Command unrelated; unrelated.bindings.push_back({ "negative", ggml::hrx::ValueId(-1) }); unrelated.bindings.push_back({ "large", ggml::hrx::ValueId(INT32_MAX) }); @@ -10303,6 +11389,158 @@ static void run_add_f32() { ggml_backend_free(backend); } +static std::vector run_rope_scale_numerical_case(ggml_backend_t backend, bool per_token, int mode) { + ggml_init_params params = {}; + params.mem_size = 512 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + constexpr int64_t head_size = 128; + constexpr int64_t n_dims = 96; + constexpr int64_t head_count = 2; + constexpr int64_t token_count = 3; + constexpr int64_t elements = head_size * head_count * token_count; + + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, head_size, head_count, token_count); + ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); + ggml_tensor * freqs = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_dims / 2); + ggml_tensor * rope = + ggml_rope_ext(ctx, input, pos, freqs, n_dims, mode, 0, 10000.0f, 1.0f, 0.0f, 1.19024f, 32.0f, 1.0f); + ggml_tensor * scales = per_token ? ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, 1, token_count) : nullptr; + ggml_tensor * output = per_token ? ggml_mul(ctx, rope, scales) : ggml_scale(ctx, rope, 0.0883883461f); + REQUIRE(input != nullptr); + REQUIRE(pos != nullptr); + REQUIRE(freqs != nullptr); + REQUIRE(rope != nullptr); + REQUIRE(output != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + REQUIRE(buffer != nullptr); + std::vector input_data(elements); + std::vector freq_data(n_dims / 2); + const std::vector positions = { 3, 7, 11 }; + const std::vector scale_data = { 0.125f, -0.25f, 0.5f }; + for (int64_t i = 0; i < elements; ++i) { + input_data[i] = static_cast((i * 17) % 113) * 0.03125f - 1.5f; + } + for (int64_t i = 0; i < n_dims / 2; ++i) { + freq_data[i] = 0.75f + static_cast(i % 7) * 0.125f; + } + ggml_backend_tensor_set(input, input_data.data(), 0, input_data.size() * sizeof(float)); + ggml_backend_tensor_set(pos, positions.data(), 0, positions.size() * sizeof(int32_t)); + ggml_backend_tensor_set(freqs, freq_data.data(), 0, freq_data.size() * sizeof(float)); + if (per_token) { + ggml_backend_tensor_set(scales, scale_data.data(), 0, scale_data.size() * sizeof(float)); + } + + REQUIRE(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + std::vector actual(elements); + ggml_backend_tensor_get(output, actual.data(), 0, actual.size() * sizeof(float)); + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + return actual; +} + +static void run_rope_scale_f32_numerics() { + ggml_backend_t hrx = ggml_backend_hrx_init(0); + ggml_backend_t cpu = ggml_backend_cpu_init(); + REQUIRE(hrx != nullptr); + REQUIRE(cpu != nullptr); + + for (const bool per_token : { false, true }) { + for (const int mode : { static_cast(GGML_ROPE_TYPE_NORMAL), static_cast(GGML_ROPE_TYPE_NEOX) }) { + const std::vector expected = run_rope_scale_numerical_case(cpu, per_token, mode); + const std::vector actual = run_rope_scale_numerical_case(hrx, per_token, mode); + REQUIRE(actual.size() == expected.size()); + for (size_t i = 0; i < actual.size(); ++i) { + const float tolerance = 2.0e-4f + 2.0e-4f * std::fabs(expected[i]); + REQUIRE(std::fabs(actual[i] - expected[i]) <= tolerance); + } + } + } + + ggml_backend_free(cpu); + ggml_backend_free(hrx); +} + +struct GatherAddRmsNormOutputs { + std::vector raw; + std::vector normalized; +}; + +static GatherAddRmsNormOutputs run_gather_add_rmsnorm_numerical_case(ggml_backend_t backend) { + ggml_init_params params = {}; + params.mem_size = 512 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + constexpr int64_t hidden_size = 128; + constexpr int64_t source_token_count = 4; + constexpr int64_t output_token_count = 5; + GatherAddRmsNormGraph tensors = + build_gather_add_rmsnorm_graph(ctx, hidden_size, source_token_count, output_token_count); + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + REQUIRE(buffer != nullptr); + + std::vector attention(hidden_size * source_token_count); + std::vector residual(hidden_size * source_token_count); + std::vector weight(hidden_size); + const std::vector row_ids = { 3, 1, 3, 0, 2 }; + for (size_t i = 0; i < attention.size(); ++i) { + attention[i] = static_cast((i * 13) % 47) * 0.03125f - 0.5f; + residual[i] = static_cast((i * 7) % 31) * -0.015625f + 0.25f; + } + for (size_t i = 0; i < weight.size(); ++i) { + weight[i] = 0.5f + static_cast(i % 11) * 0.0625f; + } + ggml_backend_tensor_set(tensors.attention, attention.data(), 0, attention.size() * sizeof(float)); + ggml_backend_tensor_set(tensors.residual, residual.data(), 0, residual.size() * sizeof(float)); + ggml_backend_tensor_set(tensors.row_ids, row_ids.data(), 0, row_ids.size() * sizeof(int32_t)); + ggml_backend_tensor_set(tensors.weight, weight.data(), 0, weight.size() * sizeof(float)); + + REQUIRE(ggml_backend_graph_compute(backend, tensors.graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + + GatherAddRmsNormOutputs outputs; + outputs.raw.resize(hidden_size * output_token_count); + outputs.normalized.resize(hidden_size * output_token_count); + ggml_backend_tensor_get(tensors.raw_output, outputs.raw.data(), 0, outputs.raw.size() * sizeof(float)); + ggml_backend_tensor_get(tensors.normalized_output, outputs.normalized.data(), 0, + outputs.normalized.size() * sizeof(float)); + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + return outputs; +} + +static void run_gather_add_rmsnorm_f32_numerics() { + ggml_backend_t hrx = ggml_backend_hrx_init(0); + ggml_backend_t cpu = ggml_backend_cpu_init(); + REQUIRE(hrx != nullptr); + REQUIRE(cpu != nullptr); + + const GatherAddRmsNormOutputs expected = run_gather_add_rmsnorm_numerical_case(cpu); + const GatherAddRmsNormOutputs actual = run_gather_add_rmsnorm_numerical_case(hrx); + REQUIRE(actual.raw.size() == expected.raw.size()); + REQUIRE(actual.normalized.size() == expected.normalized.size()); + for (size_t i = 0; i < actual.raw.size(); ++i) { + REQUIRE(actual.raw[i] == expected.raw[i]); + const float tolerance = 3.0e-4f + 3.0e-4f * std::fabs(expected.normalized[i]); + REQUIRE(std::fabs(actual.normalized[i] - expected.normalized[i]) <= tolerance); + } + + ggml_backend_free(cpu); + ggml_backend_free(hrx); +} + static void run_scale_f32() { ggml_backend_t backend = ggml_backend_hrx_init(0); REQUIRE(backend != nullptr); @@ -10748,6 +11986,7 @@ static void run_lowtoken_residual_dispatch_checks() { bool second_add; bool fused; }; + const Case cases[] = { { GGML_TYPE_Q4_K, 6144, 5120, 1, 0, 0, false, false, true }, { GGML_TYPE_Q4_K, 17408, 5120, 1, 0, 0, false, false, true }, @@ -10820,23 +12059,22 @@ static void run_lowtoken_residual_dispatch_checks() { for (const auto & command : program.commands) { const std::string name = kernel_name_for_id(command.kernel.kernel_id); binary += name == "loom_libs:ggml_binary_f32"; - const bool tiled_postops = - name == "loom_libs:ggml_mul_mat_tiled_input_f32_bias_publish_f32" || - name == "loom_libs:ggml_mul_mat_tiled_input_f32_residual_publish_f32" || - name == "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32"; - const bool skinny_postops = - name == "loom_libs:ggml_mul_mat_skinny_input_f32_bias_publish_f32" || - name == "loom_libs:ggml_mul_mat_skinny_input_f32_residual_publish_f32" || - name == "loom_libs:ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32"; - const bool vector_postops = name == "loom_libs:ggml_mul_mat_vector_bias_f32_f32" || - name == "loom_libs:ggml_mul_mat_vector_residual_f32_f32" || - name == "loom_libs:ggml_mul_mat_vector_bias_residual_f32_f32"; + const bool tiled_postops = name == "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32"; + const bool skinny_postops = name == "loom_libs:ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32"; + const bool vector_postops = name == "loom_libs:ggml_mul_mat_vector_bias_residual_f32_f32"; if (!tiled_postops && !skinny_postops && !vector_postops) { continue; } ++fused; - fused_second_add = fused_second_add || string_contains(name, "bias_residual") || - string_contains(name, "bias_add"); + if (vector_postops) { + fused_second_add = + fused_second_add || + (command.kernel.compile_parameters.at("ggml.matmul.vector_postops.apply_bias") == "1" && + command.kernel.compile_parameters.at("ggml.matmul.vector_postops.apply_residual") == "1"); + } else { + fused_second_add = + fused_second_add || command.kernel.compile_parameters.at("ggml.mul_mat_postops.epilogue") == "3"; + } REQUIRE(command.bindings.size() >= 4); REQUIRE(command.kernel.integer_parameters.at("token_count") == c.tokens); if (vector_postops) { @@ -10878,8 +12116,8 @@ static void run_lowtoken_bias_dispatch_checks() { size_t tiled_bias = 0; for (const auto & command : program.commands) { const std::string name = kernel_name_for_id(command.kernel.kernel_id); - skinny_bias += name == "loom_libs:ggml_mul_mat_skinny_input_f32_bias_publish_f32"; - tiled_bias += name == "loom_libs:ggml_mul_mat_tiled_input_f32_bias_publish_f32"; + skinny_bias += name == "loom_libs:ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32"; + tiled_bias += name == "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32"; } REQUIRE(skinny_bias == 1); REQUIRE(tiled_bias == 0); @@ -10893,12 +12131,23 @@ static void run_gdn_selected_rms_q8_dispatch_checks() { params.no_alloc = true; ggml_context * ctx = ggml_init(params); REQUIRE(ctx != nullptr); + enum { - UnreadyInput = 1, UnreadyGate = 2, GatherSide = 4, AttentionSide = 8, - GdnObserved = 16, GatherObserved = 32, Unaligned = 64, CacheRead = 128, - OutputSide = 256, CachedGate = 512, TanhGate = 1024, AttentionReshape = 2048, + UnreadyInput = 1, + UnreadyGate = 2, + GatherSide = 4, + AttentionSide = 8, + GdnObserved = 16, + GatherObserved = 32, + Unaligned = 64, + CacheRead = 128, + OutputSide = 256, + CachedGate = 512, + TanhGate = 1024, + AttentionReshape = 2048, AttentionObserved = 4096, }; + const struct { int64_t tokens, heads, sequences; int flags; @@ -10956,8 +12205,8 @@ static void run_gdn_selected_rms_q8_dispatch_checks() { } ggml_tensor * state = ggml_reshape_4d(ctx, gathered, width, width, heads, test.sequences); const auto view = [&](int64_t nheads, size_t offset) { - return ggml_view_4d(ctx, raw, width, nheads, test.tokens, test.sequences, - width * sizeof(float), hidden * sizeof(float), hidden * test.tokens * sizeof(float), offset); + return ggml_view_4d(ctx, raw, width, nheads, test.tokens, test.sequences, width * sizeof(float), + hidden * sizeof(float), hidden * test.tokens * sizeof(float), offset); }; ggml_tensor * q = ggml_l2_norm(ctx, view(qheads, 0), 1.e-6f); ggml_tensor * k = ggml_l2_norm(ctx, view(qheads, width * qheads * sizeof(float)), 1.e-6f); @@ -10967,26 +12216,32 @@ static void run_gdn_selected_rms_q8_dispatch_checks() { ggml_set_output(gdn); } const size_t attention_bytes = channels * test.tokens * test.sequences * sizeof(float); - ggml_tensor * attention = ggml_view_4d(ctx, gdn, width, heads, test.tokens, test.sequences, - width * sizeof(float), channels * sizeof(float), channels * test.tokens * sizeof(float), 0); + ggml_tensor * attention = + ggml_view_4d(ctx, gdn, width, heads, test.tokens, test.sequences, width * sizeof(float), + channels * sizeof(float), channels * test.tokens * sizeof(float), 0); if (test.flags & AttentionObserved) { ggml_set_output(attention); } - ggml_tensor * snapshots = ggml_view_3d(ctx, gdn, state_elements, test.sequences, written, - state_bytes, state_bytes * test.sequences, attention_bytes); - ggml_tensor * target = ggml_view_3d(ctx, cache, state_elements, test.sequences, written, - state_bytes, 4 * state_bytes, (test.flags & Unaligned) ? sizeof(float) : state_bytes); + ggml_tensor * snapshots = ggml_view_3d(ctx, gdn, state_elements, test.sequences, written, state_bytes, + state_bytes * test.sequences, attention_bytes); + ggml_tensor * target = ggml_view_3d(ctx, cache, state_elements, test.sequences, written, state_bytes, + 4 * state_bytes, (test.flags & Unaligned) ? sizeof(float) : state_bytes); ggml_tensor * norm_input = (test.flags & AttentionReshape) ? - ggml_reshape_2d(ctx, attention, width, heads * test.tokens * test.sequences) : attention; + ggml_reshape_2d(ctx, attention, width, heads * test.tokens * test.sequences) : + attention; ggml_tensor * rms_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, width); - ggml_tensor * raw_gate = (test.flags & CachedGate) ? + ggml_tensor * raw_gate = + (test.flags & CachedGate) ? ggml_view_4d(ctx, cache, width, heads, test.tokens, test.sequences, width * sizeof(float), channels * sizeof(float), channels * test.tokens * sizeof(float), 0) : - ggml_scale(ctx, ggml_new_tensor_4d(ctx, GGML_TYPE_F32, width, heads, test.tokens, test.sequences), 0.5f); + ggml_scale(ctx, ggml_new_tensor_4d(ctx, GGML_TYPE_F32, width, heads, test.tokens, test.sequences), + 0.5f); ggml_tensor * shaped_gate = (test.flags & AttentionReshape) ? - ggml_reshape_2d(ctx, raw_gate, width, heads * test.tokens * test.sequences) : raw_gate; + ggml_reshape_2d(ctx, raw_gate, width, heads * test.tokens * test.sequences) : + raw_gate; ggml_tensor * norm = ggml_mul(ctx, ggml_rms_norm(ctx, norm_input, 1.e-6f), rms_weight); - ggml_tensor * gated = ggml_mul(ctx, norm, (test.flags & TanhGate) ? ggml_tanh(ctx, shaped_gate) : ggml_silu(ctx, shaped_gate)); + ggml_tensor * gated = + ggml_mul(ctx, norm, (test.flags & TanhGate) ? ggml_tanh(ctx, shaped_gate) : ggml_silu(ctx, shaped_gate)); ggml_tensor * projection_input = ggml_reshape_2d(ctx, gated, channels, test.tokens * test.sequences); ggml_tensor * projection_weight = ggml_new_tensor_2d(ctx, test.type, channels, 4096); ggml_tensor * projection = ggml_mul_mat(ctx, projection_weight, projection_input); @@ -11040,7 +12295,8 @@ static void run_gdn_selected_rms_q8_dispatch_checks() { name == "loom_libs:ggml_rmsnorm_gate_f32_q8_1_x4"; })); } - const auto commands = ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + const auto commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); if (!(test.flags & UnreadyGate)) { REQUIRE(command_program_verifies(commands)); @@ -11057,6 +12313,8 @@ static void register_hrx_backend_host_cases(test_runner::Suite & suite) { suite.host_case("graph_import_mixed_backend_boundary", [] { run_graph_import_mixed_backend_boundary_checks(); }); suite.host_case("graph_view_external_use", [] { run_graph_view_external_use_checks(); }); suite.host_case("scale_f32_dispatch", [] { run_scale_f32_dispatch_checks(); }); + suite.host_case("scale_add_f32_dispatch", [] { run_scale_add_f32_dispatch_checks(); }); + suite.host_case("rmsnorm_two_binary_dispatch", [] { run_rmsnorm_two_binary_dispatch_checks(); }); suite.host_case("cont_f32_dispatch", [] { run_cont_f32_dispatch_checks(); }); suite.host_case("binary_f32_broadcast_dispatch", [] { run_binary_f32_broadcast_dispatch_checks(); }); suite.host_case("graph_snapshot_diagnostics", [] { run_graph_snapshot_diagnostics_checks(); }); @@ -11067,6 +12325,7 @@ static void register_hrx_backend_host_cases(test_runner::Suite & suite) { 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("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(); }); suite.host_case("qwen_flash_attention_dispatch", [] { run_qwen_flash_attention_dispatch_checks(); }); suite.host_case("tiled_pair_matmul_postops_dispatch", [] { run_tiled_pair_matmul_postops_dispatch_checks(); }); @@ -11083,18 +12342,24 @@ static void register_hrx_backend_host_cases(test_runner::Suite & suite) { suite.host_case("rmsnorm_gate_q8_output", [] { run_rmsnorm_gate_q8_output_checks(); }); suite.host_case("lowtoken_residual_dispatch", [] { run_lowtoken_residual_dispatch_checks(); }); suite.host_case("lowtoken_bias_dispatch", [] { run_lowtoken_bias_dispatch_checks(); }); + suite.host_case("generic_ssm_conv_binary_dispatch", [] { run_generic_ssm_conv_binary_dispatch_checks(); }); suite.host_case("quantized_conv4_dispatch", [] { run_quantized_conv4_dispatch_checks(); }); + suite.host_case("lfm_dconv3_dispatch", [] { run_lfm_dconv3_dispatch_checks(); }); suite.host_case("gdn_rmsnorm_gate_dispatch", [] { run_gdn_rmsnorm_gate_dispatch_checks(); }); suite.host_case("gdn_native_projection_pair_dispatch", [] { run_gdn_native_projection_pair_dispatch_checks(); }); suite.host_case("gdn_selected_snapshot_dispatch", [] { run_gdn_selected_snapshot_dispatch_checks(); }); suite.host_case("gdn_selected_rms_q8_dispatch", [] { run_gdn_selected_rms_q8_dispatch_checks(); }); suite.host_case("llama_attention_matmul_dispatch", [] { run_llama_attention_matmul_dispatch_checks(); }); suite.host_case("quantized_value_projection_dispatch", [] { run_quantized_value_projection_dispatch_checks(); }); + suite.host_case("prefill_value_projection_cache_dispatch", + [] { run_prefill_value_projection_cache_dispatch_checks(); }); suite.host_case("qwen_terminal_q6k_q8_command.tokens1", [] { schedule_qwen_terminal_q6k_q8_command(1); }); suite.host_case("qwen_terminal_q6k_q8_command.tokens18", [] { schedule_qwen_terminal_q6k_q8_command(18); }); suite.host_case("qwen_terminal_q6k_vector_command", [] { schedule_qwen_terminal_q6k_vector_command(); }); - suite.host_case("get_rows_q8_1_alternate_command.q4_k", [] { schedule_get_rows_q8_1_alternate_command(GGML_TYPE_Q4_K); }); - suite.host_case("get_rows_q8_1_alternate_command.bf16", [] { schedule_get_rows_q8_1_alternate_command(GGML_TYPE_BF16); }); + suite.host_case("get_rows_q8_1_alternate_command.q4_k", + [] { schedule_get_rows_q8_1_alternate_command(GGML_TYPE_Q4_K); }); + suite.host_case("get_rows_q8_1_alternate_command.bf16", + [] { schedule_get_rows_q8_1_alternate_command(GGML_TYPE_BF16); }); suite.host_case("qwen_router_top8_dispatch", [] { run_qwen_router_top8_dispatch_checks(); }); suite.host_case("common_mul_mat_id_swiglu_dispatch", [] { run_common_mul_mat_id_swiglu_dispatch_checks(); }); suite.host_case("qwen_routed_gate_up_dispatch", [] { run_qwen_routed_gate_up_dispatch_checks(); }); @@ -11104,9 +12369,11 @@ static void register_hrx_backend_host_cases(test_runner::Suite & suite) { suite.host_case("layout_alias_scheduler_elision", [] { run_layout_alias_scheduler_elision_checks(); }); suite.host_case("zero_output_scheduler_elision", [] { run_zero_output_scheduler_elision_checks(); }); suite.host_case("transient_import", [] { run_transient_import_checks(); }); - suite.host_case("graph_view_preserves_shared_output_storage", [] { run_graph_view_preserves_shared_output_storage_checks(); }); + suite.host_case("graph_view_preserves_shared_output_storage", + [] { run_graph_view_preserves_shared_output_storage_checks(); }); suite.host_case("chained_dispatch_requires_transients", [] { run_chained_dispatch_requires_transients(); }); - suite.host_case("graph_replay_host_staging_is_not_ineligible", [] { run_graph_replay_host_staging_is_not_ineligible(); }); + suite.host_case("graph_replay_host_staging_is_not_ineligible", + [] { run_graph_replay_host_staging_is_not_ineligible(); }); suite.host_case("multiple_transient_plan", [] { run_multiple_transient_plan_checks(); }); suite.host_case("disjoint_transient_plan_packing", [] { run_disjoint_transient_plan_packing_checks(); }); suite.host_case("command_program_kernel_dump", [] { run_command_program_kernel_dump_checks(); }); @@ -11124,14 +12391,19 @@ static void register_hrx_backend_host_cases(test_runner::Suite & suite) { static void register_hrx_backend_device_cases(test_runner::Suite & suite) { suite.device_case("zero_output_device_support", [] { run_zero_output_device_support_checks(); }); suite.device_case("scale_f32_device_support", [] { run_scale_f32_device_support_checks(); }); - suite.device_case("qwen_expert_table_partition_prefill_512_execution", [] { run_qwen_expert_table_partition_prefill_512_execution(); }); + suite.device_case("qwen_expert_table_partition_prefill_512_execution", + [] { run_qwen_expert_table_partition_prefill_512_execution(); }); suite.device_case("add_f32", [] { run_add_f32(); }); + suite.device_case("rope_scale_f32_numerics", [] { run_rope_scale_f32_numerics(); }); + suite.device_case("gather_add_rmsnorm_f32_numerics", [] { run_gather_add_rmsnorm_f32_numerics(); }); suite.device_case("scale_f32", [] { run_scale_f32(); }); suite.device_case("scale_f32_inplace", [] { run_scale_f32_inplace(); }); suite.device_case("two_independent_add_f32", [] { run_two_independent_add_f32(); }); suite.device_case("chained_add_f32", [] { run_chained_add_f32(); }); - suite.device_case("same_uid_distinct_graph_reuses_graph_program.default", [] { run_same_uid_distinct_graph_reuses_graph_program(); }); - suite.device_case("same_uid_distinct_graph_reuses_graph_program.distinct_uid", [] { run_same_uid_distinct_graph_reuses_graph_program(true); }); + suite.device_case("same_uid_distinct_graph_reuses_graph_program.default", + [] { run_same_uid_distinct_graph_reuses_graph_program(); }); + suite.device_case("same_uid_distinct_graph_reuses_graph_program.distinct_uid", + [] { run_same_uid_distinct_graph_reuses_graph_program(true); }); suite.device_case("unsupported_op_fails", [] { run_unsupported_op_fails(); }); } @@ -11141,8 +12413,8 @@ static void register_hrx_backend_cases(test_runner::Suite & suite) { } int main(int argc, char ** argv) { - test_runner::Suite suite(test_runner::Config::with_prefix( - "HRX backend test", "hrx-backend", "GGML_HRX_BACKEND_TEST")); + test_runner::Suite suite( + test_runner::Config::with_prefix("HRX backend test", "hrx-backend", "GGML_HRX_BACKEND_TEST")); register_hrx_backend_cases(suite); const bool has_device = ggml_backend_hrx_get_device_count() != 0; diff --git a/tests/test-hrx-loom-jit.cpp b/tests/test-hrx-loom-jit.cpp index 0c95d654cbf7..ed3391a19e47 100644 --- a/tests/test-hrx-loom-jit.cpp +++ b/tests/test-hrx-loom-jit.cpp @@ -204,6 +204,18 @@ static std::map token_count_workload(int64_t token_count) }; } +static std::map moe_routing_config() { + return { + { "ggml.moe_routing.route_count", "8" }, + { "ggml.moe_routing.expert_count", "128" }, + { "ggml.moe_routing.descriptor_expert_mask", "127" }, + { "ggml.moe_routing.descriptor_partition_shift", "7" }, + { "ggml.moe_routing.descriptor_row_count_shift", "13" }, + { "ggml.moe_routing.partition_workgroup_size", "128" }, + { "ggml.workload.token_capacity", "1" }, + }; +} + static std::map flash_attention_workload(int64_t query_token_count, int64_t key_value_token_count) { return { @@ -455,8 +467,8 @@ static void run_loom_jit_compile_checks() { const ggml::hrx::KernelDefinition & gather_add = find_kernel("ggml_gather_add_f32"); const ggml::hrx::KernelDefinition & rmsnorm = find_kernel("qwen3_moe_rmsnorm_f32"); const ggml::hrx::KernelDefinition & router_top8 = find_kernel("qwen3_moe_router_top8_f32"); - const ggml::hrx::KernelDefinition & expert_table = find_kernel("qwen3_moe_build_expert_table"); - const ggml::hrx::KernelDefinition & partition_table = find_kernel("qwen3_moe_build_expert_partition_table"); + const ggml::hrx::KernelDefinition & expert_table = find_kernel("ggml_moe_build_expert_table"); + const ggml::hrx::KernelDefinition & partition_table = find_kernel("ggml_moe_build_expert_partition_table"); const ggml::hrx::KernelDefinition & mul_mat_id_f16 = find_kernel("ggml_mul_mat_id_f16_f16_wmma"); const ggml::hrx::KernelDefinition & swiglu_f16 = find_kernel("ggml_mul_mat_id_swiglu_f16_f16_wmma"); const ggml::hrx::KernelDefinition & flash_prefill = find_kernel("ggml_flash_attention_f32_f16_wmma"); @@ -532,28 +544,23 @@ static void run_loom_jit_compile_checks() { { "route_stride", 8 }, { "expert_count", 128 }, }, - { - { "qwen3_moe.routed_gate_up.expert_count", "128" }, - { "qwen3_moe.routed_gate_up.route_count", "8" }, - { "qwen3_moe.workload.token_capacity", "1" }, - })); + moe_routing_config())); refs.push_back(compile_kernel(*async_jit, "async-partition-table-1", partition_table, { { "token_count", 1 }, { "route_count", 8 }, { "expert_count", 128 }, }, - { - { "qwen3_moe.routed_gate_up.expert_count", "128" }, - { "qwen3_moe.routed_gate_up.route_count", "8" }, - { "qwen3_moe.workload.token_capacity", "1" }, - })); + moe_routing_config())); // Exercise gather-add coverage in the async JIT path. refs.push_back(compile_kernel(*async_jit, "async-gather-add-2-to-1", gather_add, { { "source_token_count", 2 }, { "output_token_count", 1 }, { "hidden_size", 2048 }, + }, + { + { "ggml.gather_add_f32.binary_op", "3" }, })); refs.push_back(compile_kernel(*async_jit, "async-mul-mat-id-f16-q4", mul_mat_id_f16, token_count_workload(4), mul_mat_id_f16_f16_config("4"))); @@ -609,6 +616,7 @@ static void run_loom_jit_compile_checks() { std::printf("async Loom JIT compiled %zu kernels\n", refs.size()); const auto & q6_prefill = find_kernel("ggml_mul_mat_q6_k_f16_wmma_prefill_wave32"); + const struct { int64_t tokens; int64_t inputs; @@ -635,11 +643,13 @@ static void run_loom_jit_compile_checks() { {1024, 5120, 1024, 0, 512, 2, 16 }, {2048, 5120, 1024, 0, 512, 4, 16 }, }; + for (const auto & test : q6_launch_cases) { - const std::string key = "q6-prefill-" + std::to_string(test.tokens) + "-" + - std::to_string(test.inputs) + "-" + std::to_string(test.outputs) + "-" + - std::to_string(test.accumulate); - auto ref = compile_kernel(*sync_jit, key, q6_prefill, token_count_workload(test.tokens), { + const std::string key = "q6-prefill-" + std::to_string(test.tokens) + "-" + std::to_string(test.inputs) + "-" + + std::to_string(test.outputs) + "-" + std::to_string(test.accumulate); + auto ref = + compile_kernel(*sync_jit, key, q6_prefill, token_count_workload(test.tokens), + { { "ggml.mul_mat_q6_k_packed.input_size", std::to_string(test.inputs) }, { "ggml.mul_mat_q6_k_packed.output_size", std::to_string(test.outputs) }, { "ggml.mul_mat_q6_k_packed.output_accumulation", std::to_string(test.accumulate) }, @@ -657,6 +667,7 @@ static void run_loom_jit_compile_checks() { } const auto & q4_prefill = find_kernel("ggml_mul_mat_q4_k_f16_wmma_prefill_wave32"); + const struct { int64_t tokens; int64_t inputs; @@ -682,11 +693,12 @@ static void run_loom_jit_compile_checks() { {1024, 5120, 4096, 1024, 4, 16, 0 }, { 512, 17408, 5120, 512, 1, 80, 1 }, }; + for (const auto & test : q4_launch_cases) { - const std::string key = "q4-prefill-" + std::to_string(test.tokens) + "-" + - std::to_string(test.inputs) + "-" + std::to_string(test.outputs) + "-" + - std::to_string(test.input_layout); - auto ref = compile_kernel(*sync_jit, key, q4_prefill, token_count_workload(test.tokens), { + const std::string key = "q4-prefill-" + std::to_string(test.tokens) + "-" + std::to_string(test.inputs) + "-" + + std::to_string(test.outputs) + "-" + std::to_string(test.input_layout); + auto ref = compile_kernel(*sync_jit, key, q4_prefill, token_count_workload(test.tokens), + { { "ggml.mul_mat.input_size", std::to_string(test.inputs) }, { "ggml.mul_mat.output_size", std::to_string(test.outputs) }, { "ggml.workload.token_capacity", std::to_string(test.tokens) }, @@ -704,8 +716,9 @@ static void run_loom_jit_compile_checks() { const auto & packed_f16 = find_kernel("ggml_copy_f16_k16_major"); for (const auto * input_is_f16 : { "0", "1" }) { - auto ref = compile_kernel(*sync_jit, std::string("packed-f16-") + input_is_f16, - packed_f16, token_count_workload(512), { + auto ref = + compile_kernel(*sync_jit, std::string("packed-f16-") + input_is_f16, packed_f16, token_count_workload(512), + { { "ggml.copy_f16_k16_major.input_size", "256" }, { "ggml.copy_f16_k16_major.token_count", "512" }, { "ggml.copy_f16_k16_major.input_is_f16", input_is_f16 }, @@ -714,19 +727,27 @@ static void run_loom_jit_compile_checks() { } const auto & transpose_f16 = find_kernel("ggml_copy_transpose_f16"); - auto transpose_ref = compile_kernel(*sync_jit, "transpose-f16-8192x1024", transpose_f16, {}, { + auto transpose_ref = compile_kernel(*sync_jit, "transpose-f16-8192x1024", transpose_f16, + { + }, + { { "ggml.copy_transpose_f16.row_count", "8192" }, { "ggml.copy_transpose_f16.column_count", "1024" }, }); require_compiled_kernel(transpose_ref); auto transposed_attention_config = flash_attention_config("24", "4", "256", "0.0625"); transposed_attention_config.emplace("ggml.flash_attention.value_layout", "1"); - auto transposed_attention_ref = compile_kernel(*sync_jit, "flash-transposed-value", flash_prefill, - { { "query_token_count", 512 }, { "key_value_token_count", 8192 } }, + auto transposed_attention_ref = + compile_kernel(*sync_jit, "flash-transposed-value", flash_prefill, + { + { "query_token_count", 512 }, + { "key_value_token_count", 8192 } + }, transposed_attention_config); require_compiled_kernel(transposed_attention_ref); const auto & paired_q4 = find_kernel("ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32"); + const struct { int64_t tokens; int64_t outputs; @@ -740,9 +761,11 @@ static void run_loom_jit_compile_checks() { {1024, 2048, 512, 2 }, {2048,17408, 512, 4 }, }; + for (const auto & test : paired_q4_cases) { const std::string key = "paired-q4-" + std::to_string(test.tokens) + "-" + std::to_string(test.outputs); - auto ref = compile_kernel(*sync_jit, key, paired_q4, token_count_workload(test.tokens), { + auto ref = compile_kernel(*sync_jit, key, paired_q4, token_count_workload(test.tokens), + { { "ggml.mul_mat_swiglu.input_size", "5120" }, { "ggml.mul_mat_swiglu.output_size", std::to_string(test.outputs) }, { "ggml.mul_mat_swiglu.op", "4" }, @@ -761,8 +784,9 @@ static void run_loom_jit_compile_checks() { } for (const int64_t tokens : { 256, 512, 1024, 2048 }) { - auto ref = compile_kernel(*sync_jit, "paired-q4-packed-output-" + std::to_string(tokens), - paired_q4, token_count_workload(tokens), { + auto ref = compile_kernel(*sync_jit, "paired-q4-packed-output-" + std::to_string(tokens), paired_q4, + token_count_workload(tokens), + { { "ggml.mul_mat_swiglu.input_size", "5120" }, { "ggml.mul_mat_swiglu.output_size", "17408" }, { "ggml.mul_mat_swiglu.op", "4" }, @@ -772,8 +796,8 @@ static void run_loom_jit_compile_checks() { require_compiled_kernel(ref); } - const auto & paired_q4_decode = find_kernel("ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot"); + const struct { int64_t tokens; int64_t inputs; @@ -801,10 +825,12 @@ static void run_loom_jit_compile_checks() { { 5, 5632, 4096, 512, 128 }, { 5, 5888, 4096, 128, 512 }, }; + for (const auto & test : paired_q4_decode_cases) { - const std::string key = "paired-q4-decode-" + std::to_string(test.tokens) + "-" + - std::to_string(test.inputs) + "-" + std::to_string(test.outputs); - auto ref = compile_kernel(*sync_jit, key, paired_q4_decode, token_count_workload(test.tokens), { + const std::string key = "paired-q4-decode-" + std::to_string(test.tokens) + "-" + std::to_string(test.inputs) + + "-" + std::to_string(test.outputs); + auto ref = compile_kernel(*sync_jit, key, paired_q4_decode, token_count_workload(test.tokens), + { { "ggml.mul_mat_swiglu.input_size", std::to_string(test.inputs) }, { "ggml.mul_mat_swiglu.output_size", std::to_string(test.outputs) }, { "ggml.mul_mat_swiglu.op", "4" }, @@ -820,9 +846,8 @@ static void run_loom_jit_compile_checks() { compiled.reset(); } - - const auto & quant_decode = find_kernel("ggml_mul_mat_f32_f32_wmma"); + const struct { int64_t tokens; int64_t inputs; @@ -878,11 +903,13 @@ static void run_loom_jit_compile_checks() { { 5, 8192, 4032, 44, 128, 504 }, { 5, 8192, 8256, 44, 128, 1032 }, }; + for (const auto & test : quant_decode_cases) { - const std::string key = "quant-decode-" + std::to_string(test.tokens) + "-" + - std::to_string(test.inputs) + "-" + std::to_string(test.outputs) + "-" + - std::to_string(test.format); - auto ref = compile_kernel(*sync_jit, key, quant_decode, token_count_workload(test.tokens), { + const std::string key = "quant-decode-" + std::to_string(test.tokens) + "-" + std::to_string(test.inputs) + + "-" + std::to_string(test.outputs) + "-" + std::to_string(test.format); + auto ref = + compile_kernel(*sync_jit, key, quant_decode, token_count_workload(test.tokens), + { { "ggml.mul_mat.input_size", std::to_string(test.inputs) }, { "ggml.mul_mat.output_size", std::to_string(test.outputs) }, { "ggml.mul_mat.weight_format", std::to_string(test.format) }, @@ -929,7 +956,10 @@ static void run_loom_jit_compile_checks() { const auto & conv_finish = find_kernel("llm_ssm_conv_dconv4_silu_prefill_finish_f32"); REQUIRE(conv_finish.bindings.size() == 5); for (const int64_t channels : { 2112, 8192, 8256, 10240 }) { - auto ref = compile_kernel(*sync_jit, "conv-finish-" + std::to_string(channels), conv_finish, {}, { + auto ref = compile_kernel(*sync_jit, "conv-finish-" + std::to_string(channels), conv_finish, + { + }, + { { "llm.ssm_conv.generic.d_inner", std::to_string(channels) }, }); REQUIRE(ref != nullptr); @@ -942,7 +972,10 @@ static void run_loom_jit_compile_checks() { const auto & ssm_prefill = find_kernel("llm_ssm_conv_dconv4_silu_prefill_512_wg1024"); for (const int64_t channels : { 8192, 8224, 8256, 10208, 10240 }) { - auto ref = compile_kernel(*sync_jit, "ssm-prefill-" + std::to_string(channels), ssm_prefill, {}, { + auto ref = compile_kernel(*sync_jit, "ssm-prefill-" + std::to_string(channels), ssm_prefill, + { + }, + { { "llm.ssm_conv.prefill.d_conv", "4" }, { "llm.ssm_conv.prefill.d_inner", std::to_string(channels) }, { "llm.ssm_conv.prefill.n_t", "512" }, @@ -973,10 +1006,14 @@ static void run_loom_jit_compile_checks() { { "llm_gated_delta_net_f32_wmma_head128_snapshot", 5, 8 }, { "llm_gated_delta_net_f32_wmma_head128_snapshot_projection_epilogue", 5, 10 }, }; + for (const auto & test : gdn_cases) { const auto & definition = find_kernel(test.name); REQUIRE(definition.bindings.size() == test.bindings); - auto ref = compile_kernel(*sync_jit, test.name, definition, {}, { + auto ref = compile_kernel(*sync_jit, test.name, definition, + { + }, + { { "llm.gated_delta_net.head_width", "128" }, { "llm.gated_delta_net.head_count", "3" }, { "llm.gated_delta_net.token_count", std::to_string(test.tokens) }, @@ -1006,41 +1043,74 @@ static void run_loom_jit_compile_checks() { REQUIRE(compiled.launch_config.workgroup_count[1] == 1); compiled.reset(); } +} +static void run_moe_routing_compile_checks() { + static constexpr const char * kTarget = "gfx1100"; + + std::string error; + std::unique_ptr jit = ggml::hrx::create_loom_jit(kTarget, ggml::hrx::LoomJitMode::Sync, error); + REQUIRE(jit != nullptr); + + const std::map expert_table_workload = { + { "token_count", 1 }, + { "route_count", 8 }, + { "route_stride", 8 }, + { "expert_count", 128 }, + }; + const auto & expert_table = find_kernel("ggml_moe_build_expert_table"); + require_compiled_kernel( + compile_kernel(*jit, "moe-expert-table", expert_table, expert_table_workload, moe_routing_config())); + + const std::map partition_table_workload = { + { "token_count", 1 }, + { "route_count", 8 }, + { "expert_count", 128 }, + }; + const auto & partition_table = find_kernel("ggml_moe_build_expert_partition_table"); + require_compiled_kernel(compile_kernel(*jit, "moe-expert-partition-table", partition_table, + partition_table_workload, moe_routing_config())); } static void run_prefill_fusion_compile_checks() { static constexpr const char * kTarget = "gfx1100"; std::string error; - std::unique_ptr jit = - ggml::hrx::create_loom_jit(kTarget, ggml::hrx::LoomJitMode::Sync, error); + std::unique_ptr jit = ggml::hrx::create_loom_jit(kTarget, ggml::hrx::LoomJitMode::Sync, error); REQUIRE(jit != nullptr); const auto & rmsnorm_f32_f16 = find_kernel("ggml_rmsnorm_binary_f32_f16"); - require_compiled_kernel(compile_kernel(*jit, "prefill-rmsnorm-f32-f16", rmsnorm_f32_f16, - token_count_workload(512), { + require_compiled_kernel(compile_kernel(*jit, "prefill-rmsnorm-f32-f16", rmsnorm_f32_f16, token_count_workload(512), + { { "ggml.rmsnorm_binary_f32.hidden_size", "2048" }, { "ggml.rmsnorm_binary_f32.rms_epsilon", "0.00001" }, { "ggml.rmsnorm_binary_f32.op", "2" }, })); + const auto & rmsnorm_strided = find_kernel("ggml_rmsnorm_binary_strided_f32"); + require_compiled_kernel(compile_kernel(*jit, "prefill-rmsnorm-strided", rmsnorm_strided, token_count_workload(64), + { + { "ggml.rmsnorm_binary_f32.hidden_size", "256" }, + { "ggml.rmsnorm_binary_f32.input_stride", "288" }, + { "ggml.rmsnorm_binary_f32.rms_epsilon", "0.00001" }, + { "ggml.rmsnorm_binary_f32.op", "2" }, + })); + const auto & rmsnorm_q8_f16 = find_kernel("ggml_rmsnorm_binary_q8_1_x4_f16"); - require_compiled_kernel(compile_kernel(*jit, "prefill-rmsnorm-q8-f16", rmsnorm_q8_f16, - token_count_workload(512), { + require_compiled_kernel(compile_kernel(*jit, "prefill-rmsnorm-q8-f16", rmsnorm_q8_f16, token_count_workload(512), + { { "ggml.rmsnorm_binary_q8_1_x4.hidden_size", "2048" }, { "ggml.rmsnorm_binary_q8_1_x4.rms_epsilon", "0.00001" }, { "ggml.rmsnorm_binary_q8_1_x4.op", "2" }, })); const auto & flash_f16 = find_kernel("ggml_flash_attention_f32_f16_wmma_publish_f16"); - require_compiled_kernel(compile_kernel(*jit, "prefill-flash-f16", flash_f16, - flash_attention_workload(512, 512), + require_compiled_kernel(compile_kernel(*jit, "prefill-flash-f16", flash_f16, flash_attention_workload(512, 512), flash_attention_config("32", "8", "64", "0.125"))); const auto & swiglu_k16 = find_kernel("ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32_k16"); - require_compiled_kernel(compile_kernel(*jit, "prefill-swiglu-k16", swiglu_k16, - token_count_workload(512), { + require_compiled_kernel(compile_kernel(*jit, "prefill-swiglu-k16", swiglu_k16, token_count_workload(512), + { { "ggml.mul_mat_swiglu.input_size", "2048" }, { "ggml.mul_mat_swiglu.output_size", "8192" }, { "ggml.mul_mat_swiglu.gate_weight_format", "4" }, @@ -1050,16 +1120,201 @@ static void run_prefill_fusion_compile_checks() { })); const auto & get_rows_rmsnorm = find_kernel("ggml_get_rows_rmsnorm_binary_q8_1_x4_f16"); - require_compiled_kernel(compile_kernel(*jit, "prefill-get-rows-rmsnorm", get_rows_rmsnorm, { + require_compiled_kernel(compile_kernel(*jit, "prefill-get-rows-rmsnorm", get_rows_rmsnorm, + { { "token_count", 512 }, { "row_count", 128256 }, { "hidden_size", 2048 }, - }, { + }, + { { "ggml.get_rows_f32.token_capacity", "512" }, { "ggml.get_rows_f32.hidden_capacity", "2048" }, { "ggml.get_rows_f32.weight_format", "6" }, { "ggml.get_rows_rmsnorm.rms_epsilon", "0.00001" }, })); + + const auto & get_rows_scale = find_kernel("ggml_get_rows_scale_f32"); + require_compiled_kernel(compile_kernel(*jit, "prefill-get-rows-scale", get_rows_scale, + { + { "token_count", 64 }, + { "row_count", 122753 }, + { "hidden_size", 2560 }, + }, + { + { "ggml.get_rows_f32.token_capacity", "64" }, + { "ggml.get_rows_f32.hidden_capacity", "2560" }, + { "ggml.get_rows_f32.weight_format", "32" }, + { "ggml.get_rows_scale_f32.scale", "0.177800179" }, + { "ggml.get_rows_scale_f32.bias", "0" }, + })); + require_compiled_kernel(compile_kernel(*jit, "gemma-prefill-get-rows-scale-q5_1", get_rows_scale, + { + { "token_count", 64 }, + { "row_count", 262144 }, + { "hidden_size", 640 }, + }, + { + { "ggml.get_rows_f32.token_capacity", "64" }, + { "ggml.get_rows_f32.hidden_capacity", "640" }, + { "ggml.get_rows_f32.weight_format", "51" }, + { "ggml.get_rows_scale_f32.scale", "25.2982216" }, + { "ggml.get_rows_scale_f32.bias", "0" }, + })); + require_compiled_kernel(compile_kernel(*jit, "gemma-decode-get-rows-scale-q5_1", get_rows_scale, + { + { "token_count", 1 }, + { "row_count", 262144 }, + { "hidden_size", 640 }, + }, + { + { "ggml.get_rows_f32.token_capacity", "1" }, + { "ggml.get_rows_f32.hidden_capacity", "640" }, + { "ggml.get_rows_f32.weight_format", "51" }, + { "ggml.get_rows_scale_f32.scale", "25.2982216" }, + { "ggml.get_rows_scale_f32.bias", "0" }, + })); + + const auto & lfm_dconv3 = find_kernel("llm_ssm_conv_lfm_dconv3_state_f32"); + REQUIRE(lfm_dconv3.bindings.size() == 5); + for (const int64_t tokens : {1, 64}) { + auto ref = compile_kernel(*jit, "lfm-dconv3-" + std::to_string(tokens), lfm_dconv3, {}, { + {"llm.ssm_conv.lfm_dconv3.n_t", std::to_string(tokens)}, + }); + REQUIRE(ref != nullptr); + REQUIRE(resolve_with_timeout(ref, std::chrono::seconds(120))); + auto compiled = ref->take_result(); + REQUIRE(compiled.hsaco_data != nullptr); + REQUIRE(compiled.launch_config.workgroup_size[0] == 256); + REQUIRE(compiled.launch_config.workgroup_count[0] == static_cast(8 * tokens)); + } +} + +static void run_binary_fusion_compile_checks() { + std::string error; + std::unique_ptr jit = + ggml::hrx::create_loom_jit("gfx1100", ggml::hrx::LoomJitMode::Sync, error); + REQUIRE(jit != nullptr); + + const auto & gather = find_kernel("ggml_gather_add_f32"); + require_compiled_kernel(compile_kernel(*jit, "binary-fusion-gather-div", gather, + { + { "source_token_count", 2 }, + { "output_token_count", 1 }, + { "hidden_size", 2048 }, + }, + { + { "ggml.gather_add_f32.binary_op", "3" }, + })); + + const auto & scale = find_kernel("ggml_scale_add_f32"); + require_compiled_kernel(compile_kernel(*jit, "binary-fusion-scale-sub", scale, + { + { "element_count", 128 }, + }, + { + { "ggml.scale_add_f32.ne0", "128" }, + { "ggml.scale_add_f32.ne1", "1" }, + { "ggml.scale_add_f32.ne2", "1" }, + { "ggml.scale_add_f32.input_stride1", "128" }, + { "ggml.scale_add_f32.input_stride2", "128" }, + { "ggml.scale_add_f32.input_stride3", "128" }, + { "ggml.scale_add_f32.input_span", "128" }, + { "ggml.scale_add_f32.residual_stride1", "128" }, + { "ggml.scale_add_f32.residual_stride2", "128" }, + { "ggml.scale_add_f32.residual_stride3", "128" }, + { "ggml.scale_add_f32.residual_span", "128" }, + { "ggml.scale_add_f32.scale", "0.5" }, + { "ggml.scale_add_f32.bias", "0" }, + { "ggml.scale_add_f32.binary_op", "1" }, + { "ggml.scale_add_f32.scaled_lhs", "0" }, + })); + + const auto & rms = find_kernel("ggml_rmsnorm_mul_add_f32"); + require_compiled_kernel(compile_kernel(*jit, "binary-fusion-rms-div-sub", rms, + { + { "token_count", 1 }, + }, + { + { "ggml.rmsnorm_mul_add_f32.hidden_size", "128" }, + { "ggml.rmsnorm_mul_add_f32.rms_epsilon", "0.000001" }, + { "ggml.rmsnorm_mul_add_f32.normalized_op", "3" }, + { "ggml.rmsnorm_mul_add_f32.output_op", "1" }, + { "ggml.rmsnorm_mul_add_f32.binary_lhs", "0" }, + })); + + const auto & gather_rms = find_kernel("ggml_gather_add_rmsnorm_binary_f32"); + require_compiled_kernel(compile_kernel(*jit, "binary-fusion-gather-rms-mul", gather_rms, + { + { "source_token_count", 2 }, + { "output_token_count", 1 }, + }, + { + { "ggml.gather_add_rmsnorm_binary_f32.hidden_size", "128" }, + { "ggml.gather_add_rmsnorm_binary_f32.rms_epsilon", "0.000001" }, + { "ggml.gather_add_rmsnorm_binary_f32.gather_op", "2" }, + })); + + const auto & ssm = find_kernel("llm_ssm_conv_binary_f32"); + require_compiled_kernel(compile_kernel(*jit, "binary-fusion-ssm-rhs-div", ssm, + { + }, + { + { "llm.ssm_conv.generic.d_conv", "4" }, + { "llm.ssm_conv.generic.d_inner", "64" }, + { "llm.ssm_conv.generic.n_t", "8" }, + { "llm.ssm_conv.generic.n_s", "1" }, + { "llm.ssm_conv.generic.unary_op", "23" }, + { "llm.ssm_conv.generic.binary_op", "3" }, + { "llm.ssm_conv.generic.binary_lhs", "0" }, + { "llm.ssm_conv.generic.workgroup_size", "256" }, + })); + + const auto & vector_postops = find_kernel("ggml_mul_mat_vector_bias_residual_f32_f32"); + for (const int64_t epilogue : { int64_t{ 1 }, int64_t{ 2 }, int64_t{ 3 } }) { + require_compiled_kernel(compile_kernel( + *jit, "binary-fusion-vector-postops-" + std::to_string(epilogue), vector_postops, token_count_workload(1), + { + { "ggml.matmul.vector_postops.input_size", "2048" }, + { "ggml.matmul.vector_postops.output_size", "2048" }, + { "ggml.matmul.vector_postops.weight_format", "4" }, + { "ggml.matmul.vector_postops.apply_bias", epilogue & 1 ? "1" : "0" }, + { "ggml.matmul.vector_postops.apply_residual", epilogue & 2 ? "1" : "0" }, + { "ggml.workload.token_capacity", "1" }, + })); + } + + const auto compile_postops = [&](const char * kernel_name, const char * key_prefix, int64_t token_count, + int64_t token_capacity) { + const auto & kernel = find_kernel(kernel_name); + for (const int64_t epilogue : { int64_t{ 1 }, int64_t{ 2 }, int64_t{ 3 } }) { + require_compiled_kernel(compile_kernel( + *jit, std::string(key_prefix) + std::to_string(epilogue), kernel, token_count_workload(token_count), + { + { "ggml.mul_mat_postops.input_size", "2048" }, + { "ggml.mul_mat_postops.output_size", "2048" }, + { "ggml.mul_mat_postops.weight_format", "4" }, + { "ggml.mul_mat_postops.epilogue", std::to_string(epilogue) }, + { "ggml.workload.token_capacity", std::to_string(token_capacity) }, + })); + } + }; + compile_postops("ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", "binary-fusion-skinny-postops-", 2, 2); + compile_postops("ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", "binary-fusion-tiled-postops-", 32, 32); + + const auto & pair_postops = find_kernel("ggml_mul_mat_tiled_pair_input_f32_binary_bias_residual_publish_f32"); + for (const int64_t epilogue : { int64_t{ 1 }, int64_t{ 2 }, int64_t{ 3 } }) { + require_compiled_kernel(compile_kernel(*jit, "binary-fusion-pair-postops-" + std::to_string(epilogue), + pair_postops, token_count_workload(32), + { + { "ggml.mul_mat_swiglu.input_size", "2048" }, + { "ggml.mul_mat_swiglu.output_size", "2048" }, + { "ggml.mul_mat_swiglu.gate_weight_format", "4" }, + { "ggml.mul_mat_swiglu.up_weight_format", "4" }, + { "ggml.mul_mat_swiglu.op", "4" }, + { "ggml.mul_mat_swiglu.epilogue", std::to_string(epilogue) }, + { "ggml.workload.token_capacity", "32" }, + })); + } } static void run_loom_jit_materialize_checks() { @@ -1079,13 +1334,15 @@ static bool has_hrx_test_device() { static void register_hrx_loom_jit_cases(test_runner::Suite & suite) { suite.host_case("compile", [] { run_loom_jit_compile_checks(); }); + suite.host_case("moe-routing", [] { run_moe_routing_compile_checks(); }); suite.host_case("prefill-fusions", [] { run_prefill_fusion_compile_checks(); }); + suite.host_case("binary-fusions", [] { run_binary_fusion_compile_checks(); }); suite.device_case("materialize", [] { run_loom_jit_materialize_checks(); }); } int main(int argc, char ** argv) { - test_runner::Suite suite(test_runner::Config::with_prefix( - "HRX Loom JIT test", "hrx-loom-jit", "GGML_HRX_LOOM_JIT_TEST", 2)); + test_runner::Suite suite( + test_runner::Config::with_prefix("HRX Loom JIT test", "hrx-loom-jit", "GGML_HRX_LOOM_JIT_TEST", 2)); register_hrx_loom_jit_cases(suite); const bool has_device = has_hrx_test_device(); diff --git a/tests/test-hrx-ops.cpp b/tests/test-hrx-ops.cpp index 97bd2ea38d91..0d81b9c53eb2 100644 --- a/tests/test-hrx-ops.cpp +++ b/tests/test-hrx-ops.cpp @@ -10,6 +10,7 @@ #include "testing_suite.h" #include +#include #include #include #include @@ -1068,8 +1069,7 @@ static void run_rmsnorm_support_checks() { REQUIRE(q8_graph != nullptr); ggml_build_forward_expand(q8_graph, q8_output); require_kernel_subsequence(scheduled_kernel_sequence(q8_graph), - { "loom_libs:ggml_rmsnorm_binary_f32", - "loom_libs:ggml_mul_mat_vector_q6_f32_f32" }); + { "loom_libs:ggml_rmsnorm_binary_f32", "loom_libs:ggml_mul_mat_vector_q6_f32_f32" }); ggml_tensor * wrong_type_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 256, 1); ggml_tensor * wrong_type_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F16, 256); @@ -1207,6 +1207,147 @@ static void run_rmsnorm_binary_add_cpu_reference_case() { ggml_backend_free(hrx_backend); } +static void run_gemma_post_proj_cpu_reference_case(int64_t token_count) { + constexpr int64_t hidden_size = 640; + constexpr float epsilon = 1.0e-6f; + ggml_backend_t cpu_backend = init_cpu_backend(); + ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); + REQUIRE(hrx_backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = 1024 * 1024; + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); + + struct GraphTensors { + ggml_tensor * input; + ggml_tensor * weight; + ggml_tensor * residual; + ggml_tensor * raw; + ggml_tensor * raw_side; + ggml_tensor * output; + }; + + auto build = [&](ggml_context * ctx) { + GraphTensors g = {}; + g.input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, token_count); + g.weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size); + g.residual = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, token_count); + g.raw = ggml_add(ctx, g.input, g.residual); + ggml_tensor * normalized = ggml_rms_norm(ctx, g.raw, epsilon); + ggml_tensor * scaled = ggml_mul(ctx, normalized, g.weight); + g.output = ggml_add(ctx, scaled, g.residual); + g.raw_side = ggml_mul(ctx, g.raw, g.residual); + REQUIRE(g.output != nullptr); + REQUIRE(g.raw_side != nullptr); + return g; + }; + GraphTensors cpu = build(cpu_ctx); + GraphTensors hrx = build(hrx_ctx); + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu.output); + ggml_build_forward_expand(cpu_graph, cpu.raw_side); + ggml_build_forward_expand(hrx_graph, hrx.output); + ggml_build_forward_expand(hrx_graph, hrx.raw_side); + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_rmsnorm_mul_add_f32" }); + + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + const std::vector input = make_pattern_f32(hidden_size * token_count, 37, 0.02f); + const std::vector residual = make_pattern_f32(hidden_size * token_count, 43, 0.04f); + const std::vector weight = make_weight(hidden_size); + set_tensor_pair_bytes(cpu_backend, cpu.input, hrx_backend, hrx.input, input.data(), input.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu.residual, hrx_backend, hrx.residual, residual.data(), + residual.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu.weight, hrx_backend, hrx.weight, weight.data(), + weight.size() * sizeof(float)); + REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(cpu_backend); + ggml_backend_synchronize(hrx_backend); + require_close(get_f32_tensor(hrx_backend, hrx.output), get_f32_tensor(cpu_backend, cpu.output), 5.0e-4f); + require_close(get_f32_tensor(hrx_backend, hrx.raw_side), get_f32_tensor(cpu_backend, cpu.raw_side), 1.0e-5f); + + ggml_backend_buffer_free(cpu_buffer); + ggml_backend_buffer_free(hrx_buffer); + ggml_free(cpu_ctx); + ggml_free(hrx_ctx); + ggml_backend_free(cpu_backend); + ggml_backend_free(hrx_backend); +} + +static void run_gemma_post_proj_matcher_negatives() { + constexpr int64_t hidden_size = 640; + ggml_init_params params = {}; + params.mem_size = 1024 * 1024; + params.no_alloc = true; + const ggml::hrx::DispatchRegistry * registry = ggml::hrx::find_dispatch_registry({ "gfx1151" }); + REQUIRE(registry != nullptr); + + for (const int variant : { 0, 1, 2, 3, 4 }) { + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, 3); + ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size); + ggml_tensor * rhs = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, 3); + ggml_tensor * raw = ggml_add(ctx, input, rhs); + ggml_tensor * residual = variant == 2 ? ggml_view_2d(ctx, raw, hidden_size, 3, raw->nb[1], 0) : + variant == 3 ? ggml_mul(ctx, raw, rhs) : + rhs; + ggml_tensor * rms = ggml_rms_norm(ctx, raw, 1.0e-6f); + ggml_tensor * scaled = variant == 4 ? ggml_add(ctx, rms, weight) : ggml_mul(ctx, rms, weight); + ggml_tensor * output = ggml_add(ctx, scaled, residual); + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + if (variant == 1) { + ggml_build_forward_expand(graph, ggml_mul(ctx, scaled, rhs)); + } + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + const size_t rms_index = producer_index_for_tensor(imported.graph, rms); + const size_t raw_index = producer_index_for_tensor(imported.graph, raw); + ggml::hrx::CommandPlan plan; + std::vector covered(imported.graph.nodes().size(), false); + auto kernel = [&]() { + ggml::hrx::DispatchMatch match; + const ggml::hrx::DispatchMatchContext context = { + imported.graph, &imported.graph.nodes()[rms_index], rms_index, covered, + plan, next_plan_value(imported.graph, plan), + }; + REQUIRE(registry->match(context, match)); + REQUIRE(match.dispatches.size() == 1); + return kernel_name_for_id(match.dispatches[0].kernel.kernel_id); + }; + REQUIRE(kernel() == "loom_libs:ggml_rmsnorm_binary_f32"); + covered[raw_index] = true; + if (variant == 3) { + REQUIRE(kernel() == "loom_libs:ggml_rmsnorm_binary_f32"); + covered[producer_index_for_tensor(imported.graph, residual)] = true; + } + if (variant == 2) { + const size_t view_index = producer_index_for_tensor(imported.graph, residual); + covered[view_index] = true; + const ggml::hrx::Value * raw_value = imported.graph.values().find_tensor(raw); + const ggml::hrx::Value * residual_value = imported.graph.values().find_tensor(residual); + REQUIRE(raw_value != nullptr); + REQUIRE(residual_value != nullptr); + REQUIRE(raw_value->storage_root == residual_value->storage_root); + } + REQUIRE(kernel() == (variant == 0 || variant == 3 ? "loom_libs:ggml_rmsnorm_mul_add_f32" : + "loom_libs:ggml_rmsnorm_binary_f32")); + ggml_free(ctx); + } +} + static void run_rmsnorm_gate_cpu_reference_case(int64_t hidden_size, int64_t head_count, int64_t token_count, @@ -1229,6 +1370,7 @@ static void run_rmsnorm_gate_cpu_reference_case(int64_t hidden_size, ggml_tensor * output; ggml_cgraph * graph; }; + auto build = [&](ggml_context * ctx) { Graph g; g.input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, hidden_size, head_count, token_count); @@ -1410,6 +1552,77 @@ static void run_strided_rmsnorm_cpu_reference_case(int64_t token_count) { ggml_backend_free(hrx_backend); } +static void run_strided_rmsnorm_binary_cpu_reference_case() { + ggml_backend_t cpu_backend = init_cpu_backend(); + ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); + REQUIRE(hrx_backend != nullptr); + + constexpr int64_t hidden_size = 256; + constexpr int64_t token_count = 64; + constexpr int64_t input_stride = 288; + constexpr float epsilon = 1.0e-5f; + const size_t storage_elements = static_cast((token_count - 1) * input_stride + hidden_size); + + ggml_init_params params = {}; + params.mem_size = 2 * 1024 * 1024; + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); + + ggml_tensor * cpu_storage = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, storage_elements); + ggml_tensor * cpu_input = + ggml_view_2d(cpu_ctx, cpu_storage, hidden_size, token_count, input_stride * sizeof(float), 0); + ggml_tensor * cpu_scale = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, hidden_size); + ggml_tensor * cpu_output = ggml_mul(cpu_ctx, ggml_rms_norm(cpu_ctx, cpu_input, epsilon), cpu_scale); + ggml_tensor * hrx_storage = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, storage_elements); + ggml_tensor * hrx_input = + ggml_view_2d(hrx_ctx, hrx_storage, hidden_size, token_count, input_stride * sizeof(float), 0); + ggml_tensor * hrx_scale = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, hidden_size); + ggml_tensor * hrx_output = ggml_mul(hrx_ctx, ggml_rms_norm(hrx_ctx, hrx_input, epsilon), hrx_scale); + REQUIRE(cpu_output != nullptr); + REQUIRE(hrx_output != nullptr); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_output); + ggml_build_forward_expand(hrx_graph, hrx_output); + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_rmsnorm_binary_strided_f32" }); + + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + const std::vector rows = make_pattern_f32(hidden_size * token_count, 54, 0.015f); + const std::vector scale = make_pattern_f32(hidden_size, 55, 0.03f); + std::vector storage(storage_elements, -19.0f); + for (int64_t token = 0; token < token_count; ++token) { + std::memcpy(storage.data() + static_cast(token * input_stride), + rows.data() + static_cast(token * hidden_size), + static_cast(hidden_size) * sizeof(float)); + } + set_tensor_pair_bytes(cpu_backend, cpu_storage, hrx_backend, hrx_storage, storage.data(), + storage.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu_scale, hrx_backend, hrx_scale, scale.data(), scale.size() * sizeof(float)); + + REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(cpu_backend); + ggml_backend_synchronize(hrx_backend); + require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), 5.0e-4f); + + ggml_backend_buffer_free(cpu_buffer); + ggml_backend_buffer_free(hrx_buffer); + ggml_free(cpu_ctx); + ggml_free(hrx_ctx); + ggml_backend_free(cpu_backend); + ggml_backend_free(hrx_backend); +} + static void run_strided_cont_cpu_reference_case(int64_t token_count) { ggml_backend_t cpu_backend = init_cpu_backend(); ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); @@ -2819,6 +3032,75 @@ static void run_gather_add_f32_cpu_reference_case() { ggml_backend_free(hrx_backend); } +static void run_scale_add_f32_cpu_reference_case(bool strided_residual) { + ggml_backend_t cpu_backend = init_cpu_backend(); + ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); + REQUIRE(hrx_backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = 512 * 1024; + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); + + constexpr int64_t width = 257; + constexpr int64_t token_count = 3; + constexpr int64_t residual_width = 320; + ggml_tensor * cpu_input = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, width, token_count); + ggml_tensor * hrx_input = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, width, token_count); + ggml_tensor * cpu_residual_storage = + ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, strided_residual ? residual_width : width, token_count); + ggml_tensor * hrx_residual_storage = + ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, strided_residual ? residual_width : width, token_count); + ggml_tensor * cpu_residual = strided_residual ? ggml_view_2d(cpu_ctx, cpu_residual_storage, width, token_count, + cpu_residual_storage->nb[1], 0) : + cpu_residual_storage; + ggml_tensor * hrx_residual = strided_residual ? ggml_view_2d(hrx_ctx, hrx_residual_storage, width, token_count, + hrx_residual_storage->nb[1], 0) : + hrx_residual_storage; + ggml_tensor * cpu_scaled = ggml_scale_bias(cpu_ctx, cpu_input, 0.177800179f, -0.125f); + ggml_tensor * hrx_scaled = ggml_scale_bias(hrx_ctx, hrx_input, 0.177800179f, -0.125f); + ggml_tensor * cpu_output = ggml_add(cpu_ctx, cpu_scaled, cpu_residual); + ggml_tensor * hrx_output = ggml_add(hrx_ctx, hrx_scaled, hrx_residual); + REQUIRE(cpu_output != nullptr); + REQUIRE(hrx_output != nullptr); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_output); + ggml_build_forward_expand(hrx_graph, hrx_output); + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_scale_add_f32" }); + + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + const std::vector input = make_pattern_f32(width * token_count, 31, 0.03125f); + const std::vector residual = + make_pattern_f32((strided_residual ? residual_width : width) * token_count, 37, 0.046875f); + set_tensor_pair_bytes(cpu_backend, cpu_input, hrx_backend, hrx_input, input.data(), input.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu_residual_storage, hrx_backend, hrx_residual_storage, residual.data(), + residual.size() * sizeof(float)); + + REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(cpu_backend); + ggml_backend_synchronize(hrx_backend); + require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), 0.0f); + + ggml_backend_buffer_free(cpu_buffer); + ggml_backend_buffer_free(hrx_buffer); + ggml_free(cpu_ctx); + ggml_free(hrx_ctx); + ggml_backend_free(cpu_backend); + ggml_backend_free(hrx_backend); +} + static void run_get_rows_f32_cpu_reference_case(ggml_type weight_type, int64_t hidden_size = kQwenHiddenSize, int64_t vocabulary_count = 64, @@ -2922,13 +3204,77 @@ static void run_get_rows_q8_1_zero_weight_case() { ggml_backend_free(hrx_backend); } +static void run_get_rows_scale_f32_cpu_reference_case(ggml_type weight_type = GGML_TYPE_F32, + int64_t hidden_size = 256, + int64_t row_count = 8, + int64_t token_count = 6) { + ggml_backend_t cpu_backend = init_cpu_backend(); + ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); + REQUIRE(hrx_backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = 4 * 1024 * 1024; + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); + + const float scale = weight_type == GGML_TYPE_Q5_1 ? 25.2982216f : 0.177800179f; + ggml_tensor * cpu_weight = ggml_new_tensor_2d(cpu_ctx, weight_type, hidden_size, row_count); + ggml_tensor * cpu_ids = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_I32, token_count); + ggml_tensor * cpu_rows = ggml_get_rows(cpu_ctx, cpu_weight, cpu_ids); + ggml_tensor * cpu_output = ggml_scale(cpu_ctx, cpu_rows, scale); + ggml_tensor * hrx_weight = ggml_new_tensor_2d(hrx_ctx, weight_type, hidden_size, row_count); + ggml_tensor * hrx_ids = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_I32, token_count); + ggml_tensor * hrx_rows = ggml_get_rows(hrx_ctx, hrx_weight, hrx_ids); + ggml_tensor * hrx_output = ggml_scale(hrx_ctx, hrx_rows, scale); + REQUIRE(cpu_output != nullptr); + REQUIRE(hrx_output != nullptr); + ggml_set_output(cpu_rows); + ggml_set_output(hrx_rows); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_output); + ggml_build_forward_expand(hrx_graph, hrx_output); + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_get_rows_scale_f32" }); + + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + const std::vector weight = make_matmul_weight_bytes(weight_type, hidden_size, row_count, 9); + const std::vector ids = { 0, static_cast(row_count - 1), 2, + 2, static_cast(row_count - 2), 1 }; + set_tensor_pair_bytes(cpu_backend, cpu_weight, hrx_backend, hrx_weight, weight.data(), weight.size()); + set_tensor_pair_bytes(cpu_backend, cpu_ids, hrx_backend, hrx_ids, ids.data(), ids.size() * sizeof(int32_t)); + + REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(cpu_backend); + ggml_backend_synchronize(hrx_backend); + const float tolerance = weight_type == GGML_TYPE_F32 ? 0.0f : 5.0e-4f; + require_close(get_f32_tensor(hrx_backend, hrx_rows), get_f32_tensor(cpu_backend, cpu_rows), tolerance); + require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), tolerance); + + ggml_backend_buffer_free(cpu_buffer); + ggml_backend_buffer_free(hrx_buffer); + ggml_free(cpu_ctx); + ggml_free(hrx_ctx); + ggml_backend_free(cpu_backend); + ggml_backend_free(hrx_backend); +} + static auto allocate_test_weight(ggml_backend_t backend, ggml_tensor * weight) { const auto buft = ggml_backend_get_default_buffer_type(backend); const size_t alignment = ggml_backend_buft_get_alignment(buft); const size_t size = GGML_PAD(ggml_backend_buft_get_alloc_size(buft, weight), alignment); auto buffer = std::unique_ptr( - ggml_backend_buft_alloc_buffer(buft, size), - ggml_backend_buffer_free); + ggml_backend_buft_alloc_buffer(buft, size), ggml_backend_buffer_free); REQUIRE(buffer != nullptr); ggml_tallocr allocator = ggml_tallocr_new(buffer.get()); REQUIRE(ggml_tallocr_alloc(&allocator, weight) == GGML_STATUS_SUCCESS); @@ -3020,8 +3366,8 @@ static void run_dense_matmul_cpu_reference_case(ggml_type weight_type, const bool weight_is_bf16 = weight_type == GGML_TYPE_BF16; const bool weight_is_f32 = weight_type == GGML_TYPE_F32; const bool weight_is_dense_float = weight_is_f16 || weight_is_bf16 || weight_is_f32; - std::vector input = - weight_is_dense_float ? std::vector(static_cast(input_size * token_count), 0.00390625f) : + std::vector input = weight_is_dense_float ? + std::vector(static_cast(input_size * token_count), 0.00390625f) : make_pattern_f32(input_size * token_count, 7, 0.01f); if (cache_type != GGML_TYPE_COUNT) { for (size_t i = 0; i < input.size(); ++i) { @@ -3033,7 +3379,8 @@ static void run_dense_matmul_cpu_reference_case(ggml_type weight_type, set_tensor_pair_bytes(cpu_backend, cpu_input, hrx_backend, hrx_input, input.data(), input.size() * sizeof(float)); if (normalize_input) { const std::vector norm_w = make_weight(input_size); - set_tensor_pair_bytes(cpu_backend, cpu_norm_w, hrx_backend, hrx_norm_w, norm_w.data(), norm_w.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu_norm_w, hrx_backend, hrx_norm_w, norm_w.data(), + norm_w.size() * sizeof(float)); } if (cache_type != GGML_TYPE_COUNT) { const int64_t index = 7; @@ -3047,8 +3394,8 @@ static void run_dense_matmul_cpu_reference_case(ggml_type weight_type, ggml_backend_synchronize(cpu_backend); ggml_backend_synchronize(hrx_backend); const float abs_tolerance = cache_type != GGML_TYPE_COUNT ? 2.0e-3f : 1.0f; - require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), - abs_tolerance, cache_type != GGML_TYPE_COUNT ? 1.0e-4f : 3.0e-2f); + require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), abs_tolerance, + cache_type != GGML_TYPE_COUNT ? 1.0e-4f : 3.0e-2f); ggml_backend_buffer_free(cpu_buffer); ggml_backend_buffer_free(hrx_buffer); @@ -3058,6 +3405,109 @@ static void run_dense_matmul_cpu_reference_case(ggml_type weight_type, ggml_backend_free(hrx_backend); } +static void run_prefill_value_cache_repeated_indices_case() { + static constexpr const char * disable_env = "GGML_HRX_DISABLE_PREFILL_V_CACHE_FUSION"; + const char * old_env = std::getenv(disable_env); + const bool had_env = old_env != nullptr; + const std::string old_value = had_env ? old_env : ""; + REQUIRE(unsetenv(disable_env) == 0); + + constexpr int64_t input_size = 1536; + constexpr int64_t output_size = 256; + constexpr int64_t token_count = 64; + constexpr int64_t vocabulary_count = 128; + constexpr int64_t cache_row_count = 512; + + ggml_backend_t fused_backend = ggml_backend_hrx_init(0); + ggml_backend_t split_backend = ggml_backend_hrx_init(0); + REQUIRE(fused_backend != nullptr); + REQUIRE(split_backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = static_cast(32 * 1024 * 1024); + params.no_alloc = true; + ggml_context * fused_ctx = ggml_init(params); + ggml_context * split_ctx = ggml_init(params); + REQUIRE(fused_ctx != nullptr); + REQUIRE(split_ctx != nullptr); + + auto build_graph = [](ggml_context * ctx) { + ggml_tensor * embedding = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, input_size, vocabulary_count); + ggml_tensor * tokens = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); + ggml_tensor * norm = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, input_size); + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, input_size, output_size); + ggml_tensor * indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, token_count); + ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, output_size, cache_row_count); + ggml_tensor * rows = ggml_get_rows(ctx, embedding, tokens); + ggml_tensor * input = build_rmsnorm_mul_graph(ctx, rows, norm); + ggml_tensor * projected = ggml_mul_mat(ctx, weight, input); + ggml_tensor * output = ggml_set_rows(ctx, cache, projected, indices); + return std::array{ embedding, tokens, norm, weight, indices, cache, output }; + }; + + const auto fused = build_graph(fused_ctx); + const auto split = build_graph(split_ctx); + auto fused_embedding_buffer = allocate_test_weight(fused_backend, fused[0]); + auto fused_projection_buffer = allocate_test_weight(fused_backend, fused[3]); + auto split_embedding_buffer = allocate_test_weight(split_backend, split[0]); + auto split_projection_buffer = allocate_test_weight(split_backend, split[3]); + + ggml_cgraph * fused_graph = ggml_new_graph(fused_ctx); + ggml_cgraph * split_graph = ggml_new_graph(split_ctx); + REQUIRE(fused_graph != nullptr); + REQUIRE(split_graph != nullptr); + ggml_build_forward_expand(fused_graph, fused[6]); + ggml_build_forward_expand(split_graph, split[6]); + require_kernel_subsequence(scheduled_kernel_sequence(fused_graph), + { "loom_libs:ggml_get_rows_rmsnorm_binary_q8_1_x4_f16", + "loom_libs:llm_attention_v_matmul_set_rows_tiled_f32_f32" }); + REQUIRE(setenv(disable_env, "1", 1) == 0); + require_kernel_subsequence(scheduled_kernel_sequence(split_graph), + { "loom_libs:ggml_get_rows_rmsnorm_binary_q8_1_x4_f16", + "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32", "loom_libs:ggml_set_rows" }); + + ggml_backend_buffer_t fused_buffer = ggml_backend_alloc_ctx_tensors(fused_ctx, fused_backend); + ggml_backend_buffer_t split_buffer = ggml_backend_alloc_ctx_tensors(split_ctx, split_backend); + REQUIRE(fused_buffer != nullptr); + REQUIRE(split_buffer != nullptr); + + const std::vector embedding = make_matmul_weight_bytes(GGML_TYPE_Q4_K, input_size, vocabulary_count, 5); + const std::vector projection = make_matmul_weight_bytes(GGML_TYPE_Q6_K, input_size, output_size, 11); + const std::vector norm = make_weight(input_size); + const std::vector tokens(static_cast(token_count), 3); + std::vector indices(static_cast(token_count)); + for (int64_t token = 0; token < token_count; ++token) { + indices[static_cast(token)] = (token * 37 + 11) % 31; + } + const std::vector cache(static_cast(output_size * cache_row_count), ggml_fp32_to_fp16(-2.5f)); + + set_tensor_pair_bytes(fused_backend, fused[0], split_backend, split[0], embedding.data(), embedding.size()); + set_tensor_pair_bytes(fused_backend, fused[1], split_backend, split[1], tokens.data(), + tokens.size() * sizeof(int32_t)); + set_tensor_pair_bytes(fused_backend, fused[2], split_backend, split[2], norm.data(), norm.size() * sizeof(float)); + set_tensor_pair_bytes(fused_backend, fused[3], split_backend, split[3], projection.data(), projection.size()); + set_tensor_pair_bytes(fused_backend, fused[4], split_backend, split[4], indices.data(), + indices.size() * sizeof(int64_t)); + set_tensor_pair_bytes(fused_backend, fused[5], split_backend, split[5], cache.data(), + cache.size() * sizeof(ggml_fp16_t)); + + REQUIRE(unsetenv(disable_env) == 0); + REQUIRE(ggml_backend_graph_compute(fused_backend, fused_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(fused_backend); + REQUIRE(setenv(disable_env, "1", 1) == 0); + REQUIRE(ggml_backend_graph_compute(split_backend, split_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(split_backend); + require_close(get_f32_tensor(fused_backend, fused[6]), get_f32_tensor(split_backend, split[6]), 0.0f); + + ggml_backend_buffer_free(fused_buffer); + ggml_backend_buffer_free(split_buffer); + ggml_free(fused_ctx); + ggml_free(split_ctx); + ggml_backend_free(fused_backend); + ggml_backend_free(split_backend); + restore_environment_value(disable_env, had_env, old_value); +} + static void run_dense_matmul_unary_cpu_reference_case(ggml_type weight_type, const char * expected_kernel, int64_t token_count, @@ -3369,8 +3819,7 @@ static void run_vector_q8_publish_cpu_reference_case() { ggml_build_forward_expand(cpu_graph, cpu_output); ggml_build_forward_expand(hrx_graph, hrx_output); - require_kernel_subsequence( - scheduled_kernel_sequence(hrx_graph), + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_mul_mat_vector_f32_f32", "qwen3_moe:ggml_quantize_q8_1_x4_f32", "loom_libs:ggml_mul_mat_vector_f32_f32" }); @@ -3411,8 +3860,7 @@ static void run_vector_q8_publish_cpu_reference_case() { REQUIRE(cpu_buffer != nullptr); REQUIRE(hrx_buffer != nullptr); - const std::vector first_weight = - make_matmul_weight_bytes(GGML_TYPE_F32, input_size, intermediate_size, 6); + const std::vector first_weight = make_matmul_weight_bytes(GGML_TYPE_F32, input_size, intermediate_size, 6); const std::vector second_weight = make_matmul_weight_bytes(GGML_TYPE_Q4_K, intermediate_size, output_size, 11); const std::vector input = make_pattern_f32(input_size * token_count, 7, 0.01f); @@ -3667,6 +4115,157 @@ static void run_endpoint_rmsnorm_q6k_q8_cpu_reference_case() { ggml_backend_free(hrx_backend); } +static void run_lfm_head_rope_cpu_reference_case(int64_t head_count, int64_t token_count, bool cache_output) { + constexpr int64_t head_size = 64; + constexpr float epsilon = 1.0e-5f; + const int64_t row_size = head_size * head_count; + const int64_t cache_rows = token_count + 64; + + ggml_backend_t cpu_backend = init_cpu_backend(); + ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); + REQUIRE(hrx_backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = 1024 * 1024; + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); + + struct HeadGraph { + ggml_tensor * input; + ggml_tensor * weight; + ggml_tensor * positions; + ggml_tensor * cache; + ggml_tensor * row_ids; + ggml_tensor * output; + }; + + auto build = [&](ggml_context * ctx) { + HeadGraph g = {}; + g.input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, head_size, head_count, token_count); + g.weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, head_size); + g.positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); + ggml_tensor * normalized = ggml_rms_norm(ctx, g.input, epsilon); + ggml_tensor * scaled = ggml_mul(ctx, normalized, g.weight); + ggml_tensor * rope = ggml_rope_ext(ctx, scaled, g.positions, nullptr, head_size, GGML_ROPE_TYPE_NEOX, 0, + 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + g.output = rope; + if (cache_output) { + g.cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, row_size, cache_rows); + g.row_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, token_count); + ggml_tensor * rows = ggml_reshape_2d(ctx, rope, row_size, token_count); + ggml_tensor * view = ggml_view_2d(ctx, rows, row_size, token_count, rows->nb[1], 0); + g.output = ggml_set_rows(ctx, g.cache, view, g.row_ids); + } + REQUIRE(g.output != nullptr); + return g; + }; + + HeadGraph cpu = build(cpu_ctx); + HeadGraph hrx = build(hrx_ctx); + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu.output); + ggml_build_forward_expand(hrx_graph, hrx.output); + const std::vector kernels = scheduled_kernel_sequence(hrx_graph); + if (cache_output) { + REQUIRE(std::find(kernels.begin(), kernels.end(), "loom_libs:ggml_rmsnorm_mul_rope_f32") == kernels.end()); + require_kernel_subsequence(kernels, { "loom_libs:ggml_rope_f32", "loom_libs:ggml_set_rows" }); + } else { + require_kernel_subsequence(kernels, { "loom_libs:ggml_rmsnorm_mul_rope_f32" }); + } + + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + const std::vector input = make_pattern_f32(static_cast(row_size * token_count), 47, 0.03f); + const std::vector weight = make_weight(head_size); + std::vector positions(static_cast(token_count)); + std::vector row_ids(static_cast(token_count)); + for (int64_t i = 0; i < token_count; ++i) { + positions[static_cast(i)] = static_cast(3 + 2 * i); + row_ids[static_cast(i)] = (17 * i + 5) % cache_rows; + } + set_tensor_pair_bytes(cpu_backend, cpu.input, hrx_backend, hrx.input, input.data(), input.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu.weight, hrx_backend, hrx.weight, weight.data(), + weight.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu.positions, hrx_backend, hrx.positions, positions.data(), + positions.size() * sizeof(int32_t)); + if (cache_output) { + const std::vector cache(static_cast(row_size * cache_rows), ggml_fp32_to_fp16(-3.0f)); + set_tensor_pair_bytes(cpu_backend, cpu.cache, hrx_backend, hrx.cache, cache.data(), + cache.size() * sizeof(ggml_fp16_t)); + set_tensor_pair_bytes(cpu_backend, cpu.row_ids, hrx_backend, hrx.row_ids, row_ids.data(), + row_ids.size() * sizeof(int64_t)); + } + + REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(cpu_backend); + ggml_backend_synchronize(hrx_backend); + const float tolerance = cache_output ? 1.0e-3f : 1.0e-4f; + require_close(get_f32_tensor(hrx_backend, hrx.output), get_f32_tensor(cpu_backend, cpu.output), tolerance, + tolerance); + + ggml_backend_buffer_free(cpu_buffer); + ggml_backend_buffer_free(hrx_buffer); + ggml_free(cpu_ctx); + ggml_free(hrx_ctx); + ggml_backend_free(cpu_backend); + ggml_backend_free(hrx_backend); +} + +static void run_lfm_head_rope_availability_checks() { + ggml_init_params params = {}; + params.mem_size = 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 64, 8, 1); + ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 64); + ggml_tensor * pos_f32 = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); + ggml_tensor * positions = ggml_cast(ctx, pos_f32, GGML_TYPE_I32); + ggml_tensor * rms = ggml_rms_norm(ctx, input, 1.0e-5f); + ggml_tensor * scaled = ggml_mul(ctx, rms, weight); + ggml_tensor * output = ggml_rope_ext(ctx, scaled, positions, nullptr, 64, GGML_ROPE_TYPE_NEOX, 0, 10000.0f, 1.0f, + 0.0f, 1.0f, 0.0f, 0.0f); + 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()); + const ggml::hrx::DispatchRegistry * registry = ggml::hrx::find_dispatch_registry({ "gfx1151" }); + REQUIRE(registry != nullptr); + const size_t rms_index = producer_index_for_tensor(imported.graph, rms); + const size_t pos_index = producer_index_for_tensor(imported.graph, positions); + ggml::hrx::CommandPlan plan; + std::vector covered(imported.graph.nodes().size(), false); + auto matched_kernel = [&]() { + ggml::hrx::DispatchMatch match; + const ggml::hrx::DispatchMatchContext context = { + imported.graph, &imported.graph.nodes()[rms_index], rms_index, covered, + plan, next_plan_value(imported.graph, plan), + }; + REQUIRE(registry->match(context, match)); + REQUIRE(match.dispatches.size() == 1); + return kernel_name_for_id(match.dispatches[0].kernel.kernel_id); + }; + REQUIRE(matched_kernel() == "loom_libs:ggml_rmsnorm_f32"); + covered[pos_index] = true; + REQUIRE(matched_kernel() == "loom_libs:ggml_rmsnorm_mul_rope_f32"); + + ggml_free(ctx); +} + static void run_rope_set_rows_cpu_reference_case() { ggml_backend_t cpu_backend = init_cpu_backend(); ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); @@ -4495,7 +5094,12 @@ struct SsmConvPrefillGraph { ggml_cgraph * graph = nullptr; }; -static SsmConvPrefillGraph make_ssm_conv_prefill_graph(ggml_type type, int64_t k, int64_t n, bool expose, bool gathered, bool clear_state) { +static SsmConvPrefillGraph make_ssm_conv_prefill_graph(ggml_type type, + int64_t k, + int64_t n, + bool expose, + bool gathered, + bool clear_state) { ggml_init_params params = {}; params.mem_size = 32 * 1024 * 1024; params.no_alloc = true; @@ -4537,8 +5141,8 @@ static SsmConvPrefillGraph make_ssm_conv_prefill_graph(ggml_type type, int64_t k "loom_libs:ggml_mul_mat_quantized_f16_wmma_prefill_conv4"; REQUIRE((std::find(sequence.begin(), sequence.end(), fused) != sequence.end()) == !expose); if (gathered && !expose) { - REQUIRE(std::find(sequence.begin(), sequence.end(), - "loom_libs:llm_ssm_conv_dconv4_silu_prefill_finish_f32") != sequence.end()); + REQUIRE(std::find(sequence.begin(), sequence.end(), "loom_libs:llm_ssm_conv_dconv4_silu_prefill_finish_f32") != + sequence.end()); } return c; } @@ -4589,27 +5193,24 @@ static void run_vector_output_capacity_experiment() { constexpr int64_t kAboveOldOutputSize = 262208; constexpr int64_t kLiftedOutputSize = 524288; - run_dense_matmul_cpu_reference_case( - GGML_TYPE_F32, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, kAboveOldOutputSize, kSmallVectorInputSize); - run_dense_matmul_unary_cpu_reference_case( - GGML_TYPE_F32, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, kAboveOldOutputSize, kSmallVectorInputSize); - run_dense_matmul_postops_cpu_reference_case( - GGML_TYPE_F16, "loom_libs:ggml_mul_mat_vector_bias_f32_f32", 1, kAboveOldOutputSize, true, false, - false, kSmallVectorInputSize); - run_dense_matmul_cpu_reference_case( - GGML_TYPE_F16, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, kLiftedOutputSize, kSmallVectorInputSize); + run_dense_matmul_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, kAboveOldOutputSize, + kSmallVectorInputSize); + run_dense_matmul_unary_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, + kAboveOldOutputSize, kSmallVectorInputSize); + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, "loom_libs:ggml_mul_mat_vector_bias_residual_f32_f32", 1, + kAboveOldOutputSize, true, false, false, kSmallVectorInputSize); + run_dense_matmul_cpu_reference_case(GGML_TYPE_F16, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, kLiftedOutputSize, + kSmallVectorInputSize); } static void run_vector_postops_cpu_reference_case() { constexpr int64_t kVectorInputSize = 256; constexpr int64_t kVectorOutputSize = 256; - run_dense_matmul_postops_cpu_reference_case( - GGML_TYPE_F16, "loom_libs:ggml_mul_mat_vector_bias_f32_f32", 1, kVectorOutputSize, true, false, - false, kVectorInputSize); - run_dense_matmul_postops_cpu_reference_case( - GGML_TYPE_F16, "loom_libs:ggml_mul_mat_vector_bias_residual_f32_f32", 1, kVectorOutputSize, true, true, - false, kVectorInputSize); + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, "loom_libs:ggml_mul_mat_vector_bias_residual_f32_f32", 1, + kVectorOutputSize, true, false, false, kVectorInputSize); + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, "loom_libs:ggml_mul_mat_vector_bias_residual_f32_f32", 1, + kVectorOutputSize, true, true, false, kVectorInputSize); } static const char * glu_op_name(ggml_glu_op op) { @@ -4669,14 +5270,18 @@ static void register_basic_ops_cases(Suite & suite) { suite.device_case("glu_packed_f32.swiglu.outputs8192.tokens14", [] { run_glu_packed_f32_cpu_reference_case(GGML_GLU_OP_SWIGLU, false, 8192, 14); }); suite.device_case("gather_add_f32.cpu_reference", [] { run_gather_add_f32_cpu_reference_case(); }); - for (const ggml_type type : - { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K, GGML_TYPE_Q8_0, GGML_TYPE_IQ4_XS, GGML_TYPE_F16, GGML_TYPE_BF16, - GGML_TYPE_F32 }) { + suite.device_case("scale_add_f32.cpu_reference.packed", [] { run_scale_add_f32_cpu_reference_case(false); }); + suite.device_case("scale_add_f32.cpu_reference.strided", [] { run_scale_add_f32_cpu_reference_case(true); }); + for (const ggml_type type : { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K, GGML_TYPE_Q8_0, GGML_TYPE_IQ4_XS, GGML_TYPE_F16, + GGML_TYPE_BF16, GGML_TYPE_F32 }) { suite.device_case("get_rows_f32." + type_name(type), [type] { run_get_rows_f32_cpu_reference_case(type); }); } suite.device_case("get_rows_f32.q1_0.rows2048", [] { run_get_rows_f32_cpu_reference_case(GGML_TYPE_Q1_0, 2048); }); suite.device_case("get_rows_f32.q5_1.rows640", [] { run_get_rows_f32_cpu_reference_case(GGML_TYPE_Q5_1, 640); }); suite.device_case("get_rows_q8_1.zero_weight", [] { run_get_rows_q8_1_zero_weight_case(); }); + suite.device_case("get_rows_scale_f32.cpu_reference", [] { run_get_rows_scale_f32_cpu_reference_case(); }); + suite.device_case("get_rows_scale_f32.q5_1.cpu_reference", + [] { run_get_rows_scale_f32_cpu_reference_case(GGML_TYPE_Q5_1, 640); }); } static void register_dense_matmul_cases(Suite & suite) { @@ -4685,9 +5290,9 @@ static void register_dense_matmul_cases(Suite & suite) { suite.device_case( "dense_matmul.q4_k.tokens" + std::to_string(tokens) + ".outputs" + std::to_string(outputs), [tokens, outputs] { - run_dense_matmul_cpu_reference_case( - GGML_TYPE_Q4_K, - tokens == 1 ? "loom_libs:ggml_mul_mat_vector_f32_f32" : + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, + tokens == 1 ? + "loom_libs:ggml_mul_mat_vector_f32_f32" : "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", tokens, outputs); }); @@ -4728,9 +5333,8 @@ static void register_dense_matmul_cases(Suite & suite) { }); // Exercise the production K dimension and a partial output tile. suite.device_case("dense_matmul.iq4_xs.prefill.tokens512.outputs64", [] { - run_dense_matmul_cpu_reference_case(GGML_TYPE_IQ4_XS, - "loom_libs:ggml_mul_mat_q5_k_iq4_xs_q8_1_x4_wmma_token256", 512, 64, - 3072, true); + run_dense_matmul_cpu_reference_case( + GGML_TYPE_IQ4_XS, "loom_libs:ggml_mul_mat_q5_k_iq4_xs_q8_1_x4_wmma_token256", 512, 64, 3072, true); }); for (const ggml_type type : { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K }) { for (const ggml_type cache_type : { GGML_TYPE_F16, GGML_TYPE_F32 }) { @@ -4742,6 +5346,8 @@ static void register_dense_matmul_cases(Suite & suite) { }); } } + suite.device_case("attention_v_matmul.q6_k.prefill_f16.repeated_indices", + [] { run_prefill_value_cache_repeated_indices_case(); }); for (const ggml_type type : { GGML_TYPE_Q1_0, GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, GGML_TYPE_Q5_0, GGML_TYPE_Q5_1, GGML_TYPE_IQ2_S, GGML_TYPE_Q6_K, GGML_TYPE_Q8_0, GGML_TYPE_IQ4_NL, GGML_TYPE_F16, GGML_TYPE_BF16 }) { @@ -4752,11 +5358,9 @@ static void register_dense_matmul_cases(Suite & suite) { const char * route = type == GGML_TYPE_IQ4_NL ? "tiled" : "skinny"; const char * kernel = type == GGML_TYPE_IQ4_NL ? "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32" : "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32"; - suite.device_case("dense_matmul." + type_name(type) + "." + route + ".tokens2.outputs" + - std::to_string(outputs), - [type, outputs, input, kernel] { - run_dense_matmul_cpu_reference_case(type, kernel, 2, outputs, input); - }); + suite.device_case( + "dense_matmul." + type_name(type) + "." + route + ".tokens2.outputs" + std::to_string(outputs), + [type, outputs, input, kernel] { run_dense_matmul_cpu_reference_case(type, kernel, 2, outputs, input); }); suite.device_case("dense_matmul." + type_name(type) + ".vector.tokens1.outputs" + std::to_string(outputs), [type, vector_kernel, outputs, input] { run_dense_matmul_cpu_reference_case(type, vector_kernel, 1, outputs, input); @@ -4897,12 +5501,14 @@ static void register_dense_matmul_cases(Suite & suite) { GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", 1024, 2048); }); suite.device_case("dense_matmul_postops.f16.bias.tokens33.outputs256", [] { - run_dense_matmul_postops_cpu_reference_case( - GGML_TYPE_F16, "loom_libs:ggml_mul_mat_tiled_input_f32_bias_publish_f32", 33, 256, true, false, false); + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, + "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", + 33, 256, true, false, false); }); suite.device_case("dense_matmul_postops.f16.residual.tokens33.outputs256", [] { - run_dense_matmul_postops_cpu_reference_case( - GGML_TYPE_F16, "loom_libs:ggml_mul_mat_tiled_input_f32_residual_publish_f32", 33, 256, false, true, false); + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, + "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", + 33, 256, false, true, false); }); suite.device_case("dense_matmul_postops.f16.bias_residual.tokens33.outputs256", [] { run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, @@ -4910,8 +5516,9 @@ static void register_dense_matmul_cases(Suite & suite) { 33, 256, true, true, false); }); suite.device_case("dense_matmul_postops.f16.residual.alias.tokens33.outputs256", [] { - run_dense_matmul_postops_cpu_reference_case( - GGML_TYPE_F16, "loom_libs:ggml_mul_mat_tiled_input_f32_residual_publish_f32", 33, 256, false, true, true); + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, + "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", + 33, 256, false, true, true); }); suite.device_case("dense_matmul_postops.f16.bias_residual.alias.tokens33.outputs256", [] { run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, @@ -4928,14 +5535,14 @@ static void register_dense_matmul_cases(Suite & suite) { }); suite.device_case("dense_matmul_postops." + type_name(type) + ".add.tokens5.outputs4160", [type, input_size] { run_dense_matmul_postops_cpu_reference_case( - type, "loom_libs:ggml_mul_mat_skinny_input_f32_residual_publish_f32", 5, 4160, false, true, false, + type, "loom_libs:ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", 5, 4160, false, true, false, input_size + 256); }); - suite.device_case( - "dense_matmul_postops." + type_name(type) + ".add.alias.tokens5.outputs4224", [type, input_size] { + suite.device_case("dense_matmul_postops." + type_name(type) + ".add.alias.tokens5.outputs4224", + [type, input_size] { run_dense_matmul_postops_cpu_reference_case( - type, "loom_libs:ggml_mul_mat_skinny_input_f32_residual_publish_f32", 5, 4224, false, true, true, - input_size + 256); + type, "loom_libs:ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", 5, 4224, + false, true, true, input_size + 256); }); } suite.device_case("dense_matmul.q4_k.skinny.tokens5.outputs4096.input8192", [] { @@ -4943,17 +5550,16 @@ static void register_dense_matmul_cases(Suite & suite) { 4096, 8192); }); for (const int64_t input_size : { 6144, 6400, 6656 }) { - suite.device_case("dense_matmul.q4_k.skinny.tokens5.outputs5120.input" + std::to_string(input_size), - [input_size] { + suite.device_case( + "dense_matmul.q4_k.skinny.tokens5.outputs5120.input" + std::to_string(input_size), [input_size] { run_dense_matmul_cpu_reference_case( - GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 5, 5120, - input_size); + GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 5, 5120, input_size); }); - suite.device_case( - "dense_matmul_postops.q4_k.add.alias.tokens5.outputs5120.input" + std::to_string(input_size), [input_size] { + suite.device_case("dense_matmul_postops.q4_k.add.alias.tokens5.outputs5120.input" + std::to_string(input_size), + [input_size] { run_dense_matmul_postops_cpu_reference_case( - GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_skinny_input_f32_residual_publish_f32", 5, 5120, false, - true, true, input_size); + GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", + 5, 5120, false, true, true, input_size); }); } suite.device_case("dense_matmul.q4_k.skinny.tokens4.outputs4096.input4096", [] { @@ -4961,14 +5567,14 @@ static void register_dense_matmul_cases(Suite & suite) { 4096, 4096); }); suite.device_case("dense_matmul_postops.q4_k.add.tokens4.outputs4160.input6400", [] { - run_dense_matmul_postops_cpu_reference_case( - GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_skinny_input_f32_residual_publish_f32", 4, 4160, false, true, - false, 6400); + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_Q4_K, + "loom_libs:ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", + 4, 4160, false, true, false, 6400); }); suite.device_case("dense_matmul_postops.q4_k.add.alias.tokens4.outputs4224.input6656", [] { - run_dense_matmul_postops_cpu_reference_case( - GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_skinny_input_f32_residual_publish_f32", 4, 4224, false, true, - true, 6656); + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_Q4_K, + "loom_libs:ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", + 4, 4224, false, true, true, 6656); }); suite.device_case("endpoint_rmsnorm_q6k_q8.cpu_reference", [] { run_endpoint_rmsnorm_q6k_q8_cpu_reference_case(); }); @@ -4977,28 +5583,26 @@ static void register_dense_matmul_cases(Suite & suite) { 4096, 16384); }); suite.device_case("dense_matmul_postops.q6_k.add.tokens5.outputs4160.input16640", [] { - run_dense_matmul_postops_cpu_reference_case( - GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_skinny_input_f32_residual_publish_f32", 5, 4160, false, true, - false, 16640); + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_Q6_K, + "loom_libs:ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", + 5, 4160, false, true, false, 16640); }); suite.device_case("dense_matmul_postops.q6_k.add.alias.tokens5.outputs4224.input16640", [] { - run_dense_matmul_postops_cpu_reference_case( - GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_skinny_input_f32_residual_publish_f32", 5, 4224, false, true, - true, 16640); + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_Q6_K, + "loom_libs:ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", + 5, 4224, false, true, true, 16640); }); for (const int64_t tokens : { 1, 3, 5 }) { suite.device_case( "dense_matmul.q4_k.contiguous.tokens" + std::to_string(tokens) + ".outputs256.input5120", [tokens] { - run_dense_matmul_cpu_reference_case( - GGML_TYPE_Q4_K, + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, tokens == 1 ? "loom_libs:ggml_mul_mat_vector_f32_f32" : "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", tokens, 256, 5120, true); }); suite.device_case( "dense_matmul.q6_k.contiguous.tokens" + std::to_string(tokens) + ".outputs256.input5120", [tokens] { - run_dense_matmul_cpu_reference_case( - GGML_TYPE_Q6_K, + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q6_K, tokens == 1 ? "loom_libs:ggml_mul_mat_vector_f32_f32" : "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", tokens, 256, 5120, true, true); @@ -5015,12 +5619,20 @@ static void register_dense_matmul_cases(Suite & suite) { } static void register_rmsnorm_and_scheduling_cases(Suite & suite) { + suite.host_case("support.lfm_head_rope.availability", [] { run_lfm_head_rope_availability_checks(); }); + suite.host_case("gemma_post_proj.matcher_negatives", [] { run_gemma_post_proj_matcher_negatives(); }); + for (const int64_t token_count : { 1, 3, 64 }) { + suite.device_case("gemma_post_proj.tokens" + std::to_string(token_count), + [token_count] { run_gemma_post_proj_cpu_reference_case(token_count); }); + } suite.device_case("rmsnorm.default", [] { run_rmsnorm_cpu_reference_case(); }); suite.device_case("rmsnorm.hidden3840.tokens18", [] { run_rmsnorm_cpu_reference_case(3840, 18, 1.0e-6f); }); for (const int64_t token_count : { 2, 23 }) { suite.device_case("strided_rmsnorm.tokens" + std::to_string(token_count), [token_count] { run_strided_rmsnorm_cpu_reference_case(token_count); }); } + suite.device_case("strided_rmsnorm_binary.minicpm_pp512_witness", + [] { run_strided_rmsnorm_binary_cpu_reference_case(); }); suite.device_case("rmsnorm_binary_add.cpu_reference", [] { run_rmsnorm_binary_add_cpu_reference_case(); }); suite.device_case("rmsnorm_gate.silu.hidden128.tokens1.stride7", [] { run_rmsnorm_gate_cpu_reference_case(128, 1, 7, GGML_UNARY_OP_SILU); }); @@ -5042,6 +5654,18 @@ static void register_rmsnorm_and_scheduling_cases(Suite & suite) { suite.device_case("external_view_input_rmsnorm", [] { run_external_view_input_rmsnorm_case(); }); suite.device_case("split_local_view_alias_import", [] { run_split_local_view_alias_import_case(); }); suite.device_case("rope_set_rows.cpu_reference", [] { run_rope_set_rows_cpu_reference_case(); }); + for (const int64_t head_count : { 8, 32 }) { + for (const int64_t token_count : { 1, 64 }) { + suite.device_case( + "lfm_head_rope.query.heads" + std::to_string(head_count) + ".tokens" + std::to_string(token_count), + [head_count, token_count] { run_lfm_head_rope_cpu_reference_case(head_count, token_count, false); }); + } + } + suite.device_case("lfm_head_rope.query.heads7.tokens3", [] { run_lfm_head_rope_cpu_reference_case(7, 3, false); }); + for (const int64_t token_count : { 1, 64 }) { + suite.device_case("lfm_head_rope.cache.heads8.tokens" + std::to_string(token_count), + [token_count] { run_lfm_head_rope_cpu_reference_case(8, token_count, true); }); + } suite.device_case("attention_postprocess.cpu_reference", [] { run_attention_postprocess_cpu_reference_case(); }); suite.device_case("routed_moe.q4_k.next_rmsnorm", [] { run_routed_moe_cpu_reference_case(GGML_TYPE_Q4_K, true); }); suite.device_case("flash_attention.asymmetric", [] { run_asymmetric_flash_attention_cpu_reference_case(); });