Fix gated FlashAttention benchmark bindings

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",

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",

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<x> : index
+  %workitem = kernel.workitem.id<x> : 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) {

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

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<gfx11-generic> @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
+}

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",

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()