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 a2a122dea313..a99a106c9b29 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 @@ -68,6 +68,9 @@ static bool format_config_for_type(ggml_type type, FormatConfig & config) { case GGML_TYPE_F16: config = { type, 16 }; return true; + case GGML_TYPE_BF16: + config = { type, 17 }; + return true; case GGML_TYPE_F32: config = { type, 32 }; return true; diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-weight-format.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-weight-format.h index 227fc8f46646..a427826912b1 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-weight-format.h +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-weight-format.h @@ -12,6 +12,7 @@ enum class CommonMulMatWeightFormat { Q8_0, Q8_1, F16, + BF16, F32, }; @@ -32,6 +33,9 @@ inline bool common_mul_mat_format_for_type(ggml_type type, CommonMulMatWeightFor case GGML_TYPE_F16: format = CommonMulMatWeightFormat::F16; return true; + case GGML_TYPE_BF16: + format = CommonMulMatWeightFormat::BF16; + return true; case GGML_TYPE_F32: format = CommonMulMatWeightFormat::F32; return true; @@ -52,6 +56,8 @@ inline int64_t common_mul_mat_format_config_value(CommonMulMatWeightFormat forma return 81; case CommonMulMatWeightFormat::F16: return 16; + case CommonMulMatWeightFormat::BF16: + return 17; case CommonMulMatWeightFormat::F32: return 32; } 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 355918e4f060..a5fd95630942 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 @@ -73,6 +73,7 @@ enum class AttentionWeightFormat { Q8_0, Q8_1, F16, + BF16, F32, }; @@ -93,6 +94,9 @@ static bool format_for_type(ggml_type type, AttentionWeightFormat & format) { case GGML_TYPE_F16: format = AttentionWeightFormat::F16; return true; + case GGML_TYPE_BF16: + format = AttentionWeightFormat::BF16; + return true; case GGML_TYPE_F32: format = AttentionWeightFormat::F32; return true; @@ -113,6 +117,8 @@ static int64_t format_config_value(AttentionWeightFormat format) { return 81; case AttentionWeightFormat::F16: return 16; + case AttentionWeightFormat::BF16: + return 17; case AttentionWeightFormat::F32: return 32; } @@ -120,7 +126,8 @@ static int64_t format_config_value(AttentionWeightFormat format) { } static bool is_dense_float_weight_format(AttentionWeightFormat format) { - return format == AttentionWeightFormat::F16 || format == AttentionWeightFormat::F32; + return format == AttentionWeightFormat::F16 || format == AttentionWeightFormat::BF16 || + format == AttentionWeightFormat::F32; } static std::string to_config_value(int64_t value) { 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 4ce9c4e6d0cc..2c91f0419c25 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 @@ -143,6 +143,9 @@ { "path": "motifs/f16_f16.loom" }, + { + "path": "motifs/bf16_f16.loom" + }, { "path": "motifs/f32_f16.loom" } @@ -486,6 +489,7 @@ "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", "motifs/f32_f16.loom" ] }, @@ -498,6 +502,7 @@ "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", "motifs/f32_f16.loom" ] }, @@ -551,6 +556,7 @@ "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", "motifs/f32_f16.loom" ] }, @@ -565,6 +571,7 @@ "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", "motifs/f32_f16.loom" ] }, @@ -620,6 +627,7 @@ "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", "motifs/f32_f16.loom" ] }, @@ -634,6 +642,7 @@ "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", "motifs/f32_f16.loom" ] }, @@ -682,6 +691,7 @@ "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", "motifs/f32_f16.loom" ] }, @@ -695,6 +705,7 @@ "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", "motifs/f32_f16.loom" ] }, @@ -743,6 +754,7 @@ "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", "motifs/f32_f16.loom" ] }, @@ -754,6 +766,7 @@ "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", "motifs/f32_f16.loom" ] }, @@ -846,6 +859,7 @@ "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", "motifs/f32_f16.loom" ] }, @@ -856,6 +870,7 @@ "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", "motifs/f32_f16.loom" ] }, @@ -909,6 +924,7 @@ "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", "motifs/f32_f16.loom" ] }, @@ -921,6 +937,7 @@ "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", "motifs/f32_f16.loom" ] }, @@ -980,6 +997,7 @@ "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", "motifs/f32_f16.loom" ] }, @@ -992,6 +1010,7 @@ "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", "motifs/f32_f16.loom" ] }, @@ -1155,6 +1174,7 @@ "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", "motifs/f32_f16.loom" ] }, @@ -1167,6 +1187,7 @@ "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", "motifs/f32_f16.loom" ] }, @@ -1214,6 +1235,7 @@ "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", "motifs/f32_f16.loom" ] }, @@ -1226,6 +1248,7 @@ "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", "motifs/f32_f16.loom" ] }, @@ -1273,6 +1296,7 @@ "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", "motifs/f32_f16.loom" ] }, @@ -1285,6 +1309,7 @@ "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", "motifs/f32_f16.loom" ] }, @@ -1334,6 +1359,7 @@ "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", "motifs/f32_f16.loom" ] }, @@ -1346,6 +1372,7 @@ "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", "motifs/f32_f16.loom" ] }, @@ -1401,6 +1428,7 @@ "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", "motifs/f32_f16.loom" ] }, @@ -1413,6 +1441,7 @@ "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", "motifs/f32_f16.loom" ] }, @@ -1472,6 +1501,7 @@ "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", "motifs/f32_f16.loom" ] }, @@ -1482,6 +1512,7 @@ "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", "motifs/f32_f16.loom" ] }, @@ -1527,6 +1558,7 @@ "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", "motifs/f32_f16.loom" ] }, @@ -1537,6 +1569,7 @@ "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", "motifs/f32_f16.loom" ] }, @@ -1698,6 +1731,7 @@ "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", "motifs/f32_f16.loom" ] }, @@ -1710,6 +1744,7 @@ "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", "motifs/f32_f16.loom" ] }, @@ -1773,6 +1808,7 @@ "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", "motifs/f32_f16.loom" ] }, @@ -1785,6 +1821,7 @@ "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", "motifs/f32_f16.loom" ] }, diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/bf16_f16.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/bf16_f16.loom new file mode 100644 index 000000000000..1a38444c8476 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/bf16_f16.loom @@ -0,0 +1,12 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +// Loads four adjacent BF16 weights and truncates them for FP16 matrix staging. +func.def inline @ggml_bf16_f16_vector4(%weight: buffer, %row_byte_base: offset, %input_size: index, %k: index) -> (vector<4xf16>) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %row_view = buffer.view %weight[%row_byte_base] : buffer -> view<[%bounded_input_size]xbf16> + %values_bf16 = vector.load %row_view[%k] : view<[%bounded_input_size]xbf16> -> vector<4xbf16> + %values_f32 = vector.extf %values_bf16 : vector<4xbf16> to vector<4xf32> + %values = vector.fptrunc %values_f32 : vector<4xf32> to vector<4xf16> + func.return %values : vector<4xf16> +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/dequant.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/dequant.loom index 593d8fb452f3..39a2ebfbe7d7 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/dequant.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/dequant.loom @@ -11,6 +11,8 @@ func.decl @ggml_q8_1_f16_vector4(%weight: buffer, %row_byte_base: offset, %q8_bl func.decl @ggml_f16_f16_vector4(%weight: buffer, %row_byte_base: offset, %input_size: index, %k: index) -> (vector<4xf16>) +func.decl @ggml_bf16_f16_vector4(%weight: buffer, %row_byte_base: offset, %input_size: index, %k: index) -> (vector<4xf16>) + func.decl @ggml_f32_f16_vector4(%weight: buffer, %row_byte_base: offset, %input_size: index, %k: index) -> (vector<4xf16>) // Decodes four adjacent Q4_K values to F32 without the FP16 staging round used @@ -220,6 +222,14 @@ func.def inline @ggml_f16_f32_vector4(%weight: buffer, %row_byte_base: offset, % func.return %values : vector<4xf32> } +func.def inline @ggml_bf16_f32_vector4(%weight: buffer, %row_byte_base: offset, %input_size: index, %k: index) -> (vector<4xf32>) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %row_view = buffer.view %weight[%row_byte_base] : buffer -> view<[%bounded_input_size]xbf16> + %values_bf16 = vector.load %row_view[%k] : view<[%bounded_input_size]xbf16> -> vector<4xbf16> + %values = vector.extf %values_bf16 : vector<4xbf16> to vector<4xf32> + func.return %values : vector<4xf32> +} + func.def inline @ggml_f32_f32_vector4(%weight: buffer, %row_byte_base: offset, %input_size: index, %k: index) -> (vector<4xf32>) { %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index %row_view = buffer.view %weight[%row_byte_base] : buffer -> view<[%bounded_input_size]xf32> @@ -231,6 +241,7 @@ func.def inline @ggml_dequant_weight_tile_bytes(%weight_format: index) -> (offse %c4 = index.constant 4 : index %c6 = index.constant 6 : index %c16 = index.constant 16 : index + %c17 = index.constant 17 : index %c32 = index.constant 32 : index %q8_0_format = index.constant 80 : index %q8_1_format = index.constant 81 : index @@ -240,19 +251,22 @@ func.def inline @ggml_dequant_weight_tile_bytes(%weight_format: index) -> (offse %q8_0_tile_bytes = index.constant 272 : offset %q8_1_tile_bytes = index.constant 288 : offset %f16_tile_bytes = index.constant 512 : offset + %bf16_tile_bytes = index.constant 512 : offset %f32_tile_bytes = index.constant 1024 : offset %is_q4 = index.cmp eq, %weight_format, %c4 : index %is_q6 = index.cmp eq, %weight_format, %c6 : index %is_q8_0 = index.cmp eq, %weight_format, %q8_0_format : index %is_q8_1 = index.cmp eq, %weight_format, %q8_1_format : index %is_f16 = index.cmp eq, %weight_format, %c16 : index + %is_bf16 = index.cmp eq, %weight_format, %c17 : index %is_f32 = index.cmp eq, %weight_format, %c32 : index %selected_q4 = scf.select %is_q4, %q4_tile_bytes, %zero_bytes : offset %selected_q6 = scf.select %is_q6, %q6_tile_bytes, %selected_q4 : offset %selected_q8_0 = scf.select %is_q8_0, %q8_0_tile_bytes, %selected_q6 : offset %selected_q8_1 = scf.select %is_q8_1, %q8_1_tile_bytes, %selected_q8_0 : offset %selected_f16 = scf.select %is_f16, %f16_tile_bytes, %selected_q8_1 : offset - %selected = scf.select %is_f32, %f32_tile_bytes, %selected_f16 : offset + %selected_bf16 = scf.select %is_bf16, %bf16_tile_bytes, %selected_f16 : offset + %selected = scf.select %is_f32, %f32_tile_bytes, %selected_bf16 : offset func.return %selected : offset } @@ -263,6 +277,7 @@ func.def inline @ggml_dequant_f16_vector4(%weight_format: index, %weight: buffer %c6 = index.constant 6 : index %c8 = index.constant 8 : index %c16 = index.constant 16 : index + %c17 = index.constant 17 : index %c32 = index.constant 32 : index %q8_0_format = index.constant 80 : index %q8_1_format = index.constant 81 : index @@ -272,6 +287,7 @@ func.def inline @ggml_dequant_f16_vector4(%weight_format: index, %weight: buffer %is_q8_0 = index.cmp eq, %weight_format, %q8_0_format : index %is_q8_1 = index.cmp eq, %weight_format, %q8_1_format : index %is_f16 = index.cmp eq, %weight_format, %c16 : index + %is_bf16 = index.cmp eq, %weight_format, %c17 : index %is_f32 = index.cmp eq, %weight_format, %c32 : index %q8_block_base = index.mul %quant_block, %c8 : index %q8_block = index.add %q8_block_base, %quant_group : index @@ -305,6 +321,12 @@ func.def inline @ggml_dequant_f16_vector4(%weight_format: index, %weight: buffer } else { scf.yield %zero : vector<4xf16> } + %bf16_values = scf.if %is_bf16 -> (vector<4xf16>) { + %values = func.call @ggml_bf16_f16_vector4(%weight, %row_byte_base, %input_size, %k) : (buffer, offset, index, index) -> (vector<4xf16>) + scf.yield %values : vector<4xf16> + } else { + scf.yield %zero : vector<4xf16> + } %f32_values = scf.if %is_f32 -> (vector<4xf16>) { %values = func.call @ggml_f32_f16_vector4(%weight, %row_byte_base, %input_size, %k) : (buffer, offset, index, index) -> (vector<4xf16>) scf.yield %values : vector<4xf16> @@ -316,7 +338,8 @@ func.def inline @ggml_dequant_f16_vector4(%weight_format: index, %weight: buffer %selected_q8_0 = scf.select %is_q8_0, %q8_0_values, %selected_q6 : vector<4xf16> %selected_q8_1 = scf.select %is_q8_1, %q8_1_values, %selected_q8_0 : vector<4xf16> %selected_f16 = scf.select %is_f16, %f16_values, %selected_q8_1 : vector<4xf16> - %selected = scf.select %is_f32, %f32_values, %selected_f16 : vector<4xf16> + %selected_bf16 = scf.select %is_bf16, %bf16_values, %selected_f16 : vector<4xf16> + %selected = scf.select %is_f32, %f32_values, %selected_bf16 : vector<4xf16> func.return %selected : vector<4xf16> } @@ -326,6 +349,7 @@ func.def inline @ggml_dequant_f32_vector4(%weight_format: index, %weight: buffer %c6 = index.constant 6 : index %c8 = index.constant 8 : index %c16 = index.constant 16 : index + %c17 = index.constant 17 : index %c32 = index.constant 32 : index %q8_0_format = index.constant 80 : index %q8_1_format = index.constant 81 : index @@ -335,6 +359,7 @@ func.def inline @ggml_dequant_f32_vector4(%weight_format: index, %weight: buffer %is_q8_0 = index.cmp eq, %weight_format, %q8_0_format : index %is_q8_1 = index.cmp eq, %weight_format, %q8_1_format : index %is_f16 = index.cmp eq, %weight_format, %c16 : index + %is_bf16 = index.cmp eq, %weight_format, %c17 : index %is_f32 = index.cmp eq, %weight_format, %c32 : index %q8_block_base = index.mul %quant_block, %c8 : index %q8_block = index.add %q8_block_base, %quant_group : index @@ -368,6 +393,12 @@ func.def inline @ggml_dequant_f32_vector4(%weight_format: index, %weight: buffer } else { scf.yield %zero : vector<4xf32> } + %bf16_values = scf.if %is_bf16 -> (vector<4xf32>) { + %values = func.call @ggml_bf16_f32_vector4(%weight, %row_byte_base, %input_size, %k) : (buffer, offset, index, index) -> (vector<4xf32>) + scf.yield %values : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } %f32_values = scf.if %is_f32 -> (vector<4xf32>) { %values = func.call @ggml_f32_f32_vector4(%weight, %row_byte_base, %input_size, %k) : (buffer, offset, index, index) -> (vector<4xf32>) scf.yield %values : vector<4xf32> @@ -379,6 +410,7 @@ func.def inline @ggml_dequant_f32_vector4(%weight_format: index, %weight: buffer %selected_q8_0 = scf.select %is_q8_0, %q8_0_values, %selected_q6 : vector<4xf32> %selected_q8_1 = scf.select %is_q8_1, %q8_1_values, %selected_q8_0 : vector<4xf32> %selected_f16 = scf.select %is_f16, %f16_values, %selected_q8_1 : vector<4xf32> - %selected = scf.select %is_f32, %f32_values, %selected_f16 : vector<4xf32> + %selected_bf16 = scf.select %is_bf16, %bf16_values, %selected_f16 : vector<4xf32> + %selected = scf.select %is_f32, %f32_values, %selected_bf16 : vector<4xf32> func.return %selected : vector<4xf32> } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/get_rows_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/get_rows_f32.loom index 5c32608fd174..1662ab13ecc0 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/get_rows_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/get_rows_f32.loom @@ -2,7 +2,7 @@ // // One workitem publishes four adjacent output channels for one requested row. // Storage-format-specific row decoding is delegated to the common dequant -// motif so Q4_K, Q6_K, Q8_0, Q8_1, F16, and F32 share the same gather shape. +// motif so Q4_K, Q6_K, Q8_0, Q8_1, F16, BF16, and F32 share the same gather shape. template.decl @ggml.get_rows_f32.body(%token_count: index, %row_count: index, %hidden_size: index, %token_ids: buffer, %weight: buffer, %output: buffer) template.decl @ggml.get_rows_f32.launch(%hidden_capacity: index, %token_capacity: index) -> (index, index, index, index) diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_f32_f32_decode.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_f32_f32_decode.loom index 74615ac7f928..79a9494af213 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_f32_f32_decode.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_f32_f32_decode.loom @@ -4,8 +4,8 @@ // of 16 lanes each process four 256-value weight blocks in parallel, while // each lane contracts four F32 activation packets from the 0, 32, 64, and 96 // element quarters of its block. The storage-format-specific weight load is -// delegated to the common dequant motif, so Q4_K, Q6_K, Q8_0, Q8_1, F16, and -// F32 share the same skinny decode schedule. +// delegated to the common dequant motif, so Q4_K, Q6_K, Q8_0, Q8_1, F16, +// BF16, and F32 share the same skinny decode schedule. template.decl @ggml.mul_mat_f32_f32_decode.body(%weight_format: index, %publish_output: i1, %token_count: index, %token0: index, %pair: index, %input_size: index, %output_size: index, %input: buffer, %weight: buffer, %output: buffer) template.decl @ggml.mul_mat_f32_f32_decode.dispatch(%weight_format: index, %token_count: index, %input_size: index, %output_size: index, %input: buffer, %weight: buffer, %output: buffer) diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_f32_f32_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_f32_f32_wmma.loom index c67b59d454e3..edf7bcf2aa5b 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_f32_f32_wmma.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_f32_f32_wmma.loom @@ -13,6 +13,7 @@ // Q8_0: [output channel][input size / 32][34 bytes] // Q8_1: [output channel][input size / 32][36 bytes] // F16: [output channel][input size][2 bytes] +// BF16: [output channel][input size][2 bytes] // F32: [output channel][input size][4 bytes] // The configured weight format selects the storage reader before the shared // device template is instantiated, so inactive decode logic is absent from the diff --git a/tests/hrx-backend-test.cpp b/tests/hrx-backend-test.cpp index 0b25eb0716f3..b12879ef8f3c 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -225,6 +225,8 @@ static int64_t matmul_weight_format_config(ggml_type type) { return 81; case GGML_TYPE_F16: return 16; + case GGML_TYPE_BF16: + return 17; case GGML_TYPE_F32: return 32; default: @@ -2423,14 +2425,14 @@ static void schedule_qwen_terminal_q6k_q8_command(int64_t token_count) { ggml_free(ctx); } -static void schedule_get_rows_q8_1_alternate_command() { +static void schedule_get_rows_q8_1_alternate_command(ggml_type embedding_weight_type) { 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 * embedding_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 151936); + ggml_tensor * embedding_weight = ggml_new_tensor_2d(ctx, embedding_weight_type, 2048, 151936); ggml_tensor * token_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); ggml_tensor * vocab_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, 2048, 151936); REQUIRE(embedding_weight != nullptr); @@ -2465,7 +2467,8 @@ static void schedule_get_rows_q8_1_alternate_command() { REQUIRE(get_rows_dispatch.kernel.integer_parameters.at("token_count") == 1); REQUIRE(get_rows_dispatch.kernel.integer_parameters.at("row_count") == 151936); REQUIRE(get_rows_dispatch.kernel.integer_parameters.at("hidden_size") == 2048); - require_compile_parameter(get_rows_dispatch, "ggml.get_rows_f32.weight_format", "4"); + require_compile_parameter(get_rows_dispatch, "ggml.get_rows_f32.weight_format", + std::to_string(matmul_weight_format_config(embedding_weight_type))); require_compile_parameter(get_rows_dispatch, "ggml.get_rows_f32.next_format", "81"); REQUIRE(get_rows_dispatch.bindings.size() == 4); @@ -2610,7 +2613,8 @@ static bool manual_token_embedding_graph_is_supported(ggml_context * ctx, static void schedule_qwen_token_embedding_command(ggml_context * ctx, ggml_tensor * output, int64_t expected_token_count, - int64_t expected_vocabulary_count) { + int64_t expected_vocabulary_count, + ggml_type expected_weight_type) { ggml_cgraph * graph = ggml_new_graph(ctx); REQUIRE(graph != nullptr); ggml_build_forward_expand(graph, output); @@ -2631,7 +2635,8 @@ static void schedule_qwen_token_embedding_command(ggml_context * ctx, REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == expected_token_count); REQUIRE(dispatch.kernel.integer_parameters.at("row_count") == expected_vocabulary_count); REQUIRE(dispatch.kernel.integer_parameters.at("hidden_size") == 2048); - require_compile_parameter(dispatch, "ggml.get_rows_f32.weight_format", "4"); + require_compile_parameter(dispatch, "ggml.get_rows_f32.weight_format", + std::to_string(matmul_weight_format_config(expected_weight_type))); REQUIRE(dispatch.bindings.size() == 3); const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( @@ -2659,7 +2664,7 @@ static void run_qwen_token_embedding_dispatch_checks() { REQUIRE(token_ids != nullptr); ggml_tensor * output = ggml_get_rows(ctx, weight, token_ids); REQUIRE(output != nullptr); - schedule_qwen_token_embedding_command(ctx, output, 1, 151936); + schedule_qwen_token_embedding_command(ctx, output, 1, 151936, GGML_TYPE_Q4_K); } { ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 151936); @@ -2668,7 +2673,16 @@ static void run_qwen_token_embedding_dispatch_checks() { REQUIRE(token_ids != nullptr); ggml_tensor * output = ggml_get_rows(ctx, weight, token_ids); REQUIRE(output != nullptr); - schedule_qwen_token_embedding_command(ctx, output, 13, 151936); + schedule_qwen_token_embedding_command(ctx, output, 13, 151936, GGML_TYPE_Q4_K); + } + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_BF16, 2048, 151936); + ggml_tensor * token_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 5); + REQUIRE(weight != nullptr); + REQUIRE(token_ids != nullptr); + ggml_tensor * output = ggml_get_rows(ctx, weight, token_ids); + REQUIRE(output != nullptr); + schedule_qwen_token_embedding_command(ctx, output, 5, 151936, GGML_TYPE_BF16); } REQUIRE( @@ -2681,6 +2695,8 @@ static void run_qwen_token_embedding_dispatch_checks() { manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_Q8_1, GGML_TYPE_I32, GGML_TYPE_F32, 2048, 151936, 1)); REQUIRE( manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_F16, GGML_TYPE_I32, GGML_TYPE_F32, 2048, 151936, 1)); + REQUIRE(manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_BF16, GGML_TYPE_I32, GGML_TYPE_F32, 2048, + 151936, 1)); REQUIRE( !manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_Q4_K, GGML_TYPE_I64, GGML_TYPE_F32, 2048, 151936, 1)); REQUIRE( @@ -3514,6 +3530,26 @@ static void run_qwen_matmul_dispatch_checks() { schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 2048, 128); } + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_BF16, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); + REQUIRE(weight != nullptr); + 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_f32_f32_wmma", 4, 2048, 128); + } + + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_BF16, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + REQUIRE(weight != nullptr); + 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_f32_f32_decode_wave64", 1, 2048, 128); + } + { ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 256); ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); @@ -3550,6 +3586,23 @@ static void run_qwen_matmul_dispatch_checks() { { "input", "weight", "bias", "output" }, false); } + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_BF16, 2048, 256); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 33); + ggml_tensor * bias = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 256); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + REQUIRE(bias != nullptr); + ggml_tensor * projection = ggml_mul_mat(ctx, weight, input); + REQUIRE(projection != nullptr); + ggml_tensor * output = ggml_add(ctx, projection, bias); + REQUIRE(output != nullptr); + schedule_fused_matmul_postops_command(ctx, output, "loom_libs:ggml_mul_mat_bias_f32_f32_wmma", + GGML_TYPE_BF16, 33, 2048, 256, 2, + { GGML_OP_MUL_MAT, GGML_OP_ADD }, + { "input", "weight", "bias", "output" }, false); + } + { ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 2048, 256); ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 33); @@ -3649,6 +3702,7 @@ static void run_qwen_matmul_dispatch_checks() { { GGML_TYPE_Q8_0, 128 }, { GGML_TYPE_Q8_1, 128 }, { GGML_TYPE_F16, 128 }, + { GGML_TYPE_BF16, 128 }, { GGML_TYPE_F32, 256 }, }; @@ -3987,81 +4041,84 @@ static void run_llama_attention_matmul_dispatch_checks() { constexpr int64_t token_count = 7; constexpr int64_t cache_row_count = 16; - { - ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, input_size, output_size); - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, input_size, token_count); - ggml_tensor * positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); - ggml_tensor * freqs = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, head_size / 2); - REQUIRE(weight != nullptr); - REQUIRE(input != nullptr); - REQUIRE(positions != nullptr); - REQUIRE(freqs != nullptr); - ggml_tensor * projection = ggml_mul_mat(ctx, weight, input); - ggml_tensor * query = ggml_reshape_3d(ctx, projection, head_size, head_count, token_count); - ggml_tensor * output = ggml_rope_ext(ctx, query, positions, freqs, head_size, GGML_ROPE_TYPE_NORMAL, 0, - 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); - REQUIRE(projection != nullptr); - REQUIRE(query != nullptr); - REQUIRE(output != nullptr); - schedule_fused_llama_attention_matmul_command( - ctx, output, "loom_libs:llm_attention_q_matmul_rope_f32_f32_wmma", GGML_TYPE_F16, 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" }); - } + const ggml_type weight_types[] = { GGML_TYPE_F16, GGML_TYPE_BF16 }; + for (ggml_type weight_type : weight_types) { + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, weight_type, input_size, output_size); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, input_size, token_count); + ggml_tensor * positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); + ggml_tensor * freqs = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, head_size / 2); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + REQUIRE(positions != nullptr); + REQUIRE(freqs != nullptr); + ggml_tensor * projection = ggml_mul_mat(ctx, weight, input); + ggml_tensor * query = ggml_reshape_3d(ctx, projection, head_size, head_count, token_count); + ggml_tensor * output = ggml_rope_ext(ctx, query, positions, freqs, head_size, GGML_ROPE_TYPE_NORMAL, 0, + 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + REQUIRE(projection != nullptr); + REQUIRE(query != nullptr); + REQUIRE(output != nullptr); + schedule_fused_llama_attention_matmul_command( + ctx, output, "loom_libs:llm_attention_q_matmul_rope_f32_f32_wmma", 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" }); + } - { - ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, input_size, output_size); - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, input_size, token_count); - ggml_tensor * positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); - ggml_tensor * freqs = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, head_size / 2); - ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 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(positions != nullptr); - REQUIRE(freqs != nullptr); - REQUIRE(cache != nullptr); - REQUIRE(indices != nullptr); - ggml_tensor * projection = ggml_mul_mat(ctx, weight, input); - ggml_tensor * key = ggml_reshape_3d(ctx, projection, head_size, head_count, token_count); - ggml_tensor * rope = ggml_rope_ext(ctx, key, positions, freqs, head_size, GGML_ROPE_TYPE_NORMAL, 0, 10000.0f, - 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); - ggml_tensor * rows = ggml_reshape_2d(ctx, rope, output_size, token_count); - ggml_tensor * output = ggml_set_rows(ctx, cache, rows, indices); - REQUIRE(projection != nullptr); - REQUIRE(key != nullptr); - REQUIRE(rope != nullptr); - REQUIRE(rows != nullptr); - REQUIRE(output != nullptr); - schedule_fused_llama_attention_matmul_command( - ctx, output, "loom_libs:llm_attention_k_matmul_rope_set_rows_f32_f32_wmma", GGML_TYPE_F16, token_count, - input_size, output_size, head_size, head_count, cache_row_count, 16, 5, - { GGML_OP_MUL_MAT, GGML_OP_RESHAPE, GGML_OP_ROPE, GGML_OP_RESHAPE, GGML_OP_SET_ROWS }, - { "input", "weight", "positions", "indices", "theta", "freq_factors", "cache" }); - } + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, weight_type, input_size, output_size); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, input_size, token_count); + ggml_tensor * positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); + ggml_tensor * freqs = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, head_size / 2); + ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 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(positions != nullptr); + REQUIRE(freqs != nullptr); + REQUIRE(cache != nullptr); + REQUIRE(indices != nullptr); + ggml_tensor * projection = ggml_mul_mat(ctx, weight, input); + ggml_tensor * key = ggml_reshape_3d(ctx, projection, head_size, head_count, token_count); + ggml_tensor * rope = ggml_rope_ext(ctx, key, positions, freqs, head_size, GGML_ROPE_TYPE_NORMAL, 0, 10000.0f, + 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + ggml_tensor * rows = ggml_reshape_2d(ctx, rope, output_size, token_count); + ggml_tensor * output = ggml_set_rows(ctx, cache, rows, indices); + REQUIRE(projection != nullptr); + REQUIRE(key != nullptr); + REQUIRE(rope != nullptr); + REQUIRE(rows != nullptr); + REQUIRE(output != nullptr); + schedule_fused_llama_attention_matmul_command( + ctx, output, "loom_libs:llm_attention_k_matmul_rope_set_rows_f32_f32_wmma", weight_type, token_count, + input_size, output_size, head_size, head_count, cache_row_count, 16, 5, + { GGML_OP_MUL_MAT, GGML_OP_RESHAPE, GGML_OP_ROPE, GGML_OP_RESHAPE, GGML_OP_SET_ROWS }, + { "input", "weight", "positions", "indices", "theta", "freq_factors", "cache" }); + } - { - ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, input_size, output_size); - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, input_size, token_count); - ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 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, head_size, head_count, 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); - schedule_fused_llama_attention_matmul_command( - ctx, output, "loom_libs:llm_attention_v_matmul_set_rows_f32_f32_wmma", GGML_TYPE_F16, token_count, - input_size, output_size, 0, 0, cache_row_count, 16, 4, - { GGML_OP_MUL_MAT, GGML_OP_RESHAPE, GGML_OP_RESHAPE, GGML_OP_SET_ROWS }, - { "input", "weight", "indices", "cache" }); + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, weight_type, input_size, output_size); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, input_size, token_count); + ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 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, head_size, head_count, 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); + schedule_fused_llama_attention_matmul_command( + ctx, output, "loom_libs:llm_attention_v_matmul_set_rows_f32_f32_wmma", weight_type, token_count, + input_size, output_size, 0, 0, cache_row_count, 16, 4, + { GGML_OP_MUL_MAT, GGML_OP_RESHAPE, GGML_OP_RESHAPE, GGML_OP_SET_ROWS }, + { "input", "weight", "indices", "cache" }); + } } ggml_free(ctx); @@ -4589,6 +4646,8 @@ static std::string common_mul_mat_weight_format(ggml_type type) { return "81"; case GGML_TYPE_F16: return "16"; + case GGML_TYPE_BF16: + return "17"; case GGML_TYPE_F32: return "32"; default: @@ -4778,6 +4837,11 @@ static void run_common_mul_mat_id_swiglu_dispatch_checks() { build_common_mul_mat_id_swiglu_graph(ctx, GGML_TYPE_F32, GGML_TYPE_Q8_1); require_common_mul_mat_id_swiglu_match(ctx, tensors, tensors.up); } + { + const CommonMulMatIdSwiGLUTensors tensors = + build_common_mul_mat_id_swiglu_graph(ctx, GGML_TYPE_BF16, GGML_TYPE_BF16); + require_common_mul_mat_id_swiglu_match(ctx, tensors, tensors.gate); + } { const CommonMulMatIdSwiGLUTensors tensors = build_common_mul_mat_id_swiglu_graph(ctx, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, GGML_GLU_OP_GEGLU); @@ -4819,7 +4883,8 @@ static int64_t common_mul_mat_id_partition_descriptor_capacity(const CommonMulMa static CommonMulMatIdPostOpsTensors build_common_mul_mat_id_postops_graph(ggml_context * ctx, bool include_bias, bool include_residual, - bool include_rmsnorm) { + bool include_rmsnorm, + ggml_type weight_type = GGML_TYPE_Q4_K) { CommonMulMatIdPostOpsTensors tensors; constexpr int64_t token_count = 4; constexpr int64_t route_count = 8; @@ -4828,7 +4893,7 @@ static CommonMulMatIdPostOpsTensors build_common_mul_mat_id_postops_graph(ggml_c constexpr int64_t expert_count = 16; tensors.route_ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, route_count, token_count); tensors.input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, input_size, 1, token_count); - tensors.weight = ggml_new_tensor_3d(ctx, GGML_TYPE_Q4_K, input_size, output_size, expert_count); + tensors.weight = ggml_new_tensor_3d(ctx, weight_type, input_size, output_size, expert_count); REQUIRE(tensors.route_ids != nullptr); REQUIRE(tensors.input != nullptr); REQUIRE(tensors.weight != nullptr); @@ -4981,6 +5046,12 @@ static void run_common_mul_mat_id_postops_dispatch_checks() { require_common_mul_mat_id_postops_match(ctx, tensors, "loom_libs:ggml_mul_mat_id_postops_f32_f32_wmma", true, true, false); } + { + const CommonMulMatIdPostOpsTensors tensors = + build_common_mul_mat_id_postops_graph(ctx, true, true, false, GGML_TYPE_BF16); + require_common_mul_mat_id_postops_match(ctx, tensors, "loom_libs:ggml_mul_mat_id_postops_f32_f32_wmma", true, + true, false); + } { const CommonMulMatIdPostOpsTensors tensors = build_common_mul_mat_id_postops_graph(ctx, false, true, true); require_common_mul_mat_id_postops_match( @@ -5095,6 +5166,25 @@ static void run_qwen_routed_gate_up_dispatch_checks() { REQUIRE(commands.commands[3].bindings[4].name == "output"); } + { + constexpr int64_t token_count = 4; + const QwenRoutedGateUpTensors tensors = build_qwen_routed_gate_up_graph( + ctx, token_count, GGML_GLU_OP_SWIGLU, GGML_TYPE_Q4_K, true, GGML_TYPE_BF16); + ggml::hrx::GraphImportResult imported = import_qwen_routed_gate_up_graph(ctx, tensors); + std::vector covered_nodes(imported.graph.nodes().size(), false); + ggml::hrx::CommandPlan plan = build_qwen_router_plan_for_graph(imported.graph, covered_nodes); + ggml::hrx::DispatchMatch down_match = require_common_mul_mat_id_for_graph( + imported.graph, plan, covered_nodes, tensors.output, tensors.glu, tensors.down_weight, tensors.route_ids); + append_match_to_plan(plan, down_match, covered_nodes, &imported.graph); + + REQUIRE(plan.dispatches.size() == 4); + const ggml::hrx::CommandProgram commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.size() == 4); + REQUIRE(command_program_verifies(commands)); + } + { constexpr int64_t token_count = 4; const QwenRoutedGateUpTensors tensors = build_qwen_routed_gate_up_graph(ctx, token_count); @@ -7114,7 +7204,8 @@ int main() { run_llama_attention_matmul_dispatch_checks(); schedule_qwen_terminal_q6k_q8_command(1); schedule_qwen_terminal_q6k_q8_command(18); - schedule_get_rows_q8_1_alternate_command(); + schedule_get_rows_q8_1_alternate_command(GGML_TYPE_Q4_K); + schedule_get_rows_q8_1_alternate_command(GGML_TYPE_BF16); run_qwen_router_top8_dispatch_checks(); run_common_mul_mat_id_swiglu_dispatch_checks(); run_qwen_routed_gate_up_dispatch_checks(); diff --git a/tests/test-hrx-ops.cpp b/tests/test-hrx-ops.cpp index 14e5b6bca6bb..064361042068 100644 --- a/tests/test-hrx-ops.cpp +++ b/tests/test-hrx-ops.cpp @@ -576,6 +576,15 @@ static std::vector make_matmul_weight_bytes(ggml_type type, int64_t row return bytes; } + if (type == GGML_TYPE_BF16) { + (void) seed; + const std::vector weights(static_cast(row_length * row_count), + ggml_fp32_to_bf16(0.25f)); + std::vector bytes(weights.size() * sizeof(ggml_bf16_t)); + std::memcpy(bytes.data(), weights.data(), bytes.size()); + return bytes; + } + return make_quantized_rows(type, row_length, row_count, seed); } @@ -1964,11 +1973,14 @@ static void run_dense_matmul_cpu_reference_case(ggml_type weight_type, REQUIRE(cpu_buffer != nullptr); REQUIRE(hrx_buffer != nullptr); - const std::vector weight = make_matmul_weight_bytes(weight_type, input_size, output_size, 6); + const std::vector weight = make_matmul_weight_bytes(weight_type, input_size, output_size, 6); + const bool weight_is_f16 = weight_type == GGML_TYPE_F16; + 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; const std::vector input = - weight_type == GGML_TYPE_F16 || weight_type == GGML_TYPE_F32 ? - std::vector(static_cast(input_size * token_count), 0.00390625f) : - make_pattern_f32(input_size * token_count, 7, 0.01f); + weight_is_dense_float ? std::vector(static_cast(input_size * token_count), 0.00390625f) : + make_pattern_f32(input_size * token_count, 7, 0.01f); set_tensor_pair_bytes(cpu_backend, cpu_weight, hrx_backend, hrx_weight, weight.data(), weight.size()); set_tensor_pair_bytes(cpu_backend, cpu_input, hrx_backend, hrx_input, input.data(), input.size() * sizeof(float)); @@ -3030,6 +3042,7 @@ int main() { run_get_rows_f32_cpu_reference_case(GGML_TYPE_Q8_0); run_get_rows_q8_1_zero_weight_case(); run_get_rows_f32_cpu_reference_case(GGML_TYPE_F16); + run_get_rows_f32_cpu_reference_case(GGML_TYPE_BF16); run_get_rows_f32_cpu_reference_case(GGML_TYPE_F32); run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 128); @@ -3041,6 +3054,8 @@ int main() { run_dense_matmul_zero_weight_case(GGML_TYPE_Q8_1, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1); run_dense_matmul_cpu_reference_case(GGML_TYPE_F16, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); run_dense_matmul_cpu_reference_case(GGML_TYPE_F16, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 128); + run_dense_matmul_cpu_reference_case(GGML_TYPE_BF16, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); + run_dense_matmul_cpu_reference_case(GGML_TYPE_BF16, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 128); run_dense_matmul_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 256); run_dense_matmul_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 256); run_dense_matmul_unary_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128);