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 b6d35c974bac..741bd7c110f1 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rmsnorm.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rmsnorm.cpp @@ -21,6 +21,8 @@ namespace ggml::hrx { namespace { static constexpr KernelCatalogRef kRmsNormBinaryF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_rmsnorm_binary_f32"); +static constexpr KernelCatalogRef kRmsNormBinaryF32LowTokenStripedKernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_rmsnorm_binary_f32_lowtoken_striped"); static constexpr KernelCatalogRef kRmsNormBinaryStridedF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_rmsnorm_binary_strided_f32"); static constexpr KernelCatalogRef kRmsNormMulAddF32Kernel = @@ -974,10 +976,14 @@ static bool match_rmsnorm_binary_f32_dispatch(const DispatchMatchContext & conte publication_demand->format == CommonActivationPublicationFormat::F16K16Major; const bool f16_output = publication_demand != nullptr && publication_demand->format == CommonActivationPublicationFormat::F16Row; + const bool lowtoken_striped = !strided_input && publication_demand == nullptr && rms_match.op == BinaryKind::Mul && + rms_match.token_count < 6 && rms_match.hidden_size >= 2048 && + rms_match.hidden_size <= 5120; Dispatch dispatch; dispatch.kernel = make_kernel_specialization(strided_input ? kRmsNormBinaryStridedF32Kernel : packed_output ? kRmsNormBinaryF32K16Kernel : f16_output ? kRmsNormBinaryF32F16Kernel : + lowtoken_striped ? kRmsNormBinaryF32LowTokenStripedKernel : kRmsNormBinaryF32Kernel); dispatch.kernel.integer_parameters.emplace("token_count", rms_match.token_count); dispatch.kernel.compile_parameters.emplace("ggml.rmsnorm_binary_f32.hidden_size", 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 bf1928273c30..625075d1fbc9 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 @@ -71,6 +71,9 @@ { "path": "ops/normalization/rmsnorm_binary_f32.loom" }, + { + "path": "ops/normalization/rmsnorm_binary_f32_lowtoken_striped.loom" + }, { "path": "ops/normalization/rmsnorm_binary_strided_f32.loom" }, @@ -1797,6 +1800,51 @@ "motifs/elementwise/binary_apply.loom" ] }, + { + "name": "ggml_rmsnorm_binary_f32_lowtoken_striped", + "family": "loom_libs", + "symbol": "ggml_rmsnorm_binary_f32_lowtoken_striped", + "source": "ops/normalization/rmsnorm_binary_f32_lowtoken_striped.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_f32_lowtoken_striped.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_binary_strided_f32", "family": "loom_libs", 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 f3ee58daa77f..ace56e1c8440 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 @@ -25,6 +25,8 @@ config.decl @ggml.rmsnorm_binary_q8_1_x4.op : %value: index where [range(%value, 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_f32.lowtoken_striped_body(%token_count: index, %input: buffer, %rhs: buffer, %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) { @@ -59,6 +61,43 @@ template.def<@ggml.rmsnorm_binary_f32.body> device @ggml_rmsnorm_binary_f32_body template.return } +template.def<@ggml.rmsnorm_binary_f32.lowtoken_striped_body> device @ggml_rmsnorm_binary_f32_lowtoken_striped_body(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer) { + %hidden_size0 = config.get @ggml.rmsnorm_binary_f32.hidden_size : index + %epsilon = config.get @ggml.rmsnorm_binary_f32.rms_epsilon : f32 + %op = config.get @ggml.rmsnorm_binary_f32.op : index + %hidden_size = index.assume %hidden_size0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128)] : index + %group = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + %c1023 = index.constant 1023 : index + %c1024 = index.constant 1024 : index + %c0_offset = index.constant 0 : offset + %rounded_hidden = index.add %hidden_size, %c1023 : index + %stripe_count = index.div %rounded_hidden, %c1024 : index + %token0 = index.div %group, %stripe_count : index + %stripe = index.rem %group, %stripe_count : index + %stripe_base = index.mul %stripe, %c1024 : index + %stripe_limit0 = index.add %stripe_base, %c1024 : index + %stripe_limit = index.min %stripe_limit0, %hidden_size : 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, %hidden_size, %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_size]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 { + %first_channel = index.add %stripe_base, %workitem : index + scf.for %channel = [%first_channel to %stripe_limit step %c256] { + %value = view.load %input_view[%token, %channel] : view<[%launch_token_count]x[%hidden_size]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 +} + 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) { diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_swiglu_f32_f32_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_swiglu_f32_f32_wmma.loom index 36a5a15cc629..f81e0f3f6593 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_swiglu_f32_f32_wmma.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_swiglu_f32_f32_wmma.loom @@ -413,12 +413,19 @@ func.def inline @ggml_mul_mat_swiglu_lowtoken_dot(%op: index, %gate_format: inde %c2 = index.constant 2 : index %c4 = index.constant 4 : index %c8 = index.constant 8 : index + %c10 = index.constant 10 : index %c256 = index.constant 256 : index %q5_format = index.constant 5 : index %iq4xs_format = index.constant 23 : index + %q8_0_format = index.constant 80 : index %gate_is_q5 = index.cmp eq, %gate_format, %q5_format : index %gate_is_iq4xs = index.cmp eq, %gate_format, %iq4xs_format : index %same_weight_format = index.cmp eq, %gate_format, %up_format : index + %gate_is_q8_0 = index.cmp eq, %gate_format, %q8_0_format : index + %q8_0_pair = scalar.andi %gate_is_q8_0, %same_weight_format : i1 + %token_one = index.cmp eq, %token_count, %c1 : index + %q8_0_token_one = scalar.andi %q8_0_pair, %token_one : i1 + %fallback_unroll = scf.select %q8_0_token_one, %c10, %c1 : index %base = index.constant 0 : offset %zero = vector.constant 0.0 : vector<4xf32> %zero_scalar = scalar.constant 0.0 : f32 @@ -512,7 +519,7 @@ func.def inline @ggml_mul_mat_swiglu_lowtoken_dot(%op: index, %gate_format: inde } scf.yield %iq4_total_g0, %iq4_total_u0, %iq4_total_g1, %iq4_total_u1, %iq4_total_g2, %iq4_total_u2, %iq4_total_g3, %iq4_total_u3, %iq4_total_g4, %iq4_total_u4 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> } else { - %fallback_total_g0, %fallback_total_u0, %fallback_total_g1, %fallback_total_u1, %fallback_total_g2, %fallback_total_u2, %fallback_total_g3, %fallback_total_u3, %fallback_total_g4, %fallback_total_u4 = scf.for %step = [%c0 to %steps step %c1](%acc_g0 = %zero : vector<4xf32>, %acc_u0 = %zero : vector<4xf32>, %acc_g1 = %zero : vector<4xf32>, %acc_u1 = %zero : vector<4xf32>, %acc_g2 = %zero : vector<4xf32>, %acc_u2 = %zero : vector<4xf32>, %acc_g3 = %zero : vector<4xf32>, %acc_u3 = %zero : vector<4xf32>, %acc_g4 = %zero : vector<4xf32>, %acc_u4 = %zero : vector<4xf32>) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) { + %fallback_total_g0, %fallback_total_u0, %fallback_total_g1, %fallback_total_u1, %fallback_total_g2, %fallback_total_u2, %fallback_total_g3, %fallback_total_u3, %fallback_total_g4, %fallback_total_u4 = scf.for %step = [%c0 to %steps step %c1](%acc_g0 = %zero : vector<4xf32>, %acc_u0 = %zero : vector<4xf32>, %acc_g1 = %zero : vector<4xf32>, %acc_u1 = %zero : vector<4xf32>, %acc_g2 = %zero : vector<4xf32>, %acc_u2 = %zero : vector<4xf32>, %acc_g3 = %zero : vector<4xf32>, %acc_u3 = %zero : vector<4xf32>, %acc_g4 = %zero : vector<4xf32>, %acc_u4 = %zero : vector<4xf32>) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) unroll(%fallback_unroll) { %block = index.div %step, %c4 : index %quarter = index.rem %step, %c4 : index %quarter_group = index.mul %quarter, %c2 : index diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_binary_f32_lowtoken_striped.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_binary_f32_lowtoken_striped.loom new file mode 100644 index 000000000000..655cd00ea03b --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_binary_f32_lowtoken_striped.loom @@ -0,0 +1,27 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.rmsnorm_binary_f32.lowtoken_striped_body(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer) + +amdgpu.target @ggml_rmsnorm_binary_lowtoken_striped_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.rms_epsilon : f32 + +config.decl @ggml.rmsnorm_binary_f32.op : %value: index where [range(%value, 0, 3)] + +kernel.def target(@ggml_rmsnorm_binary_lowtoken_striped_gfx11_wave32) export("ggml_rmsnorm_binary_f32_lowtoken_striped") @ggml_rmsnorm_binary_f32_lowtoken_striped(%token_count: index) { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + %c1023 = index.constant 1023 : index + %c1024 = index.constant 1024 : index + %hidden_size = config.get @ggml.rmsnorm_binary_f32.hidden_size : index + %rounded_hidden = index.add %hidden_size, %c1023 : index + %stripe_count = index.div %rounded_hidden, %c1024 : index + %group_count = index.mul %token_count, %stripe_count : index + kernel.launch.config workgroups(%group_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer) where [range(%token_count, 1, 5)] { + template.apply<@ggml.rmsnorm_binary_f32.lowtoken_striped_body>(%token_count, %input, %rhs, %output) : (index, buffer, buffer, buffer) + kernel.return +} 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 a3c1ca06f630..4b9084678f91 100755 --- a/ggml/src/ggml-hrx/tools/benchmarks/generate-model-benchmarks.py +++ b/ggml/src/ggml-hrx/tools/benchmarks/generate-model-benchmarks.py @@ -427,11 +427,33 @@ def case_flash_attention(symbol: str, command: dict[str, Any], kernel: str, pref output_shape = f"{query_tokens}x{query_heads}x{value_head_size}" mask_shape = f"{query_tokens}x{kv_tokens}" if prefix == "ggml": + gate_bytes = binding_length(command, "gate") + if gate_bytes % 4 != 0: + fail( + f"command {command.get('ordinal')} {command.get('kernel')} has " + f"non-f32 gate binding length {gate_bytes}" + ) + gate_elements = gate_bytes // 4 + apply_gate = config_int(command, "ggml.flash_attention.apply_gate", 0) + if apply_gate: + gate_stride_head = config_int(command, "ggml.flash_attention.gate_stride_head") + gate_stride_token = config_int(command, "ggml.flash_attention.gate_stride_token") + required_gate_elements = ( + (query_tokens - 1) * gate_stride_token + + (query_heads - 1) * gate_stride_head + + value_head_size + ) + if gate_elements < required_gate_elements: + fail( + f"command {command.get('ordinal')} {command.get('kernel')} gate binding " + f"has {gate_elements} f32 elements but the strided access requires " + f"{required_gate_elements}" + ) launch = ( f" kernel.launch @{kernel}[%query_token_count, %key_value_token_count]" - f"(%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %query, %output)" + f"(%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %gate, %output)" f" : [index, index](index, index, tensor<{query_shape}xf32>, tensor<{key_shape}xf16>, " - f"tensor<{value_shape}xf16>, tensor<{mask_shape}xf16>, tensor<{query_shape}xf32>, " + f"tensor<{value_shape}xf16>, tensor<{mask_shape}xf16>, tensor<{gate_elements}xf32>, " f"tensor<{output_shape}xf32>)" ) else: @@ -450,6 +472,7 @@ def case_flash_attention(symbol: str, command: dict[str, Any], kernel: str, pref fill_tensor("key", "0.0", key_shape, "f16"), fill_tensor("value", "0.0", value_shape, "f16"), fill_tensor("mask", "0.0", mask_shape, "f16"), + *([fill_tensor("gate", "0.0", f"{gate_elements}", "f32")] if prefix == "ggml" else []), fill_tensor("output", "1.0", output_shape, "f32"), launch, " check.return", diff --git a/ggml/src/ggml-hrx/tools/benchmarks/test-generate-model-benchmarks.py b/ggml/src/ggml-hrx/tools/benchmarks/test-generate-model-benchmarks.py new file mode 100644 index 000000000000..097adb476524 --- /dev/null +++ b/ggml/src/ggml-hrx/tools/benchmarks/test-generate-model-benchmarks.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python3 + +from __future__ import annotations + +import importlib.util +import sys +import unittest +from pathlib import Path + + +SCRIPT = Path(__file__).with_name("generate-model-benchmarks.py") +SPEC = importlib.util.spec_from_file_location("generate_model_benchmarks", SCRIPT) +if SPEC is None or SPEC.loader is None: + raise RuntimeError(f"could not load {SCRIPT}") +MODULE = importlib.util.module_from_spec(SPEC) +sys.modules[SPEC.name] = MODULE +SPEC.loader.exec_module(MODULE) + + +def gated_flash_attention_command(gate_bytes: int = 48128) -> dict: + return { + "ordinal": 59, + "kernel": "loom_libs:ggml_flash_attention_f32_f16_wmma", + "integer_parameters": { + "query_token_count": 1, + "key_value_token_count": 256, + }, + "compile_parameters": { + "ggml.flash_attention.apply_gate": "1", + "ggml.flash_attention.query_head_count": "24", + "ggml.flash_attention.key_value_head_count": "4", + "ggml.flash_attention.qk_head_size": "256", + "ggml.flash_attention.value_head_size": "256", + "ggml.flash_attention.gate_stride_head": "512", + "ggml.flash_attention.gate_stride_token": "12288", + }, + "bindings": [ + {"name": "query", "length": 24576}, + {"name": "key", "length": 524288}, + {"name": "value", "length": 524288}, + {"name": "mask", "length": 512}, + {"name": "gate", "length": gate_bytes}, + {"name": "output", "length": 24576}, + ], + } + + +class FlashAttentionCaseTest(unittest.TestCase): + def test_materializes_distinct_gate_with_exact_binding_capacity(self) -> None: + case = MODULE.case_flash_attention( + "qwen38_tg8_flash_attention", + gated_flash_attention_command(), + "ggml_flash_attention_f32_f16_wmma", + "ggml", + ) + + self.assertIn("%gate = check.generate.fill value(0.0) : tensor<12032xf32>", case) + self.assertIn("%mask, %gate, %output)", case) + self.assertIn("tensor<12032xf32>", case) + self.assertNotIn("%mask, %query, %output)", case) + + def test_rejects_gate_binding_smaller_than_strided_access(self) -> None: + with self.assertRaisesRegex(SystemExit, "strided access requires 12032"): + MODULE.case_flash_attention( + "qwen38_tg8_flash_attention", + gated_flash_attention_command(gate_bytes=24576), + "ggml_flash_attention_f32_f16_wmma", + "ggml", + ) + + def test_rejects_gate_binding_not_aligned_to_f32(self) -> None: + with self.assertRaisesRegex(SystemExit, "non-f32 gate binding length"): + MODULE.case_flash_attention( + "qwen38_tg8_flash_attention", + gated_flash_attention_command(gate_bytes=48127), + "ggml_flash_attention_f32_f16_wmma", + "ggml", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/hrx-backend-test.cpp b/tests/hrx-backend-test.cpp index 2c2f930b7f10..cfe17209aef0 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -1654,6 +1654,50 @@ static ggml_tensor * make_test_binary(ggml_context * ctx, } } +static const char * expected_plain_rmsnorm_binary_kernel(int64_t hidden_size, int64_t token_count) { + return token_count < 6 && hidden_size >= 2048 && hidden_size <= 5120 ? + "loom_libs:ggml_rmsnorm_binary_f32_lowtoken_striped" : + "loom_libs:ggml_rmsnorm_binary_f32"; +} + +static void run_rmsnorm_lowtoken_striped_dispatch_checks() { + const struct { + int64_t hidden_size; + int64_t token_count; + ggml::hrx::BinaryKind op; + const char * expected_kernel; + } cases[] = { + { 1024, 1, ggml::hrx::BinaryKind::Mul, "loom_libs:ggml_rmsnorm_binary_f32" }, + { 2048, 1, ggml::hrx::BinaryKind::Mul, "loom_libs:ggml_rmsnorm_binary_f32_lowtoken_striped" }, + { 3072, 2, ggml::hrx::BinaryKind::Mul, "loom_libs:ggml_rmsnorm_binary_f32_lowtoken_striped" }, + { 5120, 4, ggml::hrx::BinaryKind::Mul, "loom_libs:ggml_rmsnorm_binary_f32_lowtoken_striped" }, + { 5120, 6, ggml::hrx::BinaryKind::Mul, "loom_libs:ggml_rmsnorm_binary_f32" }, + { 2048, 1, ggml::hrx::BinaryKind::Add, "loom_libs:ggml_rmsnorm_binary_f32" }, + }; + + for (const auto & test : cases) { + ggml_init_params params = {}; + params.mem_size = static_cast(test.hidden_size * test.token_count * sizeof(float) * 8 + 1024 * 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, test.hidden_size, test.token_count); + ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, test.hidden_size); + ggml_tensor * rms = ggml_rms_norm(ctx, input, 1.0e-6f); + ggml_tensor * output = make_test_binary(ctx, rms, weight, test.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) == test.expected_kernel); + REQUIRE(dispatch.bindings.size() == 3); + require_compile_parameter(dispatch, "ggml.rmsnorm_binary_f32.hidden_size", + std::to_string(test.hidden_size)); + require_compile_parameter(dispatch, "ggml.rmsnorm_binary_f32.op", + std::to_string(ggml::hrx::binary_kind_config_value(test.op))); + ggml_free(ctx); + } +} + static void run_rmsnorm_two_binary_dispatch_checks() { const struct { ggml::hrx::BinaryKind normalized_op; @@ -3706,7 +3750,8 @@ static void schedule_fused_matmul_postops_command(ggml_context * REQUIRE(dispatch.kernel.compile_parameters.count("ggml.mul_mat_postops.rms_epsilon") == 0); if (expected_rmsnorm) { const ggml::hrx::Dispatch & rmsnorm_dispatch = scheduler.plan().dispatches[1]; - REQUIRE(kernel_name_for_id(rmsnorm_dispatch.kernel.kernel_id) == "loom_libs:ggml_rmsnorm_binary_f32"); + REQUIRE(kernel_name_for_id(rmsnorm_dispatch.kernel.kernel_id) == + expected_plain_rmsnorm_binary_kernel(expected_output_size, expected_token_count)); } const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( @@ -13710,6 +13755,7 @@ static void register_hrx_backend_host_cases(test_runner::Suite & suite) { 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("rmsnorm_lowtoken_striped_dispatch", [] { run_rmsnorm_lowtoken_striped_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("binary_q8_publication", [] { run_binary_q8_publication_checks(); }); diff --git a/tests/test-hrx-loom-jit.cpp b/tests/test-hrx-loom-jit.cpp index ee5272eaef95..d90ae6a4ee98 100644 --- a/tests/test-hrx-loom-jit.cpp +++ b/tests/test-hrx-loom-jit.cpp @@ -1195,6 +1195,14 @@ static void run_prefill_fusion_compile_checks() { { "ggml.rmsnorm_binary_f32.op", "2" }, })); + const auto & rmsnorm_lowtoken_striped = find_kernel("ggml_rmsnorm_binary_f32_lowtoken_striped"); + require_compiled_kernel(compile_kernel(*jit, "rmsnorm-lowtoken-striped", rmsnorm_lowtoken_striped, + token_count_workload(1), { + { "ggml.rmsnorm_binary_f32.hidden_size", "5120" }, + { "ggml.rmsnorm_binary_f32.rms_epsilon", "0.000001" }, + { "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), { diff --git a/tests/test-hrx-ops.cpp b/tests/test-hrx-ops.cpp index 1eddae831ead..ff745578a939 100644 --- a/tests/test-hrx-ops.cpp +++ b/tests/test-hrx-ops.cpp @@ -103,6 +103,12 @@ static std::vector make_weight(int64_t hidden_size) { return data; } +static const char * expected_plain_rmsnorm_binary_kernel(int64_t hidden_size, int64_t token_count) { + return token_count < 6 && hidden_size >= 2048 && hidden_size <= 5120 ? + "loom_libs:ggml_rmsnorm_binary_f32_lowtoken_striped" : + "loom_libs:ggml_rmsnorm_binary_f32"; +} + static std::vector make_router_input(int64_t hidden_size, int64_t token_count) { std::vector data(hidden_size * token_count); for (int64_t i = 0; i < static_cast(data.size()); ++i) { @@ -1637,7 +1643,8 @@ static void run_rmsnorm_mul_case(int64_t hidden_size, int64_t token_count) { REQUIRE(graph != nullptr); ggml_build_forward_expand(graph, output); - require_kernel_subsequence(scheduled_kernel_sequence(graph), { "loom_libs:ggml_rmsnorm_binary_f32" }); + require_kernel_subsequence(scheduled_kernel_sequence(graph), + { expected_plain_rmsnorm_binary_kernel(hidden_size, token_count) }); ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); REQUIRE(buffer != nullptr); @@ -2217,7 +2224,10 @@ static void run_strided_cont_cpu_reference_case(int64_t token_count) { ggml_backend_free(hrx_backend); } -static void run_rmsnorm_mul_cpu_reference_case(int64_t hidden_size, int64_t token_count, float epsilon) { +static void run_rmsnorm_mul_cpu_reference_case(int64_t hidden_size, + int64_t token_count, + float epsilon, + bool adversarial = false) { ggml_backend_t cpu_backend = init_cpu_backend(); ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); REQUIRE(hrx_backend != nullptr); @@ -2249,14 +2259,29 @@ static void run_rmsnorm_mul_cpu_reference_case(int64_t hidden_size, int64_t toke 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_f32" }); + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), + { expected_plain_rmsnorm_binary_kernel(hidden_size, token_count) }); 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, 23, 0.02f); + std::vector input = make_pattern_f32(hidden_size * token_count, 23, 0.02f); + if (adversarial) { + for (size_t i = 0; i < input.size(); ++i) { + switch (i % 8) { + case 0: input[i] = 0.0f; break; + case 1: input[i] = 1.0e-6f; break; + case 2: input[i] = -1.0e-6f; break; + case 3: input[i] = 8.0f; break; + case 4: input[i] = -8.0f; break; + case 5: input[i] = 0.125f; break; + case 6: input[i] = -0.25f; break; + default: input[i] = 1.5f; break; + } + } + } 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_weight, hrx_backend, hrx_weight, weight.data(), @@ -6389,6 +6414,21 @@ static void register_rmsnorm_and_scheduling_cases(Suite & suite) { suite.device_case("rmsnorm_gate.silu.hidden128.tokens48.stride512", [] { run_rmsnorm_gate_cpu_reference_case(128, 48, 512, GGML_UNARY_OP_SILU); }); suite.device_case("rmsnorm_mul.hidden3840.tokens18", [] { run_rmsnorm_mul_cpu_reference_case(3840, 18, 1.0e-6f); }); + for (const int64_t hidden_size : { 1024, 2048, 2560, 3072, 4096, 5120 }) { + suite.device_case("rmsnorm_mul.lowtoken.random.hidden" + std::to_string(hidden_size) + ".tokens1", + [hidden_size] { run_rmsnorm_mul_cpu_reference_case(hidden_size, 1, 1.0e-6f); }); + suite.device_case("rmsnorm_mul.lowtoken.adversarial.hidden" + std::to_string(hidden_size) + ".tokens1", + [hidden_size] { run_rmsnorm_mul_cpu_reference_case(hidden_size, 1, 1.0e-5f, true); }); + } + for (const int64_t hidden_size : { 2048, 5120 }) { + for (const int64_t token_count : { 2, 4 }) { + suite.device_case("rmsnorm_mul.lowtoken.random.hidden" + std::to_string(hidden_size) + ".tokens" + + std::to_string(token_count), + [hidden_size, token_count] { + run_rmsnorm_mul_cpu_reference_case(hidden_size, token_count, 1.0e-6f); + }); + } + } suite.device_case("gemma_scaled_rmsnorm.cpu_reference", [] { run_gemma_scaled_rmsnorm_cpu_reference_case(); }); suite.device_case("gemma_scaled_rmsnorm_mul.cpu_reference", [] { run_gemma_scaled_rmsnorm_mul_cpu_reference_case(); });