sync 91d990483a17

Browse files

Files changed (10) hide show

README.md +13 -9
build/webgpu/bench.json +336 -1
build/webgpu/manifest.json +165 -240
build/webgpu/metadata.json +21 -11
build/webgpu/norm-row-stats.wgsl.jinja +142 -13
build/webgpu/rms-normalization-splitk-normalize.wgsl.jinja +23 -21
build/webgpu/rms-normalization-splitk-partials.wgsl.jinja +0 -3
build/webgpu/rms-normalization-stash-f16-serial.wgsl.jinja +8 -8
build/webgpu/rms-normalization.wgsl.jinja +8 -11
build/webgpu/test.json +550 -4

README.md CHANGED Viewed

@@ -12,22 +12,22 @@ tags:
 ## Description
-Computes RMS normalization over the suffix dimensions of `X` starting at `axis`: `Y = X / sqrt(mean(X^2) + epsilon) * scale`. The normalization stage supports TensorProto `stash_type` values `1` (float32) and `10` (float16), and is cast back to the dtype of `X` before `scale` is applied. The input type `T` and scale/output type `V` may independently be float16 or float32; ONNX's bfloat16 and double cases are not yet implemented.
 See the [ONNX `RMSNormalization` spec](https://onnx.ai/onnx/operators/onnx__RMSNormalization.html) for the reference semantics.
 ## Inputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `X` | `x` | `T` | — | — | Input tensor to be normalized; the RMS is taken over the last dimensions starting at `axis`. | required |
-| `scale` | `scale` | `V` | — | — | Scale tensor, unidirectionally broadcastable to `X`; its dtype `V` may differ from the input dtype `T`. | required |
 ## Outputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `Y` | `y` | `V` | same as `X` | same as `X` | Normalized and scaled output tensor; same shape as `X` and same dtype `V` as `scale`. | required |
 ## Attributes
@@ -48,7 +48,7 @@ Default values (overridable per request):
 ## Files
-- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
 - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
@@ -60,10 +60,14 @@ Default values (overridable per request):
 ## Use with `@huggingface/kernels`
-The loader derives every required output's shape and logical dtype from the manifest contract and this call.
-It then allocates the result tensors automatically.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

 ## Description
+Computes RMS normalization over the suffix dimensions of `X` starting at `axis`: `Y = X / sqrt(mean(X^2) + epsilon) * scale`. The normalization stage supports TensorProto `stash_type` values `1` (float32) and `10` (float16), and is cast back to the dtype of `X` before `scale` is applied. The input type `T` and scale/output type `V` may independently be float16 or float32; ONNX bfloat16 and double cases are unsupported.
 See the [ONNX `RMSNormalization` spec](https://onnx.ai/onnx/operators/onnx__RMSNormalization.html) for the reference semantics.
 ## Inputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `x` | `X` | `T` | — | — | Input tensor to be normalized; the RMS is taken over the last dimensions starting at `axis`. | required |
+| `scale` | — | `V` | — | — | Scale tensor, unidirectionally broadcastable to `X`; its dtype `V` may differ from the input dtype `T`. | required |
 ## Outputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `y` | `Y` | `V` | same as `x` | same as `x` | Normalized and scaled output tensor; same shape as `X` and same dtype `V` as `scale`. | required |
 ## Attributes
 ## Files
+- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
 - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 ## Use with `@huggingface/kernels`
+```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.2
+```
+Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

build/webgpu/bench.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "ai.onnx.RMSNormalization",
   "cases": [
     {
       "name": "rmsnorm-f32-256x1024",
@@ -168,6 +167,342 @@
       "bench": {
         "metrics": [{ "type": "bandwidth", "value": "(args.rows * args.dim * 2 + args.dim) * dtypeBytes(args.dtype)" }]
       }
     }
   ]
 }

 {
   "cases": [
     {
       "name": "rmsnorm-f32-256x1024",
       "bench": {
         "metrics": [{ "type": "bandwidth", "value": "(args.rows * args.dim * 2 + args.dim) * dtypeBytes(args.dtype)" }]
       }
+    },
+    {
+      "name": "split_f32-1x16384",
+      "preset": "stress",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": { "shape": [1, 16384], "dtype": "float32", "dist": "normal", "seed": 8123, "scale": 0.5 },
+        "scale": { "shape": [], "dtype": "float32", "dist": "uniform", "seed": 8124, "scale": 0.25, "offset": 1 }
+      },
+      "outputs": { "y": { "shape": [1, 16384], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 131076 }] }
+    },
+    {
+      "name": "split_f32-1x16385",
+      "preset": "stress",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": { "shape": [1, 16385], "dtype": "float32", "dist": "normal", "seed": 8123, "scale": 0.5 },
+        "scale": { "shape": [], "dtype": "float32", "dist": "uniform", "seed": 8124, "scale": 0.25, "offset": 1 }
+      },
+      "outputs": { "y": { "shape": [1, 16385], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 131084 }] }
+    },
+    {
+      "name": "split_f32-1x32769",
+      "preset": "stress",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": { "shape": [1, 32769], "dtype": "float32", "dist": "normal", "seed": 8123, "scale": 0.5 },
+        "scale": { "shape": [], "dtype": "float32", "dist": "uniform", "seed": 8124, "scale": 0.25, "offset": 1 }
+      },
+      "outputs": { "y": { "shape": [1, 32769], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 262156 }] }
+    },
+    {
+      "name": "split_f32-1x65536",
+      "preset": "stress",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": { "shape": [1, 65536], "dtype": "float32", "dist": "normal", "seed": 8123, "scale": 0.5 },
+        "scale": { "shape": [], "dtype": "float32", "dist": "uniform", "seed": 8124, "scale": 0.25, "offset": 1 }
+      },
+      "outputs": { "y": { "shape": [1, 65536], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 524292 }] }
+    },
+    {
+      "name": "split_f32-1x131072",
+      "preset": "stress",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": { "shape": [1, 131072], "dtype": "float32", "dist": "normal", "seed": 8123, "scale": 0.5 },
+        "scale": { "shape": [], "dtype": "float32", "dist": "uniform", "seed": 8124, "scale": 0.25, "offset": 1 }
+      },
+      "outputs": { "y": { "shape": [1, 131072], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 1048580 }] }
+    },
+    {
+      "name": "split_f32-1x262144",
+      "preset": "stress",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": { "shape": [1, 262144], "dtype": "float32", "dist": "normal", "seed": 8123, "scale": 0.5 },
+        "scale": { "shape": [], "dtype": "float32", "dist": "uniform", "seed": 8124, "scale": 0.25, "offset": 1 }
+      },
+      "outputs": { "y": { "shape": [1, 262144], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 2097156 }] }
+    },
+    {
+      "name": "split_f32-1x524288",
+      "preset": "stress",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": { "shape": [1, 524288], "dtype": "float32", "dist": "normal", "seed": 8123, "scale": 0.5 },
+        "scale": { "shape": [], "dtype": "float32", "dist": "uniform", "seed": 8124, "scale": 0.25, "offset": 1 }
+      },
+      "outputs": { "y": { "shape": [1, 524288], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 4194308 }] }
+    },
+    {
+      "name": "split_f32-1x1048576",
+      "preset": "stress",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": { "shape": [1, 1048576], "dtype": "float32", "dist": "normal", "seed": 8123, "scale": 0.5 },
+        "scale": { "shape": [], "dtype": "float32", "dist": "uniform", "seed": 8124, "scale": 0.25, "offset": 1 }
+      },
+      "outputs": { "y": { "shape": [1, 1048576], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 8388612 }] }
+    },
+    {
+      "name": "split_f32-1x2097152",
+      "preset": "stress",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": { "shape": [1, 2097152], "dtype": "float32", "dist": "normal", "seed": 8123, "scale": 0.5 },
+        "scale": { "shape": [], "dtype": "float32", "dist": "uniform", "seed": 8124, "scale": 0.25, "offset": 1 }
+      },
+      "outputs": { "y": { "shape": [1, 2097152], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 16777220 }] }
+    },
+    {
+      "name": "split_f32-2x131073",
+      "preset": "stress",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": { "shape": [2, 131073], "dtype": "float32", "dist": "normal", "seed": 8123, "scale": 0.5 },
+        "scale": { "shape": [], "dtype": "float32", "dist": "uniform", "seed": 8124, "scale": 0.25, "offset": 1 }
+      },
+      "outputs": { "y": { "shape": [2, 131073], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 2097172 }] }
+    },
+    {
+      "name": "split_f32-3x524289",
+      "preset": "stress",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": { "shape": [3, 524289], "dtype": "float32", "dist": "normal", "seed": 8123, "scale": 0.5 },
+        "scale": { "shape": [], "dtype": "float32", "dist": "uniform", "seed": 8124, "scale": 0.25, "offset": 1 }
+      },
+      "outputs": { "y": { "shape": [3, 524289], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 12582940 }] }
+    },
+    {
+      "name": "split_f32-4x1048576",
+      "preset": "stress",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": { "shape": [4, 1048576], "dtype": "float32", "dist": "normal", "seed": 8123, "scale": 0.5 },
+        "scale": { "shape": [], "dtype": "float32", "dist": "uniform", "seed": 8124, "scale": 0.25, "offset": 1 }
+      },
+      "outputs": { "y": { "shape": [4, 1048576], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 33554436 }] }
+    },
+    {
+      "name": "split_f32-8x16384",
+      "preset": "stress",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": { "shape": [8, 16384], "dtype": "float32", "dist": "normal", "seed": 8123, "scale": 0.5 },
+        "scale": { "shape": [], "dtype": "float32", "dist": "uniform", "seed": 8124, "scale": 0.25, "offset": 1 }
+      },
+      "outputs": { "y": { "shape": [8, 16384], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 1048580 }] }
+    },
+    {
+      "name": "split_f32-16x65536",
+      "preset": "stress",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": { "shape": [16, 65536], "dtype": "float32", "dist": "normal", "seed": 8123, "scale": 0.5 },
+        "scale": { "shape": [], "dtype": "float32", "dist": "uniform", "seed": 8124, "scale": 0.25, "offset": 1 }
+      },
+      "outputs": { "y": { "shape": [16, 65536], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 8388612 }] }
+    },
+    {
+      "name": "split_f32-32x32769",
+      "preset": "stress",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": { "shape": [32, 32769], "dtype": "float32", "dist": "normal", "seed": 8123, "scale": 0.5 },
+        "scale": { "shape": [], "dtype": "float32", "dist": "uniform", "seed": 8124, "scale": 0.25, "offset": 1 }
+      },
+      "outputs": { "y": { "shape": [32, 32769], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 8388868 }] }
+    },
+    {
+      "name": "split_f32-128x16384",
+      "preset": "stress",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": { "shape": [128, 16384], "dtype": "float32", "dist": "normal", "seed": 8123, "scale": 0.5 },
+        "scale": { "shape": [], "dtype": "float32", "dist": "uniform", "seed": 8124, "scale": 0.25, "offset": 1 }
+      },
+      "outputs": { "y": { "shape": [128, 16384], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 16777220 }] }
+    },
+    {
+      "name": "split_f16-1x16385",
+      "preset": "stress",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": { "shape": [1, 16385], "dtype": "float16", "dist": "normal", "seed": 8123, "scale": 0.5 },
+        "scale": { "shape": [], "dtype": "float16", "dist": "uniform", "seed": 8124, "scale": 0.25, "offset": 1 }
+      },
+      "outputs": { "y": { "shape": [1, 16385], "dtype": "float16" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 65542 }] }
+    },
+    {
+      "name": "split_f16-1x131072",
+      "preset": "stress",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": { "shape": [1, 131072], "dtype": "float16", "dist": "normal", "seed": 8123, "scale": 0.5 },
+        "scale": { "shape": [], "dtype": "float16", "dist": "uniform", "seed": 8124, "scale": 0.25, "offset": 1 }
+      },
+      "outputs": { "y": { "shape": [1, 131072], "dtype": "float16" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 524290 }] }
+    },
+    {
+      "name": "split_f16-1x524288",
+      "preset": "stress",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": { "shape": [1, 524288], "dtype": "float16", "dist": "normal", "seed": 8123, "scale": 0.5 },
+        "scale": { "shape": [], "dtype": "float16", "dist": "uniform", "seed": 8124, "scale": 0.25, "offset": 1 }
+      },
+      "outputs": { "y": { "shape": [1, 524288], "dtype": "float16" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 2097154 }] }
+    },
+    {
+      "name": "split_f16-1x2097152",
+      "preset": "stress",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": { "shape": [1, 2097152], "dtype": "float16", "dist": "normal", "seed": 8123, "scale": 0.5 },
+        "scale": { "shape": [], "dtype": "float16", "dist": "uniform", "seed": 8124, "scale": 0.25, "offset": 1 }
+      },
+      "outputs": { "y": { "shape": [1, 2097152], "dtype": "float16" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 8388610 }] }
+    },
+    {
+      "name": "split_f16-3x524289",
+      "preset": "stress",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": { "shape": [3, 524289], "dtype": "float16", "dist": "normal", "seed": 8123, "scale": 0.5 },
+        "scale": { "shape": [], "dtype": "float16", "dist": "uniform", "seed": 8124, "scale": 0.25, "offset": 1 }
+      },
+      "outputs": { "y": { "shape": [3, 524289], "dtype": "float16" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 6291470 }] }
+    },
+    {
+      "name": "split_f16-16x65536",
+      "preset": "stress",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": { "shape": [16, 65536], "dtype": "float16", "dist": "normal", "seed": 8123, "scale": 0.5 },
+        "scale": { "shape": [], "dtype": "float16", "dist": "uniform", "seed": 8124, "scale": 0.25, "offset": 1 }
+      },
+      "outputs": { "y": { "shape": [16, 65536], "dtype": "float16" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 4194306 }] }
+    },
+    {
+      "name": "split_f32-f16-2x32769",
+      "preset": "stress",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": { "shape": [2, 32769], "dtype": "float32", "dist": "normal", "seed": 8123, "scale": 0.5 },
+        "scale": { "shape": [], "dtype": "float16", "dist": "uniform", "seed": 8124, "scale": 0.25, "offset": 1 }
+      },
+      "outputs": { "y": { "shape": [2, 32769], "dtype": "float16" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 393230 }] }
+    },
+    {
+      "name": "split_f16-f32-2x32769",
+      "preset": "stress",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": { "shape": [2, 32769], "dtype": "float16", "dist": "normal", "seed": 8123, "scale": 0.5 },
+        "scale": { "shape": [], "dtype": "float32", "dist": "uniform", "seed": 8124, "scale": 0.25, "offset": 1 }
+      },
+      "outputs": { "y": { "shape": [2, 32769], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 393232 }] }
+    },
+    {
+      "name": "split_f32-1x16384-split1",
+      "preset": "stress",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": { "MAX_SPLITS": 1 },
+      "inputs": {
+        "x": { "shape": [1, 16384], "dtype": "float32", "dist": "normal", "seed": 8123, "scale": 0.5 },
+        "scale": { "shape": [], "dtype": "float32", "dist": "uniform", "seed": 8124, "scale": 0.25, "offset": 1 }
+      },
+      "outputs": { "y": { "shape": [1, 16384], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 131076 }] }
+    },
+    {
+      "name": "split_f32-1x16385-split3",
+      "preset": "stress",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": { "MAX_SPLITS": 3 },
+      "inputs": {
+        "x": { "shape": [1, 16385], "dtype": "float32", "dist": "normal", "seed": 8123, "scale": 0.5 },
+        "scale": { "shape": [], "dtype": "float32", "dist": "uniform", "seed": 8124, "scale": 0.25, "offset": 1 }
+      },
+      "outputs": { "y": { "shape": [1, 16385], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 131084 }] }
+    },
+    {
+      "name": "split_f32-2x262145-wg64",
+      "preset": "stress",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": { "WORKGROUP_SIZE": 64 },
+      "inputs": {
+        "x": { "shape": [2, 262145], "dtype": "float32", "dist": "normal", "seed": 8123, "scale": 0.5 },
+        "scale": { "shape": [], "dtype": "float32", "dist": "uniform", "seed": 8124, "scale": 0.25, "offset": 1 }
+      },
+      "outputs": { "y": { "shape": [2, 262145], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 4194324 }] }
+    },
+    {
+      "name": "split_f32-2x262145-wg128",
+      "preset": "stress",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": { "WORKGROUP_SIZE": 128 },
+      "inputs": {
+        "x": { "shape": [2, 262145], "dtype": "float32", "dist": "normal", "seed": 8123, "scale": 0.5 },
+        "scale": { "shape": [], "dtype": "float32", "dist": "uniform", "seed": 8124, "scale": 0.25, "offset": 1 }
+      },
+      "outputs": { "y": { "shape": [2, 262145], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 4194324 }] }
     }
   ]
 }

build/webgpu/manifest.json CHANGED Viewed

@@ -2,47 +2,17 @@
   "domain": "ai.onnx",
   "name": "RMSNormalization",
   "sinceVersion": 23,
-  "description": "Computes RMS normalization over the suffix dimensions of `X` starting at `axis`: `Y = X / sqrt(mean(X^2) + epsilon) * scale`. The normalization stage supports TensorProto `stash_type` values `1` (float32) and `10` (float16), and is cast back to the dtype of `X` before `scale` is applied. The input type `T` and scale/output type `V` may independently be float16 or float32; ONNX's bfloat16 and double cases are not yet implemented.",
-  "inputs": [
-    {
-      "role": "X",
-      "dtype": "T",
-      "description": "Input tensor to be normalized; the RMS is taken over the last dimensions starting at `axis`."
-    },
-    {
-      "role": "scale",
-      "dtype": "V",
-      "description": "Scale tensor, unidirectionally broadcastable to `X`; its dtype `V` may differ from the input dtype `T`."
-    }
-  ],
-  "outputs": [
-    {
-      "role": "Y",
-      "dtype": "V",
-      "rank": "ranks.X",
-      "shape": "shapes.X",
-      "description": "Normalized and scaled output tensor; same shape as `X` and same dtype `V` as `scale`."
-    }
-  ],
-  "attributes": { "axis": -1, "epsilon": 0.00001, "stash_type": 1 },
-  "attributeDescriptions": {
-    "axis": "The first dimension of the normalization suffix; negative values count from the end, so the default `-1` normalizes over only the last dimension.",
-    "epsilon": "Small constant added to the mean square before taking the square root to avoid division by zero.",
-    "stash_type": "TensorProto element type used for normalization: `1` computes in float32, while `10` computes in float16."
-  },
   "attributeConstraints": { "stash_type": { "values": [1, 10] } },
   "typeConstraints": { "T": ["float32", "float16"], "V": ["float32", "float16"] },
-  "args": {
-    "x": { "kind": "tensor", "semantic": "X", "role": "input" },
-    "scale": { "kind": "tensor", "semantic": "scale", "role": "input" },
-    "y": { "kind": "tensor", "semantic": "Y", "role": "output" }
-  },
   "tunables": {
-    "WORKGROUP_SIZE": 256,
-    "SPLIT_MAX_ROWS": 256,
-    "SPLIT_MIN_HIDDEN": 16384,
-    "SPLIT_TARGET_ELEMENTS": 4096,
-    "MAX_SPLITS": 64
   },
   "derive": {
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
@@ -50,118 +20,51 @@
     "reportedNonWave32Adapter": "not wave32Adapter and (has(device.adapterInfo, \"subgroupMinSize\") or has(device.adapterInfo, \"subgroupMaxSize\"))",
     "normMaxWorkgroup": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
     "hasSubgroupId": "device.features.has(\"subgroups\") and device.wgslLanguageFeatures.has(\"subgroup_id\")",
-    "axisNorm": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.X",
-    "normalizedRows": "outer(shapes.X, axisNorm)",
-    "normalizedHidden": "dim(shapes.X, axisNorm) * inner(shapes.X, axisNorm)",
     "normalizedDispatchRows": "0 if normalizedHidden == 0 else normalizedRows",
     "normalizedWorkgroupHidden": "max(1, normalizedHidden)",
-    "normalizationShapeOk": "ranks.X >= 1 and ranks.scale >= 0 and ranks.scale <= ranks.X and sameShape(shapes.Y, shapes.X) and attrs.axis + ranks.X >= 0 and attrs.axis < ranks.X and broadcastable(shapes.scale, shapes.X) and f16Ok(dtypes.T) and f16Ok(dtypes.V)",
     "baseOk": "normalizationShapeOk and attrs.stash_type == onnxDtypeCode(\"float32\")",
     "stashF16Ok": "normalizationShapeOk and attrs.stash_type == onnxDtypeCode(\"float16\")",
-    "lastAxisOk": "baseOk and (attrs.axis == -1 or attrs.axis == ranks.X - 1)",
-    "suffixAxisOk": "baseOk and ranks.X >= 2 and not (attrs.axis == -1 or attrs.axis == ranks.X - 1)"
   },
-  "bindingSets": {
-    "rows": [
-      {
-        "name": "x",
-        "arg": "x",
-        "semantic": "X",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$xElement"
-      },
-      {
-        "name": "scale",
-        "arg": "scale",
-        "semantic": "scale",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$ioElement"
-      },
-      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$ioElement" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "normalizedRows" },
-            {
-              "name": "rowStride",
-              "type": "u32",
-              "value": "max(1, min(normalizedRows, device.limits.maxComputeWorkgroupsPerDimension))"
-            }
-          ]
-        }
-      }
-    ],
-    "splitPartials": [
-      {
-        "name": "x",
-        "arg": "x",
-        "semantic": "X",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$xElement"
-      },
-      { "name": "partials", "semantic": "partials", "buffer": { "type": "storage" }, "elementType": "f32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "splitRows" },
-            {
-              "name": "rowStride",
-              "type": "u32",
-              "value": "max(1, min(splitRows, device.limits.maxComputeWorkgroupsPerDimension))"
-            }
-          ]
         }
-      }
-    ],
-    "splitNormalize": [
-      {
-        "name": "x",
-        "arg": "x",
-        "semantic": "X",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$xElement"
-      },
-      {
-        "name": "scale",
-        "arg": "scale",
-        "semantic": "scale",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$ioElement"
-      },
-      { "name": "partials", "semantic": "partials", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
-      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$ioElement" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "splitRows" },
-            {
-              "name": "rowStride",
-              "type": "u32",
-              "value": "max(1, min(splitRows, device.limits.maxComputeWorkgroupsPerDimension))"
-            }
-          ]
         }
-      }
-    ]
   },
   "variants": [
     {
       "id": "stash_f16_serial",
       "priority": 1000,
-      "when": "stashF16Ok",
-      "constants": {
         "scalar": "dtypes.V",
         "xElement": "dtypes.T",
         "ioElement": "dtypes.V",
@@ -173,73 +76,87 @@
         {
           "id": "main",
           "name": "RMSNormalization.StashF16Serial",
-          "source": {
-            "shader": "rms-normalization-stash-f16-serial.wgsl.jinja",
-            "inputs": {
-              "xShape": "shapes.X",
-              "scaleShape": "shapes.scale",
-              "xRank": "ranks.X",
-              "scaleRank": "ranks.scale"
-            }
           },
-          "bindings": "rows",
-          "dispatch": { "workgroups": "normalizedDispatchRows" }
         }
       ]
     },
     {
       "id": "suffix_axis_splitk",
       "priority": 15,
       "derive": {
         "splitRows": "normalizedRows",
         "splitHidden": "normalizedHidden",
-        "split": "min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(splitHidden, tunables.SPLIT_TARGET_ELEMENTS)))"
-      },
-      "when": ["baseOk", "ranks.X >= 2", "normalizedRows <= tunables.SPLIT_MAX_ROWS", "normalizedHidden >= tunables.SPLIT_MIN_HIDDEN", "min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(normalizedHidden, tunables.SPLIT_TARGET_ELEMENTS))) <= device.limits.maxComputeWorkgroupsPerDimension", "normalizedRows * min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(normalizedHidden, tunables.SPLIT_TARGET_ELEMENTS))) * 4 <= device.limits.maxStorageBufferBindingSize", "normalizedRows * min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(normalizedHidden, tunables.SPLIT_TARGET_ELEMENTS))) * 4 <= device.limits.maxBufferSize"],
-      "demoteWhen": ["reportedNonWave32Adapter"],
-      "constants": {
         "scalar": "dtypes.V",
         "xElement": "dtypes.T",
         "ioElement": "dtypes.V",
         "usesF16": "dtypes.T == \"f16\" or dtypes.V == \"f16\"",
         "hiddenSize": "splitHidden",
         "workgroupSize": "normMaxWorkgroup",
-        "split": "split",
-        "epsilon": "attrs.epsilon"
       },
       "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[splitRows * split]" }],
       "passes": [
         {
           "id": "partials",
           "name": "RMSNormalization.SplitKPartials",
-          "source": { "shader": "rms-normalization-splitk-partials.wgsl.jinja", "inputs": {} },
-          "bindings": "splitPartials",
-          "dispatch": { "workgroups": "splitRows", "z": "split" }
         },
         {
           "id": "normalize",
           "name": "RMSNormalization.SplitKNormalize",
-          "source": {
-            "shader": "rms-normalization-splitk-normalize.wgsl.jinja",
-            "inputs": {
-              "xShape": "shapes.X",
-              "scaleShape": "shapes.scale",
-              "xRank": "ranks.X",
-              "scaleRank": "ranks.scale",
-              "writeStats": false,
-              "rmsScaleAfterCast": true
-            }
           },
-          "bindings": "splitNormalize",
-          "dispatch": { "workgroups": "splitRows", "z": "split" }
         }
       ]
     },
     {
       "id": "last_axis",
       "priority": 0,
-      "when": "lastAxisOk",
-      "constants": {
         "scalar": "dtypes.V",
         "xElement": "dtypes.T",
         "usesF16": "dtypes.T == \"f16\" or dtypes.V == \"f16\"",
@@ -252,27 +169,29 @@
         {
           "id": "main",
           "name": "RMSNormalization",
-          "source": {
-            "shader": "rms-normalization.wgsl.jinja",
-            "inputs": {
-              "xShape": "shapes.X",
-              "scaleShape": "shapes.scale",
-              "xRank": "ranks.X",
-              "scaleRank": "ranks.scale",
-              "writeStats": false,
-              "rmsScaleAfterCast": true
-            }
           },
-          "bindings": "rows",
-          "dispatch": { "workgroups": "normalizedDispatchRows" }
         }
       ]
     },
     {
       "id": "suffix_axis",
       "priority": 10,
-      "when": "suffixAxisOk",
-      "constants": {
         "scalar": "dtypes.V",
         "xElement": "dtypes.T",
         "usesF16": "dtypes.T == \"f16\" or dtypes.V == \"f16\"",
@@ -285,83 +204,89 @@
         {
           "id": "main",
           "name": "RMSNormalization.SuffixAxis",
-          "source": {
-            "shader": "rms-normalization.wgsl.jinja",
-            "inputs": {
-              "xShape": "shapes.X",
-              "scaleShape": "shapes.scale",
-              "xRank": "ranks.X",
-              "scaleRank": "ranks.scale",
-              "writeStats": false,
-              "rmsScaleAfterCast": true
-            }
           },
-          "bindings": "rows",
-          "dispatch": { "workgroups": "normalizedDispatchRows" }
         }
       ]
     },
     {
       "id": "last_axis_row_vec4",
       "priority": 110,
-      "when": ["lastAxisOk", "dtypes.T == dtypes.V", "ranks.scale >= 1", "numel(shapes.scale) == dim(shapes.X, -1)", "dim(shapes.scale, -1) == dim(shapes.X, -1)", "dim(shapes.X, -1) % 4 == 0"],
-      "constants": { "xElement": "\"vec4<\" ~ dtypes.T ~ \">\"", "ioElement": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "RMSNormalization.LastAxisRow",
-          "source": {
-            "shader": "norm-row-stats.wgsl.jinja",
-            "inputs": {
-              "mode": "\"rms\"",
-              "vec4": true,
-              "writeStats": false,
-              "rmsScaleAfterCast": true,
-              "scalar": "dtypes.T",
-              "usesF16": "dtypes.T == \"f16\"",
-              "hidden": "dim(shapes.X, -1)",
-              "wg": "min(normMaxWorkgroup, pow2ceil(max(1, dim(shapes.X, -1) / 4)))",
-              "epsilon": "attrs.epsilon",
-              "hiddenVec": "dim(shapes.X, -1) / 4",
-              "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
-              "combineSubgroups": "hasSubgroupId"
-            }
           },
-          "subgroupCollectivesWidth": "portable",
-          "bindings": "rows",
-          "dispatch": { "workgroups": "normalizedDispatchRows" }
         }
       ]
     },
     {
       "id": "last_axis_row",
       "priority": 100,
-      "when": ["lastAxisOk", "dtypes.T == dtypes.V", "ranks.scale >= 1", "numel(shapes.scale) == dim(shapes.X, -1)", "dim(shapes.scale, -1) == dim(shapes.X, -1)", "true"],
-      "constants": { "xElement": "dtypes.T", "ioElement": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "RMSNormalization.LastAxisRow",
-          "source": {
-            "shader": "norm-row-stats.wgsl.jinja",
-            "inputs": {
-              "mode": "\"rms\"",
-              "vec4": false,
-              "writeStats": false,
-              "rmsScaleAfterCast": true,
-              "scalar": "dtypes.T",
-              "usesF16": "dtypes.T == \"f16\"",
-              "hidden": "dim(shapes.X, -1)",
-              "wg": "min(normMaxWorkgroup, pow2ceil(max(1, dim(shapes.X, -1))))",
-              "epsilon": "attrs.epsilon",
-              "hiddenVec": 1,
-              "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
-              "combineSubgroups": "hasSubgroupId"
-            }
           },
-          "subgroupCollectivesWidth": "portable",
-          "bindings": "rows",
-          "dispatch": { "workgroups": "normalizedDispatchRows" }
         }
       ]
     }

   "domain": "ai.onnx",
   "name": "RMSNormalization",
   "sinceVersion": 23,
+  "inputs": { "x": { "onnx": "X", "dtype": "T" }, "scale": { "dtype": "V" } },
+  "outputs": { "y": { "onnx": "Y", "dtype": "V", "rank": "ranks.x", "shape": "shapes.x" } },
+  "attributes": { "axis": { "default": -1 }, "epsilon": { "default": 0.00001 }, "stash_type": { "default": 1 } },
   "attributeConstraints": { "stash_type": { "values": [1, 10] } },
   "typeConstraints": { "T": ["float32", "float16"], "V": ["float32", "float16"] },
   "tunables": {
+    "WORKGROUP_SIZE": { "default": 256 },
+    "SPLIT_MAX_ROWS": { "default": 256 },
+    "SPLIT_MIN_HIDDEN": { "default": 16384 },
+    "SPLIT_TARGET_ELEMENTS": { "default": 4096 },
+    "MAX_SPLITS": { "default": 64 }
   },
   "derive": {
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
     "reportedNonWave32Adapter": "not wave32Adapter and (has(device.adapterInfo, \"subgroupMinSize\") or has(device.adapterInfo, \"subgroupMaxSize\"))",
     "normMaxWorkgroup": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
     "hasSubgroupId": "device.features.has(\"subgroups\") and device.wgslLanguageFeatures.has(\"subgroup_id\")",
+    "axisNorm": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x",
+    "normalizedRows": "outer(shapes.x, axisNorm)",
+    "normalizedHidden": "dim(shapes.x, axisNorm) * inner(shapes.x, axisNorm)",
     "normalizedDispatchRows": "0 if normalizedHidden == 0 else normalizedRows",
     "normalizedWorkgroupHidden": "max(1, normalizedHidden)",
+    "normalizationShapeOk": "ranks.x >= 1 and ranks.scale >= 0 and ranks.scale <= ranks.x and sameShape(shapes.y, shapes.x) and attrs.axis + ranks.x >= 0 and attrs.axis < ranks.x and broadcastable(shapes.scale, shapes.x) and f16Ok(dtypes.T) and f16Ok(dtypes.V)",
     "baseOk": "normalizationShapeOk and attrs.stash_type == onnxDtypeCode(\"float32\")",
     "stashF16Ok": "normalizationShapeOk and attrs.stash_type == onnxDtypeCode(\"float16\")",
+    "lastAxisOk": "baseOk and (attrs.axis == -1 or attrs.axis == ranks.x - 1)",
+    "suffixAxisOk": "baseOk and ranks.x >= 2 and not (attrs.axis == -1 or attrs.axis == ranks.x - 1)"
   },
+  "bindings": {
+    "x": { "buffer": "read-only-storage", "elementType": "$xElement" },
+    "scale": { "buffer": "read-only-storage", "elementType": "$ioElement" },
+    "y": { "buffer": "storage", "elementType": "$ioElement" },
+    "params": {
+      "buffer": "uniform",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "normalizedRows" },
+        {
+          "name": "rowStride",
+          "type": "u32",
+          "value": "max(1, min(normalizedRows, min(device.limits.maxComputeWorkgroupsPerDimension, 65535)))"
         }
+      ]
+    },
+    "params_2": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "splitRows" },
+        {
+          "name": "rowStride",
+          "type": "u32",
+          "value": "max(1, min(splitRows, min(device.limits.maxComputeWorkgroupsPerDimension, 65535)))"
         }
+      ]
+    }
   },
   "variants": [
     {
       "id": "stash_f16_serial",
       "priority": 1000,
+      "when": ["stashF16Ok"],
+      "derive": {
         "scalar": "dtypes.V",
         "xElement": "dtypes.T",
         "ioElement": "dtypes.V",
         {
           "id": "main",
           "name": "RMSNormalization.StashF16Serial",
+          "shader": "rms-normalization-stash-f16-serial.wgsl.jinja",
+          "derive": {
+            "xShape": "shapes.x",
+            "scaleShape": "shapes.scale",
+            "xRank": "ranks.x",
+            "scaleRank": "ranks.scale"
           },
+          "bindings": ["x", "scale", "y", "params"],
+          "dispatch": {
+            "x": "min(normalizedDispatchRows, 65535)",
+            "y": "ceilDiv(normalizedDispatchRows, 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "suffix_axis_splitk",
       "priority": 15,
+      "when": ["baseOk", "ranks.x >= 2", "normalizedRows <= tunables.SPLIT_MAX_ROWS", "normalizedHidden >= tunables.SPLIT_MIN_HIDDEN", "min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(normalizedHidden, tunables.SPLIT_TARGET_ELEMENTS))) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "normalizedRows * min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(normalizedHidden, tunables.SPLIT_TARGET_ELEMENTS))) * 4 <= device.limits.maxStorageBufferBindingSize", "normalizedRows * min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(normalizedHidden, tunables.SPLIT_TARGET_ELEMENTS))) * 4 <= device.limits.maxBufferSize"],
+      "demoteWhen": ["reportedNonWave32Adapter"],
       "derive": {
         "splitRows": "normalizedRows",
         "splitHidden": "normalizedHidden",
+        "split": "min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(splitHidden, tunables.SPLIT_TARGET_ELEMENTS)))",
         "scalar": "dtypes.V",
         "xElement": "dtypes.T",
         "ioElement": "dtypes.V",
         "usesF16": "dtypes.T == \"f16\" or dtypes.V == \"f16\"",
         "hiddenSize": "splitHidden",
         "workgroupSize": "normMaxWorkgroup",
+        "epsilon": "attrs.epsilon",
+        "normalizeRows": "splitRows"
       },
       "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[splitRows * split]" }],
       "passes": [
         {
           "id": "partials",
           "name": "RMSNormalization.SplitKPartials",
+          "shader": "rms-normalization-splitk-partials.wgsl.jinja",
+          "bindings": ["x", { "name": "partials", "buffer": "storage", "elementType": "f32" }, "params_2"],
+          "dispatch": {
+            "x": "min(splitRows, DISPATCH_FOLD_WIDTH)",
+            "y": "ceilDiv(splitRows, DISPATCH_FOLD_WIDTH)",
+            "z": "split"
+          }
         },
         {
           "id": "normalize",
           "name": "RMSNormalization.SplitKNormalize",
+          "shader": "rms-normalization-splitk-normalize.wgsl.jinja",
+          "derive": {
+            "xShape": "shapes.x",
+            "scaleShape": "shapes.scale",
+            "xRank": "ranks.x",
+            "scaleRank": "ranks.scale",
+            "writeStats": false,
+            "rmsScaleAfterCast": true,
+            "normalizeBlocks": "max(split, min(min(device.limits.maxComputeWorkgroupsPerDimension, 65535), ceilDiv(workgroupSize, max(1, normalizeRows)), ceilDiv(hiddenSize, workgroupSize * 4)))",
+            "normalizeChunk": "ceilDiv(hiddenSize, normalizeBlocks)"
           },
+          "bindings": [
+            "x",
+            "scale",
+            { "name": "partials", "buffer": "read-only-storage", "elementType": "f32" },
+            "y",
+            "params_2"
+          ],
+          "dispatch": {
+            "x": "min(splitRows, DISPATCH_FOLD_WIDTH)",
+            "y": "ceilDiv(splitRows, DISPATCH_FOLD_WIDTH)",
+            "z": "normalizeBlocks"
+          }
         }
       ]
     },
     {
       "id": "last_axis",
       "priority": 0,
+      "when": ["lastAxisOk"],
+      "derive": {
         "scalar": "dtypes.V",
         "xElement": "dtypes.T",
         "usesF16": "dtypes.T == \"f16\" or dtypes.V == \"f16\"",
         {
           "id": "main",
           "name": "RMSNormalization",
+          "shader": "rms-normalization.wgsl.jinja",
+          "derive": {
+            "xShape": "shapes.x",
+            "scaleShape": "shapes.scale",
+            "xRank": "ranks.x",
+            "scaleRank": "ranks.scale",
+            "writeStats": false,
+            "rmsScaleAfterCast": true
           },
+          "bindings": ["x", "scale", "y", "params"],
+          "dispatch": {
+            "x": "min(normalizedDispatchRows, 65535)",
+            "y": "ceilDiv(normalizedDispatchRows, 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "suffix_axis",
       "priority": 10,
+      "when": ["suffixAxisOk"],
+      "derive": {
         "scalar": "dtypes.V",
         "xElement": "dtypes.T",
         "usesF16": "dtypes.T == \"f16\" or dtypes.V == \"f16\"",
         {
           "id": "main",
           "name": "RMSNormalization.SuffixAxis",
+          "shader": "rms-normalization.wgsl.jinja",
+          "derive": {
+            "xShape": "shapes.x",
+            "scaleShape": "shapes.scale",
+            "xRank": "ranks.x",
+            "scaleRank": "ranks.scale",
+            "writeStats": false,
+            "rmsScaleAfterCast": true
           },
+          "bindings": ["x", "scale", "y", "params"],
+          "dispatch": {
+            "x": "min(normalizedDispatchRows, 65535)",
+            "y": "ceilDiv(normalizedDispatchRows, 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "last_axis_row_vec4",
       "priority": 110,
+      "when": ["lastAxisOk", "dtypes.T == dtypes.V", "ranks.scale >= 1", "numel(shapes.scale) == dim(shapes.x, -1)", "dim(shapes.scale, -1) == dim(shapes.x, -1)", "dim(shapes.x, -1) % 4 == 0"],
+      "derive": { "xElement": "\"vec4<\" ~ dtypes.T ~ \">\"", "ioElement": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "RMSNormalization.LastAxisRow",
+          "shader": "norm-row-stats.wgsl.jinja",
+          "derive": {
+            "modeSpec": "\"rms\"",
+            "vec4": true,
+            "writeStats": false,
+            "rmsScaleAfterCast": true,
+            "scalar": "dtypes.T",
+            "usesF16Spec": "dtypes.T == \"f16\"",
+            "hidden": "dim(shapes.x, -1)",
+            "wg": "min(normMaxWorkgroup, pow2ceil(max(1, dim(shapes.x, -1) / 4)))",
+            "epsilon": "attrs.epsilon",
+            "hiddenVec": "dim(shapes.x, -1) / 4",
+            "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
+            "combineSubgroups": "hasSubgroupId"
+          },
+          "bindings": ["x", "scale", "y", "params"],
+          "dispatch": {
+            "x": "min(normalizedDispatchRows, 65535)",
+            "y": "ceilDiv(normalizedDispatchRows, 65535)",
+            "z": 1
           },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "last_axis_row",
       "priority": 100,
+      "when": ["lastAxisOk", "dtypes.T == dtypes.V", "ranks.scale >= 1", "numel(shapes.scale) == dim(shapes.x, -1)", "dim(shapes.scale, -1) == dim(shapes.x, -1)", "true"],
+      "derive": { "xElement": "dtypes.T", "ioElement": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "RMSNormalization.LastAxisRow",
+          "shader": "norm-row-stats.wgsl.jinja",
+          "derive": {
+            "modeSpec": "\"rms\"",
+            "vec4": false,
+            "writeStats": false,
+            "rmsScaleAfterCast": true,
+            "scalar": "dtypes.T",
+            "usesF16Spec": "dtypes.T == \"f16\"",
+            "hidden": "dim(shapes.x, -1)",
+            "wg": "min(normMaxWorkgroup, pow2ceil(max(1, dim(shapes.x, -1))))",
+            "epsilon": "attrs.epsilon",
+            "hiddenVec": 1,
+            "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
+            "combineSubgroups": "hasSubgroupId"
+          },
+          "bindings": ["x", "scale", "y", "params"],
+          "dispatch": {
+            "x": "min(normalizedDispatchRows, 65535)",
+            "y": "ceilDiv(normalizedDispatchRows, 65535)",
+            "z": 1
           },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,22 +1,32 @@
 {
   "name": "ai.onnx.RMSNormalization",
-  "id": "_ai_onnx_rmsnormalization_webgpu_78878c2",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "aQ8RZqdCxZw7VI+7n2XQneyE8kM1Qzire3++ZtvURxg=",
-      "manifest.json": "Nl952QCnY+Pw6VgYi/lKfYZEh3PpBL2LijyCQqtMTTg=",
-      "norm-row-stats.wgsl.jinja": "uxWNeifDOH5P8jfrDfE95ow3QGhewBtScBqM76ps2YI=",
-      "rms-normalization-splitk-normalize.wgsl.jinja": "TXhEDbPnIGP5Ba5SckUHyuuAdCWCip2BH5sJLZ43DdE=",
-      "rms-normalization-splitk-partials.wgsl.jinja": "GbYu7Cj/LLO9y2Rvhw30RJAUAduI0xfN5YS3H/4Z1k8=",
-      "rms-normalization-stash-f16-serial.wgsl.jinja": "tzxVAVThZ5QRnu0dtHoYwqYErZwth2dViidXHaqsVGo=",
-      "rms-normalization.wgsl.jinja": "2YtdapcEAWQUL/wotiMPs8i6wfRV3VG+y+CPhdu47nM=",
-      "test.json": "09ASLvCprHAheaBj+hX9jc5xDu/05mUj7mOoFwJWWyk="
     }
   },
-  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
-  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.RMSNormalization" }
 }

 {
   "name": "ai.onnx.RMSNormalization",
+  "id": "_ai_onnx_rmsnormalization_webgpu_9f92cf5",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "aQq6UMoOU5W6oOnS6tWV1CjBpS0e4r93laCsON1jvqk=",
+      "manifest.json": "ayMwG0Hl94mrw8rLqrQus65vGITfIySQBSRfgp+6dK0=",
+      "norm-row-stats.wgsl.jinja": "eRBO50QnNhvyqRW/Wdw6rzfJRgVlRT0P6oWVqcDPf5w=",
+      "rms-normalization-splitk-normalize.wgsl.jinja": "Vs708VdMX/fAhYyv9Q7EZbVRNllO6XVhaUaxt+aSJGg=",
+      "rms-normalization-splitk-partials.wgsl.jinja": "Vs4HNsa9ZbLPsW/ZOunRg64qFmbg7WfJ/uE6gqALmSQ=",
+      "rms-normalization-stash-f16-serial.wgsl.jinja": "stNcILDu/EW/WEVtmhYEcIpgE0vOVp0XdvomibMkXBU=",
+      "rms-normalization.wgsl.jinja": "5+i/fHAcHlyZ+Co5hpTNZFHHw1YO+QdqsGg2har/Pa4=",
+      "test.json": "B1fCbyIY2ig58fLsPyM9j4x2gyi1uUuD0TTRN6cUesQ="
     }
   },
+  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.0",
+    "variants": {
+      "stash_f16_serial": ["rms-normalization-stash-f16-serial.wgsl.jinja"],
+      "suffix_axis_splitk": ["rms-normalization-splitk-normalize.wgsl.jinja", "rms-normalization-splitk-partials.wgsl.jinja"],
+      "last_axis": ["rms-normalization.wgsl.jinja"],
+      "suffix_axis": ["rms-normalization.wgsl.jinja"],
+      "last_axis_row_vec4": ["norm-row-stats.wgsl.jinja"],
+      "last_axis_row": ["norm-row-stats.wgsl.jinja"]
+    }
+  }
 }

build/webgpu/norm-row-stats.wgsl.jinja CHANGED Viewed

@@ -1,10 +1,25 @@
-{% if source.usesF16 %}
 enable f16;
 {% endif %}
-{% set combineSubgroups = source.combineSubgroups %}
-{% set scalarIo = source.scalarIo if source.scalarIo is defined else false %}
 {% set rmsScaleVec = "vec4<f32>(scale[i])" %}
 {% set rmsScaleScalar = "f32(scale[i])" %}
 {% set reduceThreadParameters = ", sg_lane: u32, sg_id: u32, num_sg: u32"
     if combineSubgroups else ", tid: u32" %}
 {% set reduceThreadArguments = ", sg_lane, sg_id, num_sg"
@@ -23,14 +38,57 @@ enable subgroups;
 // tree, then every thread applies the fused normalize + affine write.
 //
 // RMS mode uses sum_sq / HIDDEN without computing or subtracting a mean.
-const HIDDEN: u32 = {{ source.hidden }}u;
-{% if source.vec4 %}
-const HIDDEN_V: u32 = {{ source.hiddenVec }}u;
 {% endif %}
-const WG: u32 = {{ source.wg }}u;
-const EPSILON: f32 = {{ source.epsilon }};
 {% if combineSubgroups %}
 var<workgroup> sg_partials: array<f32, WG>;
@@ -84,7 +142,14 @@ fn main(
     return;
   }
   let tid = lid.x;
-{% if source.vec4 and not scalarIo %}
   let base = row * HIDDEN_V;
 {% else %}
   let base = row * HIDDEN;
@@ -92,14 +157,26 @@ fn main(
   var acc = 0.0;
-{% if source.vec4 %}
   for (var i = tid; i < HIDDEN_V; i = i + WG) {
     let v = vec4<f32>(x[base + i]);
     acc = acc + dot(v, v);
   }
 {% else %}
   for (var i = tid; i < HIDDEN; i = i + WG) {
     let v = f32(x[base + i]);
     acc = acc + v * v;
   }
 {% endif %}
@@ -107,18 +184,70 @@ fn main(
   let total = reduce_scalar(acc{{ reduceThreadArguments }});
   let inv = inverseSqrt(total / f32(HIDDEN) + EPSILON);
-{% if source.vec4 %}
   for (var i = tid; i < HIDDEN_V; i = i + WG) {
     let idx = base + i;
     let v = vec4<f32>(x[idx]);
-    y[idx] = {{ source.vecType }}(v * inv) * {{ source.vecType }}({{ rmsScaleVec }});
   }
 {% else %}
   for (var i = tid; i < HIDDEN; i = i + WG) {
     let idx = base + i;
     let v = f32(x[idx]);
-    y[idx] = {{ source.scalar }}(v * inv) * {{ source.scalar }}({{ rmsScaleScalar }});
   }
 {% endif %}
 }

+{% if usesF16Spec %}
 enable f16;
 {% endif %}
+{% set combineSubgroups = combineSubgroups %}
+{% set scalarIo = scalarIo if scalarIo is defined else false %}
+{% set packedBf16Embedding = packedBf16Embedding if packedBf16Embedding is defined else false %}
+{% set writeStats = writeStats if writeStats is defined else false %}
+{% set rmsWeightOffset = rmsWeightOffset if rmsWeightOffset is defined else false %}
+{% set rmsScaleAfterCast = rmsScaleAfterCast if rmsScaleAfterCast is defined else false %}
+{% set rmsResidualAdd = rmsResidualAdd if rmsResidualAdd is defined else false %}
+{% set rmsChainNorm = rmsChainNorm if rmsChainNorm is defined else false %}
+{% set hiddenPairs = hiddenPairs | default(0) %}
+{% set numRows = numRows | default(0) %}
+{% set epsilon = epsilon | default("0.0") %}
+{% set epsilon2 = epsilon2 | default("0.0") %}
+{% if rmsWeightOffset %}
+{% set rmsScaleVec = "(vec4<f32>(1.0) + vec4<f32>(scale[i]))" %}
+{% set rmsScaleScalar = "(1.0 + f32(scale[i]))" %}
+{% else %}
 {% set rmsScaleVec = "vec4<f32>(scale[i])" %}
 {% set rmsScaleScalar = "f32(scale[i])" %}
+{% endif %}
 {% set reduceThreadParameters = ", sg_lane: u32, sg_id: u32, num_sg: u32"
     if combineSubgroups else ", tid: u32" %}
 {% set reduceThreadArguments = ", sg_lane, sg_id, num_sg"
 // tree, then every thread applies the fused normalize + affine write.
 //
 // RMS mode uses sum_sq / HIDDEN without computing or subtracting a mean.
+const HIDDEN: u32 = {{ hidden }}u;
+{% if vec4 %}
+const HIDDEN_V: u32 = {{ hiddenVec }}u;
+{% endif %}
+{% if packedBf16Embedding %}
+const HIDDEN_PAIRS: u32 = {{ hiddenPairs }}u;
+const NUM_ROWS: u32 = {{ numRows }}u;
+{% endif %}
+const WG: u32 = {{ wg }}u;
+const EPSILON: f32 = {{ epsilon }};
+{% if rmsChainNorm %}
+const EPSILON2: f32 = {{ epsilon2 }};
+{% endif %}
+{% if packedBf16Embedding %}
+{% if vec4 %}
+fn unpack_bf16_pair(word: u32) -> vec2<f32> {
+  let bits = vec2<u32>(word & 0xffffu, word >> 16u);
+  return bitcast<vec2<f32>>(bits << vec2<u32>(16u));
+}
+{% endif %}
+{% if not vec4 %}
+fn embedding_scalar(source_row: u32, hidden: u32) -> f32 {
+  if (source_row >= NUM_ROWS) {
+    return 0.0;
+  }
+  let word = x[source_row * HIDDEN_PAIRS + (hidden >> 1u)];
+  let bits = select(word & 0xffffu, word >> 16u, (hidden & 1u) != 0u);
+  return bitcast<f32>(bits << 16u);
+}
 {% endif %}
+{% if vec4 %}
+fn embedding_vec4(source_row: u32, hidden_vec: u32) -> vec4<f32> {
+  if (source_row >= NUM_ROWS) {
+    return vec4<f32>(0.0);
+  }
+  let base = source_row * HIDDEN_PAIRS + hidden_vec * 2u;
+  let low = unpack_bf16_pair(x[base]);
+  let high = unpack_bf16_pair(x[base + 1u]);
+  return vec4<f32>(low, high);
+}
+{% endif %}
+{% endif %}
+{% if vec4 and scalarIo %}
+fn load_vec4(index: u32) -> vec4<f32> {
+  return vec4<f32>(x[index], x[index + 1u], x[index + 2u], x[index + 3u]);
+}
+{% endif %}
 {% if combineSubgroups %}
 var<workgroup> sg_partials: array<f32, WG>;
     return;
   }
   let tid = lid.x;
+{% if packedBf16Embedding %}
+  let source_row = indices[row];
+{% if vec4 %}
+  let base = row * HIDDEN_V;
+{% else %}
+  let base = row * HIDDEN;
+{% endif %}
+{% elif vec4 and not scalarIo %}
   let base = row * HIDDEN_V;
 {% else %}
   let base = row * HIDDEN;
   var acc = 0.0;
+{% if vec4 %}
   for (var i = tid; i < HIDDEN_V; i = i + WG) {
+{% if packedBf16Embedding %}
+    let v = embedding_vec4(source_row, i);
+    embedding_out[base + i] = v;
+{% elif scalarIo %}
+    let v = load_vec4(base + i * 4u);
+{% else %}
     let v = vec4<f32>(x[base + i]);
+{% endif %}
     acc = acc + dot(v, v);
   }
 {% else %}
   for (var i = tid; i < HIDDEN; i = i + WG) {
+{% if packedBf16Embedding %}
+    let v = embedding_scalar(source_row, i);
+    embedding_out[base + i] = v;
+{% else %}
     let v = f32(x[base + i]);
+{% endif %}
     acc = acc + v * v;
   }
 {% endif %}
   let total = reduce_scalar(acc{{ reduceThreadArguments }});
   let inv = inverseSqrt(total / f32(HIDDEN) + EPSILON);
+{% if writeStats %}
+  if (tid == 0u) {
+    inv_std_out[row] = inv;
+  }
+{% endif %}
+{% if rmsChainNorm %}
+  var acc2 = 0.0;
+{% endif %}
+{% if vec4 %}
   for (var i = tid; i < HIDDEN_V; i = i + WG) {
+{% if packedBf16Embedding %}
+    let idx = base + i;
+    let v = embedding_vec4(source_row, i);
+{% elif scalarIo %}
+    let idx = base + i * 4u;
+    let v = load_vec4(idx);
+{% else %}
     let idx = base + i;
     let v = vec4<f32>(x[idx]);
+{% endif %}
+{% if rmsScaleAfterCast %}
+    y[idx] = {{ vecType }}(v * inv) * {{ vecType }}({{ rmsScaleVec }});
+{% elif rmsChainNorm %}
+    // fma(a, b, 0.0) rounds the weighted product exactly as the decomposed pair's store does,
+    // and prevents the compiler from re-contracting it into the residual add.
+    let hv = y[idx] + fma(v * inv, {{ rmsScaleVec }}, vec4<f32>(0.0));
+    y[idx] = hv;
+    acc2 = acc2 + dot(hv, hv);
+{% elif rmsResidualAdd %}
+    // See the chained branch: fma(a, b, 0.0) pins the pre-add rounding of the decomposed pair.
+    y[idx] = y[idx] + fma(v * inv, {{ rmsScaleVec }}, vec4<f32>(0.0));
+{% else %}
+    y[idx] = {{ vecType }}(v * inv * {{ rmsScaleVec }});
+{% endif %}
   }
+{% if rmsChainNorm %}
+  // The chained second norm reads the residual row this loop just stored. This
+  // barrier completes those stores and any preceding shared-scratch use before
+  // the next reduction reuses its scratch; each lane then re-reads only the
+  // elements it wrote itself.
+  workgroupBarrier();
+  let total2 = reduce_scalar(acc2{{ reduceThreadArguments }});
+  let inv2 = inverseSqrt(total2 / f32(HIDDEN) + EPSILON2);
+  for (var i = tid; i < HIDDEN_V; i = i + WG) {
+    let idx = base + i;
+    let hv = vec4<f32>(y[idx]);
+    normed2[idx] = {{ vecType }}(hv * inv2 * vec4<f32>(scale2[i]));
+  }
+{% endif %}
 {% else %}
   for (var i = tid; i < HIDDEN; i = i + WG) {
     let idx = base + i;
+{% if packedBf16Embedding %}
+    let v = embedding_scalar(source_row, i);
+{% else %}
     let v = f32(x[idx]);
+{% endif %}
+{% if rmsScaleAfterCast %}
+    y[idx] = {{ scalar }}(v * inv) * {{ scalar }}({{ rmsScaleScalar }});
+{% else %}
+    y[idx] = {{ scalar }}(v * inv * {{ rmsScaleScalar }});
+{% endif %}
   }
 {% endif %}
 }

build/webgpu/rms-normalization-splitk-normalize.wgsl.jinja CHANGED Viewed

@@ -1,11 +1,7 @@
-// Split-K normalize pass. Each workgroup (row = wg.x, split index = wg.z)
-// folds the SPLIT per-row partial sums of squares into the RMS scale, then
-// normalizes its HIDDEN/SPLIT slice. SPLIT is small (<=64), so the serial fold
-// avoids a third combine pass. Scale offsets follow the suffix-axis broadcast
-// contract.
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 const HIDDEN: u32 = {{ hiddenSize }}u;
@@ -13,11 +9,11 @@ const EPSILON: f32 = {{ epsilon }};
 const WG: u32 = {{ workgroupSize }}u;
 const SPLIT: u32 = {{ split }}u;
-{% if source.scaleRank > 0 %}
-const X_RANK: u32 = {{ source.xRank }}u;
-const SCALE_RANK: u32 = {{ source.scaleRank }}u;
-const X_SHAPE: array<u32, {{ source.xRank }}> = array<u32, {{ source.xRank }}>({% for d in source.xShape %}{{ d }}u{% if not loop.last %}, {% endif %}{% endfor %});
-const SCALE_SHAPE: array<u32, {{ source.scaleRank }}> = array<u32, {{ source.scaleRank }}>({% for d in source.scaleShape %}{{ d }}u{% if not loop.last %}, {% endif %}{% endfor %});
 fn x_stride(axis: u32) -> u32 {
   var stride = 1u;
@@ -36,8 +32,8 @@ fn scale_stride(axis: u32) -> u32 {
 }
 {% endif %}
-fn scale_offset({% if source.scaleRank > 0 %}out_index: u32{% endif %}) -> u32 {
-{% if source.scaleRank == 0 %}
   return 0u;
 {% else %}
   var rem = out_index;
@@ -59,6 +55,8 @@ fn scale_offset({% if source.scaleRank > 0 %}out_index: u32{% endif %}) -> u32 {
 }
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
   let row = wg.x + wg.y * params.rowStride;
@@ -68,12 +66,16 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
   let k = wg.z;
   let tid = lid.x;
-  var total = 0.0;
-  for (var i = 0u; i < SPLIT; i = i + 1u) {
-    total = total + partials[row * SPLIT + i];
   }
-  let inv = inverseSqrt(total / f32(HIDDEN) + EPSILON);
-  let chunk = (HIDDEN + SPLIT - 1u) / SPLIT;
   let start = k * chunk;
   var end = start + chunk;
   if (end > HIDDEN) { end = HIDDEN; }
@@ -86,7 +88,7 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
     // Preserve the ONNX stage boundary: round Normalized to X's dtype before
     // the affine scale is applied.
     let normalized = {{ xElement }}(f32(x[index]) * inv);
-    let value = f32(normalized) * f32(scale[scale_offset({% if source.scaleRank > 0 %}index{% endif %})]);
     y[index] = {{ scalar }}(value);
     d = d + WG;
   }

+// Split-K normalize pass. One lane folds the per-row partials in their original
+// order and shares the inverse RMS with its workgroup. Output tiles can outnumber
+// the reduction splits: their independent work does not need additional scratch.
+// Scale offsets follow the suffix-axis broadcast contract.
 {{ env.wgsl.resourceDeclarations }}
 const HIDDEN: u32 = {{ hiddenSize }}u;
 const WG: u32 = {{ workgroupSize }}u;
 const SPLIT: u32 = {{ split }}u;
+{% if scaleRank > 0 %}
+const X_RANK: u32 = {{ xRank }}u;
+const SCALE_RANK: u32 = {{ scaleRank }}u;
+const X_SHAPE: array<u32, {{ xRank }}> = array<u32, {{ xRank }}>({% for d in xShape %}{{ d }}u{% if not loop.last %}, {% endif %}{% endfor %});
+const SCALE_SHAPE: array<u32, {{ scaleRank }}> = array<u32, {{ scaleRank }}>({% for d in scaleShape %}{{ d }}u{% if not loop.last %}, {% endif %}{% endfor %});
 fn x_stride(axis: u32) -> u32 {
   var stride = 1u;
 }
 {% endif %}
+fn scale_offset({% if scaleRank > 0 %}out_index: u32{% endif %}) -> u32 {
+{% if scaleRank == 0 %}
   return 0u;
 {% else %}
   var rem = out_index;
 }
+var<workgroup> shared_inv: f32;
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
   let row = wg.x + wg.y * params.rowStride;
   let k = wg.z;
   let tid = lid.x;
+  if (tid == 0u) {
+    var total = 0.0;
+    for (var i = 0u; i < SPLIT; i = i + 1u) {
+      total = total + partials[row * SPLIT + i];
+    }
+    shared_inv = inverseSqrt(total / f32(HIDDEN) + EPSILON);
   }
+  workgroupBarrier();
+  let inv = shared_inv;
+  let chunk = {{ normalizeChunk }}u;
   let start = k * chunk;
   var end = start + chunk;
   if (end > HIDDEN) { end = HIDDEN; }
     // Preserve the ONNX stage boundary: round Normalized to X's dtype before
     // the affine scale is applied.
     let normalized = {{ xElement }}(f32(x[index]) * inv);
+    let value = f32(normalized) * f32(scale[scale_offset({% if scaleRank > 0 %}index{% endif %})]);
     y[index] = {{ scalar }}(value);
     d = d + WG;
   }

build/webgpu/rms-normalization-splitk-partials.wgsl.jinja CHANGED Viewed

@@ -51,9 +51,6 @@
    sum-of-squares over its HIDDEN/SPLIT slice and writes one partial to scratch.
    The normalize pass folds the SPLIT partials per row. Split-K reassociates the
    f32 sum, so this route is not bit-identical to the unsplit reduction. */
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 const HIDDEN: u32 = {{ hiddenSize }}u;

    sum-of-squares over its HIDDEN/SPLIT slice and writes one partial to scratch.
    The normalize pass folds the SPLIT partials per row. Split-K reassociates the
    f32 sum, so this route is not bit-identical to the unsplit reduction. */
 {{ env.wgsl.resourceDeclarations }}
 const HIDDEN: u32 = {{ hiddenSize }}u;

build/webgpu/rms-normalization-stash-f16-serial.wgsl.jinja CHANGED Viewed

@@ -59,11 +59,11 @@ fn widen_f16_bits(value: u32) -> f32 {
 }
-{% if source.scaleRank > 0 %}
-const X_RANK: u32 = {{ source.xRank }}u;
-const SCALE_RANK: u32 = {{ source.scaleRank }}u;
-const X_SHAPE: array<u32, {{ source.xRank }}> = array<u32, {{ source.xRank }}>({% for d in source.xShape %}{{ d }}u{% if not loop.last %}, {% endif %}{% endfor %});
-const SCALE_SHAPE: array<u32, {{ source.scaleRank }}> = array<u32, {{ source.scaleRank }}>({% for d in source.scaleShape %}{{ d }}u{% if not loop.last %}, {% endif %}{% endfor %});
 fn x_stride(axis: u32) -> u32 {
   var stride = 1u;
@@ -82,8 +82,8 @@ fn scale_stride(axis: u32) -> u32 {
 }
 {% endif %}
-fn scale_offset({% if source.scaleRank > 0 %}out_index: u32{% endif %}) -> u32 {
-{% if source.scaleRank == 0 %}
   return 0u;
 {% else %}
   var rem = out_index;
@@ -134,7 +134,7 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>) {
     let value_f16 = round_f16_bits_rte(f32(x[index]));
     let normalized = round_f16_bits_rte(widen_f16_bits(value_f16) / widen_f16_bits(rms));
     let value = widen_f16_bits(normalized)
-      * f32(scale[scale_offset({% if source.scaleRank > 0 %}index{% endif %})]);
     y[index] = {{ scalar }}(value);
   }
 }

 }
+{% if scaleRank > 0 %}
+const X_RANK: u32 = {{ xRank }}u;
+const SCALE_RANK: u32 = {{ scaleRank }}u;
+const X_SHAPE: array<u32, {{ xRank }}> = array<u32, {{ xRank }}>({% for d in xShape %}{{ d }}u{% if not loop.last %}, {% endif %}{% endfor %});
+const SCALE_SHAPE: array<u32, {{ scaleRank }}> = array<u32, {{ scaleRank }}>({% for d in scaleShape %}{{ d }}u{% if not loop.last %}, {% endif %}{% endfor %});
 fn x_stride(axis: u32) -> u32 {
   var stride = 1u;
 }
 {% endif %}
+fn scale_offset({% if scaleRank > 0 %}out_index: u32{% endif %}) -> u32 {
+{% if scaleRank == 0 %}
   return 0u;
 {% else %}
   var rem = out_index;
     let value_f16 = round_f16_bits_rte(f32(x[index]));
     let normalized = round_f16_bits_rte(widen_f16_bits(value_f16) / widen_f16_bits(rms));
     let value = widen_f16_bits(normalized)
+      * f32(scale[scale_offset({% if scaleRank > 0 %}index{% endif %})]);
     y[index] = {{ scalar }}(value);
   }
 }

build/webgpu/rms-normalization.wgsl.jinja CHANGED Viewed

@@ -1,6 +1,3 @@
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 const HIDDEN: u32 = {{ hiddenSize }}u;
@@ -9,11 +6,11 @@ const WG: u32 = {{ workgroupSize }}u;
 var<workgroup> partial: array<f32, WG>;
-{% if source.scaleRank > 0 %}
-const X_RANK: u32 = {{ source.xRank }}u;
-const SCALE_RANK: u32 = {{ source.scaleRank }}u;
-const X_SHAPE: array<u32, {{ source.xRank }}> = array<u32, {{ source.xRank }}>({% for d in source.xShape %}{{ d }}u{% if not loop.last %}, {% endif %}{% endfor %});
-const SCALE_SHAPE: array<u32, {{ source.scaleRank }}> = array<u32, {{ source.scaleRank }}>({% for d in source.scaleShape %}{{ d }}u{% if not loop.last %}, {% endif %}{% endfor %});
 fn x_stride(axis: u32) -> u32 {
   var stride = 1u;
@@ -32,8 +29,8 @@ fn scale_stride(axis: u32) -> u32 {
 }
 {% endif %}
-fn scale_offset({% if source.scaleRank > 0 %}out_index: u32{% endif %}) -> u32 {
-{% if source.scaleRank == 0 %}
   return 0u;
 {% else %}
   var rem = out_index;
@@ -138,7 +135,7 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
     // ONNX stage one ends by casting Normalized back to X's dtype; Scale is
     // applied only after that rounding point.
     let normalized = {{ xElement }}(f32(x[index]) * inv);
-    let value = f32(normalized) * f32(scale[scale_offset({% if source.scaleRank > 0 %}index{% endif %})]);
     y[base + d] = {{ scalar }}(value);
   }
 }

 {{ env.wgsl.resourceDeclarations }}
 const HIDDEN: u32 = {{ hiddenSize }}u;
 var<workgroup> partial: array<f32, WG>;
+{% if scaleRank > 0 %}
+const X_RANK: u32 = {{ xRank }}u;
+const SCALE_RANK: u32 = {{ scaleRank }}u;
+const X_SHAPE: array<u32, {{ xRank }}> = array<u32, {{ xRank }}>({% for d in xShape %}{{ d }}u{% if not loop.last %}, {% endif %}{% endfor %});
+const SCALE_SHAPE: array<u32, {{ scaleRank }}> = array<u32, {{ scaleRank }}>({% for d in scaleShape %}{{ d }}u{% if not loop.last %}, {% endif %}{% endfor %});
 fn x_stride(axis: u32) -> u32 {
   var stride = 1u;
 }
 {% endif %}
+fn scale_offset({% if scaleRank > 0 %}out_index: u32{% endif %}) -> u32 {
+{% if scaleRank == 0 %}
   return 0u;
 {% else %}
   var rem = out_index;
     // ONNX stage one ends by casting Normalized back to X's dtype; Scale is
     // applied only after that rounding point.
     let normalized = {{ xElement }}(f32(x[index]) * inv);
+    let value = f32(normalized) * f32(scale[scale_offset({% if scaleRank > 0 %}index{% endif %})]);
     y[base + d] = {{ scalar }}(value);
   }
 }

build/webgpu/test.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "ai.onnx.RMSNormalization",
   "fixtureArrays": {
     "f16_scalar_cast_x": [-1.1103515625, 2.982421875, 1.248046875, -1.8544921875],
     "f16_scalar_cast_scale": [2.015625],
@@ -46,7 +45,7 @@
       "provenance": {
         "source": "onnx/defs/nn/defs.cc",
         "test": "RMSNormalization-23 schema",
-        "notes": "Pins the default-epsilon path: no epsilon attribute is passed, so manifest and oracle defaults must both match the ONNX schema default 1e-5. Small-magnitude rows make epsilon dominate the mean-square, so a wrong default (e.g. 1e-6) diverges by >2x."
       }
     },
     {
@@ -153,7 +152,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/nn/rms_norm_op_test.cc",
         "test": "RMSNormalizationOpTest.RMSNorm_Scale",
-        "notes": "Odd hidden-size companion for valid subnormal scale outputs."
       },
       "attrs": { "epsilon": 0, "axis": -1 },
       "inputs": {
@@ -1372,7 +1371,7 @@
     {
       "name": "f16_lastaxis_unaligned_hidden4094",
       "provenance": {
-        "notes": "Compact correctness sibling for the f16 unaligned-hidden benchmark cliff; dim=4094 exercises the scalar tail path next to the vec4/subgroup path."
       },
       "attrs": { "epsilon": 0.000001, "axis": -1 },
       "inputs": {
@@ -1585,6 +1584,553 @@
         }
       },
       "outputs": { "y": { "dtype": "float32", "shape": [2, 8], "tolerance": 0.000002 } }
     }
   ]
 }

 {
   "fixtureArrays": {
     "f16_scalar_cast_x": [-1.1103515625, 2.982421875, 1.248046875, -1.8544921875],
     "f16_scalar_cast_scale": [2.015625],
       "provenance": {
         "source": "onnx/defs/nn/defs.cc",
         "test": "RMSNormalization-23 schema",
+        "notes": "With no epsilon attribute, the ONNX default 1e-5 applies. Small-magnitude rows make epsilon dominate the mean square, so using a different default changes the normalized values substantially."
       }
     },
     {
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/nn/rms_norm_op_test.cc",
         "test": "RMSNormalizationOpTest.RMSNorm_Scale",
+        "notes": "An odd hidden size exercises scalar-tail normalization with valid subnormal scale outputs."
       },
       "attrs": { "epsilon": 0, "axis": -1 },
       "inputs": {
     {
       "name": "f16_lastaxis_unaligned_hidden4094",
       "provenance": {
+        "notes": "A float16 hidden dimension of 4,094 exercises scalar-tail handling after vectorized normalization."
       },
       "attrs": { "epsilon": 0.000001, "axis": -1 },
       "inputs": {
         }
       },
       "outputs": { "y": { "dtype": "float32", "shape": [2, 8], "tolerance": 0.000002 } }
+    },
+    {
+      "name": "split_f32-1x16384",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 16384],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.023, "scale": 0.5, "offset": 0.0 }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [],
+          "data": { "kind": "fillFloat32", "sinStep": 0.031, "cosStep": 0.009, "scale": 0.25, "offset": 1.0 }
+        }
+      },
+      "outputs": { "y": { "shape": [1, 16384], "dtype": "float32", "tolerance": 0.00001 } }
+    },
+    {
+      "name": "split_f32-1x16385",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 16385],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.023, "scale": 0.5, "offset": 0.0 }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [],
+          "data": { "kind": "fillFloat32", "sinStep": 0.031, "cosStep": 0.009, "scale": 0.25, "offset": 1.0 }
+        }
+      },
+      "outputs": { "y": { "shape": [1, 16385], "dtype": "float32", "tolerance": 0.00001 } }
+    },
+    {
+      "name": "split_f32-1x32769",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 32769],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.023, "scale": 0.5, "offset": 0.0 }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [],
+          "data": { "kind": "fillFloat32", "sinStep": 0.031, "cosStep": 0.009, "scale": 0.25, "offset": 1.0 }
+        }
+      },
+      "outputs": { "y": { "shape": [1, 32769], "dtype": "float32", "tolerance": 0.00001 } }
+    },
+    {
+      "name": "split_f32-1x65536",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 65536],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.023, "scale": 0.5, "offset": 0.0 }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [],
+          "data": { "kind": "fillFloat32", "sinStep": 0.031, "cosStep": 0.009, "scale": 0.25, "offset": 1.0 }
+        }
+      },
+      "outputs": { "y": { "shape": [1, 65536], "dtype": "float32", "tolerance": 0.00001 } }
+    },
+    {
+      "name": "split_f32-1x131072",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 131072],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.023, "scale": 0.5, "offset": 0.0 }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [],
+          "data": { "kind": "fillFloat32", "sinStep": 0.031, "cosStep": 0.009, "scale": 0.25, "offset": 1.0 }
+        }
+      },
+      "outputs": { "y": { "shape": [1, 131072], "dtype": "float32", "tolerance": 0.00001 } }
+    },
+    {
+      "name": "split_f32-1x262144",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 262144],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.023, "scale": 0.5, "offset": 0.0 }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [],
+          "data": { "kind": "fillFloat32", "sinStep": 0.031, "cosStep": 0.009, "scale": 0.25, "offset": 1.0 }
+        }
+      },
+      "outputs": { "y": { "shape": [1, 262144], "dtype": "float32", "tolerance": 0.00001 } }
+    },
+    {
+      "name": "split_f32-1x524288",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 524288],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.023, "scale": 0.5, "offset": 0.0 }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [],
+          "data": { "kind": "fillFloat32", "sinStep": 0.031, "cosStep": 0.009, "scale": 0.25, "offset": 1.0 }
+        }
+      },
+      "outputs": { "y": { "shape": [1, 524288], "dtype": "float32", "tolerance": 0.00001 } }
+    },
+    {
+      "name": "split_f32-1x1048576",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 1048576],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.023, "scale": 0.5, "offset": 0.0 }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [],
+          "data": { "kind": "fillFloat32", "sinStep": 0.031, "cosStep": 0.009, "scale": 0.25, "offset": 1.0 }
+        }
+      },
+      "outputs": { "y": { "shape": [1, 1048576], "dtype": "float32", "tolerance": 0.00001 } }
+    },
+    {
+      "name": "split_f32-1x2097152",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 2097152],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.023, "scale": 0.5, "offset": 0.0 }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [],
+          "data": { "kind": "fillFloat32", "sinStep": 0.031, "cosStep": 0.009, "scale": 0.25, "offset": 1.0 }
+        }
+      },
+      "outputs": { "y": { "shape": [1, 2097152], "dtype": "float32", "tolerance": 0.00001 } }
+    },
+    {
+      "name": "split_f32-2x131073",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 131073],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.023, "scale": 0.5, "offset": 0.0 }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [],
+          "data": { "kind": "fillFloat32", "sinStep": 0.031, "cosStep": 0.009, "scale": 0.25, "offset": 1.0 }
+        }
+      },
+      "outputs": { "y": { "shape": [2, 131073], "dtype": "float32", "tolerance": 0.00001 } }
+    },
+    {
+      "name": "split_f32-3x524289",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 524289],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.023, "scale": 0.5, "offset": 0.0 }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [],
+          "data": { "kind": "fillFloat32", "sinStep": 0.031, "cosStep": 0.009, "scale": 0.25, "offset": 1.0 }
+        }
+      },
+      "outputs": { "y": { "shape": [3, 524289], "dtype": "float32", "tolerance": 0.00001 } }
+    },
+    {
+      "name": "split_f32-4x1048576",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [4, 1048576],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.023, "scale": 0.5, "offset": 0.0 }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [],
+          "data": { "kind": "fillFloat32", "sinStep": 0.031, "cosStep": 0.009, "scale": 0.25, "offset": 1.0 }
+        }
+      },
+      "outputs": { "y": { "shape": [4, 1048576], "dtype": "float32", "tolerance": 0.00001 } }
+    },
+    {
+      "name": "split_f32-8x16384",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [8, 16384],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.023, "scale": 0.5, "offset": 0.0 }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [],
+          "data": { "kind": "fillFloat32", "sinStep": 0.031, "cosStep": 0.009, "scale": 0.25, "offset": 1.0 }
+        }
+      },
+      "outputs": { "y": { "shape": [8, 16384], "dtype": "float32", "tolerance": 0.00001 } }
+    },
+    {
+      "name": "split_f32-16x65536",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [16, 65536],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.023, "scale": 0.5, "offset": 0.0 }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [],
+          "data": { "kind": "fillFloat32", "sinStep": 0.031, "cosStep": 0.009, "scale": 0.25, "offset": 1.0 }
+        }
+      },
+      "outputs": { "y": { "shape": [16, 65536], "dtype": "float32", "tolerance": 0.00001 } }
+    },
+    {
+      "name": "split_f32-32x32769",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [32, 32769],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.023, "scale": 0.5, "offset": 0.0 }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [],
+          "data": { "kind": "fillFloat32", "sinStep": 0.031, "cosStep": 0.009, "scale": 0.25, "offset": 1.0 }
+        }
+      },
+      "outputs": { "y": { "shape": [32, 32769], "dtype": "float32", "tolerance": 0.00001 } }
+    },
+    {
+      "name": "split_f32-128x16384",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [128, 16384],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.023, "scale": 0.5, "offset": 0.0 }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [],
+          "data": { "kind": "fillFloat32", "sinStep": 0.031, "cosStep": 0.009, "scale": 0.25, "offset": 1.0 }
+        }
+      },
+      "outputs": { "y": { "shape": [128, 16384], "dtype": "float32", "tolerance": 0.00001 } }
+    },
+    {
+      "name": "split_f16-1x16385",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [1, 16385],
+          "data": { "kind": "cycle", "values": [-0.75, -0.5, -0.25, 0.0, 0.25, 0.5, 0.75, 1.0] }
+        },
+        "scale": {
+          "dtype": "float16",
+          "shape": [],
+          "data": { "kind": "fillFloat32", "sinStep": 0.031, "cosStep": 0.009, "scale": 0.25, "offset": 1.0 }
+        }
+      },
+      "outputs": { "y": { "shape": [1, 16385], "dtype": "float16", "tolerance": 0.01 } }
+    },
+    {
+      "name": "split_f16-1x131072",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [1, 131072],
+          "data": { "kind": "cycle", "values": [-0.75, -0.5, -0.25, 0.0, 0.25, 0.5, 0.75, 1.0] }
+        },
+        "scale": {
+          "dtype": "float16",
+          "shape": [],
+          "data": { "kind": "fillFloat32", "sinStep": 0.031, "cosStep": 0.009, "scale": 0.25, "offset": 1.0 }
+        }
+      },
+      "outputs": { "y": { "shape": [1, 131072], "dtype": "float16", "tolerance": 0.01 } }
+    },
+    {
+      "name": "split_f16-1x524288",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [1, 524288],
+          "data": { "kind": "cycle", "values": [-0.75, -0.5, -0.25, 0.0, 0.25, 0.5, 0.75, 1.0] }
+        },
+        "scale": {
+          "dtype": "float16",
+          "shape": [],
+          "data": { "kind": "fillFloat32", "sinStep": 0.031, "cosStep": 0.009, "scale": 0.25, "offset": 1.0 }
+        }
+      },
+      "outputs": { "y": { "shape": [1, 524288], "dtype": "float16", "tolerance": 0.01 } }
+    },
+    {
+      "name": "split_f16-1x2097152",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [1, 2097152],
+          "data": { "kind": "cycle", "values": [-0.75, -0.5, -0.25, 0.0, 0.25, 0.5, 0.75, 1.0] }
+        },
+        "scale": {
+          "dtype": "float16",
+          "shape": [],
+          "data": { "kind": "fillFloat32", "sinStep": 0.031, "cosStep": 0.009, "scale": 0.25, "offset": 1.0 }
+        }
+      },
+      "outputs": { "y": { "shape": [1, 2097152], "dtype": "float16", "tolerance": 0.01 } }
+    },
+    {
+      "name": "split_f16-3x524289",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [3, 524289],
+          "data": { "kind": "cycle", "values": [-0.75, -0.5, -0.25, 0.0, 0.25, 0.5, 0.75, 1.0] }
+        },
+        "scale": {
+          "dtype": "float16",
+          "shape": [],
+          "data": { "kind": "fillFloat32", "sinStep": 0.031, "cosStep": 0.009, "scale": 0.25, "offset": 1.0 }
+        }
+      },
+      "outputs": { "y": { "shape": [3, 524289], "dtype": "float16", "tolerance": 0.01 } }
+    },
+    {
+      "name": "split_f16-16x65536",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [16, 65536],
+          "data": { "kind": "cycle", "values": [-0.75, -0.5, -0.25, 0.0, 0.25, 0.5, 0.75, 1.0] }
+        },
+        "scale": {
+          "dtype": "float16",
+          "shape": [],
+          "data": { "kind": "fillFloat32", "sinStep": 0.031, "cosStep": 0.009, "scale": 0.25, "offset": 1.0 }
+        }
+      },
+      "outputs": { "y": { "shape": [16, 65536], "dtype": "float16", "tolerance": 0.01 } }
+    },
+    {
+      "name": "split_f32-f16-2x32769",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 32769],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.023, "scale": 0.5, "offset": 0.0 }
+        },
+        "scale": {
+          "dtype": "float16",
+          "shape": [],
+          "data": { "kind": "fillFloat32", "sinStep": 0.031, "cosStep": 0.009, "scale": 0.25, "offset": 1.0 }
+        }
+      },
+      "outputs": { "y": { "shape": [2, 32769], "dtype": "float16", "tolerance": 0.01 } }
+    },
+    {
+      "name": "split_f16-f32-2x32769",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [2, 32769],
+          "data": { "kind": "cycle", "values": [-0.75, -0.5, -0.25, 0.0, 0.25, 0.5, 0.75, 1.0] }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [],
+          "data": { "kind": "fillFloat32", "sinStep": 0.031, "cosStep": 0.009, "scale": 0.25, "offset": 1.0 }
+        }
+      },
+      "outputs": { "y": { "shape": [2, 32769], "dtype": "float32", "tolerance": 0.00001 } }
+    },
+    {
+      "name": "split_f32-1x16384-split1",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": { "MAX_SPLITS": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 16384],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.023, "scale": 0.5, "offset": 0.0 }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [],
+          "data": { "kind": "fillFloat32", "sinStep": 0.031, "cosStep": 0.009, "scale": 0.25, "offset": 1.0 }
+        }
+      },
+      "outputs": { "y": { "shape": [1, 16384], "dtype": "float32", "tolerance": 0.00001 } }
+    },
+    {
+      "name": "split_f32-1x16385-split3",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": { "MAX_SPLITS": 3 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 16385],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.023, "scale": 0.5, "offset": 0.0 }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [],
+          "data": { "kind": "fillFloat32", "sinStep": 0.031, "cosStep": 0.009, "scale": 0.25, "offset": 1.0 }
+        }
+      },
+      "outputs": { "y": { "shape": [1, 16385], "dtype": "float32", "tolerance": 0.00001 } }
+    },
+    {
+      "name": "split_f32-2x262145-wg64",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": { "WORKGROUP_SIZE": 64 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 262145],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.023, "scale": 0.5, "offset": 0.0 }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [],
+          "data": { "kind": "fillFloat32", "sinStep": 0.031, "cosStep": 0.009, "scale": 0.25, "offset": 1.0 }
+        }
+      },
+      "outputs": { "y": { "shape": [2, 262145], "dtype": "float32", "tolerance": 0.00001 } }
+    },
+    {
+      "name": "split_f32-2x262145-wg128",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": { "WORKGROUP_SIZE": 128 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 262145],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.023, "scale": 0.5, "offset": 0.0 }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [],
+          "data": { "kind": "fillFloat32", "sinStep": 0.031, "cosStep": 0.009, "scale": 0.25, "offset": 1.0 }
+        }
+      },
+      "outputs": { "y": { "shape": [2, 262145], "dtype": "float32", "tolerance": 0.00001 } }
+    },
+    {
+      "name": "split_f32_suffix_broadcast_tail",
+      "attrs": { "axis": 1, "epsilon": 0.000001 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 17, 4097],
+          "data": { "kind": "cycle", "values": [-0.75, -0.5, -0.25, 0.0, 0.25, 0.5, 0.75, 1.0] }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [1, 17, 1],
+          "data": { "kind": "fillFloat32", "sinStep": 0.31, "cosStep": 0.09, "scale": 0.25, "offset": 1.0 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [3, 17, 4097], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "split_f16_suffix_broadcast_tail",
+      "attrs": { "axis": 1, "epsilon": 0.000001 },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [3, 17, 4097],
+          "data": { "kind": "cycle", "values": [-0.75, -0.5, -0.25, 0.0, 0.25, 0.5, 0.75, 1.0] }
+        },
+        "scale": {
+          "dtype": "float16",
+          "shape": [1, 17, 1],
+          "data": { "kind": "fillFloat32", "sinStep": 0.31, "cosStep": 0.09, "scale": 0.25, "offset": 1.0 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float16", "shape": [3, 17, 4097], "tolerance": 0.01 } }
+    },
+    {
+      "name": "split_empty_outer_long_suffix",
+      "attrs": { "axis": 1 },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [0, 2, 16384], "data": { "kind": "values", "values": [] } },
+        "scale": { "dtype": "float32", "shape": [], "data": { "kind": "values", "values": [1.0] } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [0, 2, 16384], "tolerance": 0 } }
     }
   ]
 }