sync 91d990483a17

Browse files

Files changed (9) hide show

README.md +15 -11
build/webgpu/bench.json +468 -1
build/webgpu/manifest.json +214 -386
build/webgpu/metadata.json +24 -10
build/webgpu/norm-row-stats.wgsl.jinja +137 -14
build/webgpu/rms-normalization-splitk-normalize.wgsl.jinja +24 -22
build/webgpu/rms-normalization-splitk-partials.wgsl.jinja +0 -3
build/webgpu/rms-normalization.wgsl.jinja +9 -12
build/webgpu/test.json +713 -2

README.md CHANGED Viewed

@@ -12,21 +12,21 @@ tags:
 ## Description
-Implements ONNX Runtime's legacy RMS normalization for models that serialize `SimplifiedLayerNormalization` in the default ONNX domain: `Y = (X / sqrt(mean(X^2) + epsilon)) * scale`. It shares kernels with [ONNX `RMSNormalization`](https://onnx.ai/onnx/operators/onnx__RMSNormalization.html), but retains scale-before-output-cast ordering and optional `inv_std_var`. See the [compatibility discussion](https://github.com/onnx/onnx/issues/6582#issuecomment-3591862327). Supports float16/float32, float32 statistics, and a nonempty normalization suffix.
 ## Inputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `X` | `x` | `T` | — | — | Input tensor; the mean square is taken over the suffix dimensions starting at `axis`. | required |
-| `scale` | `scale` | `V` | — | — | Scale tensor, unidirectionally broadcastable to `X`; output `Y` has this tensor's dtype. | required |
 ## Outputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `Y` | `y` | `V` | same as `X` | same as `X` | Normalized and scaled output; same shape as `X` and same dtype as `scale`. | required |
-| `inv_std_var` | `invStdVar` | `U` | same as `X` | derived; see description | Optional inverse RMS statistic `1 / sqrt(mean(X^2) + epsilon)`; same rank as `X`, with the dimensions from `axis` onward set to one. | optional |
 ## Attributes
@@ -36,8 +36,8 @@ Default values (overridable per request):
 | --- | --- | --- |
 | `axis` | `-1` | The first dimension of the normalization suffix; negative values count from the end, so the default `-1` normalizes only the last dimension. |
 | `epsilon` | `0.00001` | Small constant added to the mean square before taking the square root to avoid division by zero. |
 | `stash_type` | `1` | TensorProto element type of `inv_std_var`; this package implements the legacy float32 value `1`. |
-| `keep_dims` | `1` | Legacy compatibility attribute accepted and ignored by ONNX Runtime. The supported value `1` records the runtime's fixed behavior: optional statistics retain the input rank with normalized dimensions set to one. |
 ## Type constraints
@@ -49,7 +49,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
+Implements ONNX Runtime's legacy RMS normalization for models that serialize `SimplifiedLayerNormalization` in the default ONNX domain: `Y = (X / sqrt(mean(X^2) + epsilon)) * scale`. Its arithmetic corresponds to [ONNX `RMSNormalization`](https://onnx.ai/onnx/operators/onnx__RMSNormalization.html), while retaining scale-before-output-cast ordering and optional `inv_std_var`. See the [compatibility discussion](https://github.com/onnx/onnx/issues/6582#issuecomment-3591862327). Supports float16/float32, float32 statistics, and a nonempty normalization suffix.
 ## Inputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `x` | `X` | `T` | — | — | Input tensor; the mean square is taken over the suffix dimensions starting at `axis`. | required |
+| `scale` | — | `V` | — | — | Scale tensor, unidirectionally broadcastable to `X`; output `Y` has this tensor's dtype. | required |
 ## Outputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `y` | `Y` | `V` | same as `x` | same as `x` | Normalized and scaled output; same shape as `X` and same dtype as `scale`. | required |
+| `invStdVar` | `inv_std_var` | `U` | same as `x` | derived | Optional inverse RMS statistic `1 / sqrt(mean(X^2) + epsilon)`; same rank as `X`, with the dimensions from `axis` onward set to one. | optional |
 ## Attributes
 | --- | --- | --- |
 | `axis` | `-1` | The first dimension of the normalization suffix; negative values count from the end, so the default `-1` normalizes only the last dimension. |
 | `epsilon` | `0.00001` | Small constant added to the mean square before taking the square root to avoid division by zero. |
+| `keep_dims` | `1` | Legacy compatibility attribute. Only `1` is supported: optional statistics retain the input rank with normalized dimensions set to one. |
 | `stash_type` | `1` | TensorProto element type of `inv_std_var`; this package implements the legacy float32 value `1`. |
 ## Type constraints
 ## 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.SimplifiedLayerNormalization",
   "cases": [
     {
       "name": "simplifiedln-f32-4096x4096",
@@ -61,6 +60,474 @@
           }
         ]
       }
     }
   ]
 }

 {
   "cases": [
     {
       "name": "simplifiedln-f32-4096x4096",
           }
         ]
       }
+    },
+    {
+      "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" },
+        "invStdVar": { "shape": [1, 1], "dtype": "float32" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 131080 }] }
+    },
+    {
+      "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" },
+        "invStdVar": { "shape": [1, 1], "dtype": "float32" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 131088 }] }
+    },
+    {
+      "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" },
+        "invStdVar": { "shape": [1, 1], "dtype": "float32" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 262160 }] }
+    },
+    {
+      "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" },
+        "invStdVar": { "shape": [1, 1], "dtype": "float32" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 524296 }] }
+    },
+    {
+      "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" },
+        "invStdVar": { "shape": [1, 1], "dtype": "float32" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 1048584 }] }
+    },
+    {
+      "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" },
+        "invStdVar": { "shape": [1, 1], "dtype": "float32" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 2097160 }] }
+    },
+    {
+      "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" },
+        "invStdVar": { "shape": [1, 1], "dtype": "float32" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 4194312 }] }
+    },
+    {
+      "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" },
+        "invStdVar": { "shape": [1, 1], "dtype": "float32" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 8388616 }] }
+    },
+    {
+      "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" },
+        "invStdVar": { "shape": [1, 1], "dtype": "float32" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 16777224 }] }
+    },
+    {
+      "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" },
+        "invStdVar": { "shape": [2, 1], "dtype": "float32" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 2097180 }] }
+    },
+    {
+      "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" },
+        "invStdVar": { "shape": [3, 1], "dtype": "float32" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 12582952 }] }
+    },
+    {
+      "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" },
+        "invStdVar": { "shape": [4, 1], "dtype": "float32" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 33554452 }] }
+    },
+    {
+      "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" },
+        "invStdVar": { "shape": [8, 1], "dtype": "float32" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 1048612 }] }
+    },
+    {
+      "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" },
+        "invStdVar": { "shape": [16, 1], "dtype": "float32" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 8388676 }] }
+    },
+    {
+      "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" },
+        "invStdVar": { "shape": [32, 1], "dtype": "float32" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 8388996 }] }
+    },
+    {
+      "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" },
+        "invStdVar": { "shape": [128, 1], "dtype": "float32" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 16777732 }] }
+    },
+    {
+      "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" },
+        "invStdVar": { "shape": [1, 1], "dtype": "float32" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 65546 }] }
+    },
+    {
+      "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" },
+        "invStdVar": { "shape": [1, 1], "dtype": "float32" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 524294 }] }
+    },
+    {
+      "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" },
+        "invStdVar": { "shape": [1, 1], "dtype": "float32" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 2097158 }] }
+    },
+    {
+      "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" },
+        "invStdVar": { "shape": [1, 1], "dtype": "float32" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 8388614 }] }
+    },
+    {
+      "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" },
+        "invStdVar": { "shape": [3, 1], "dtype": "float32" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 6291482 }] }
+    },
+    {
+      "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" },
+        "invStdVar": { "shape": [16, 1], "dtype": "float32" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 4194370 }] }
+    },
+    {
+      "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" },
+        "invStdVar": { "shape": [2, 1], "dtype": "float32" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 393238 }] }
+    },
+    {
+      "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" },
+        "invStdVar": { "shape": [2, 1], "dtype": "float32" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 393240 }] }
+    },
+    {
+      "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" },
+        "invStdVar": { "shape": [1, 1], "dtype": "float32" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 131080 }] }
+    },
+    {
+      "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" },
+        "invStdVar": { "shape": [1, 1], "dtype": "float32" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 131088 }] }
+    },
+    {
+      "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" },
+        "invStdVar": { "shape": [2, 1], "dtype": "float32" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 4194332 }] }
+    },
+    {
+      "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" },
+        "invStdVar": { "shape": [2, 1], "dtype": "float32" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 4194332 }] }
+    },
+    {
+      "name": "split_f32-1x524288-no-stats",
+      "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-1x2097152-no-stats",
+      "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_f16-1x524288-no-stats",
+      "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-no-stats",
+      "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 }] }
     }
   ]
 }

build/webgpu/manifest.json CHANGED Viewed

@@ -3,57 +3,31 @@
   "name": "SimplifiedLayerNormalization",
   "conformance": "legacy-default-domain",
   "sinceVersion": 1,
-  "description": "Implements ONNX Runtime's legacy RMS normalization for models that serialize `SimplifiedLayerNormalization` in the default ONNX domain: `Y = (X / sqrt(mean(X^2) + epsilon)) * scale`. It shares kernels with [ONNX `RMSNormalization`](https://onnx.ai/onnx/operators/onnx__RMSNormalization.html), but retains scale-before-output-cast ordering and optional `inv_std_var`. See the [compatibility discussion](https://github.com/onnx/onnx/issues/6582#issuecomment-3591862327). Supports float16/float32, float32 statistics, and a nonempty normalization suffix.",
-  "inputs": [
-    {
-      "role": "X",
-      "dtype": "T",
-      "description": "Input tensor; the mean square is taken over the suffix dimensions starting at `axis`."
-    },
-    {
-      "role": "scale",
-      "dtype": "V",
-      "description": "Scale tensor, unidirectionally broadcastable to `X`; output `Y` has this tensor's dtype."
-    }
-  ],
-  "outputs": [
-    {
-      "role": "Y",
-      "dtype": "V",
-      "rank": "ranks.X",
-      "shape": "shapes.X",
-      "description": "Normalized and scaled output; same shape as `X` and same dtype as `scale`."
-    },
-    {
-      "role": "inv_std_var",
       "dtype": "U",
-      "rank": "ranks.X",
-      "shape": "prefix(shapes.X, axisNorm) + fill(1, ranks.X - axisNorm)",
       "optional": true,
-      "description": "Optional inverse RMS statistic `1 / sqrt(mean(X^2) + epsilon)`; same rank as `X`, with the dimensions from `axis` onward set to one."
     }
-  ],
-  "attributes": { "axis": -1, "epsilon": 0.00001, "stash_type": 1, "keep_dims": 1 },
-  "attributeDescriptions": {
-    "axis": "The first dimension of the normalization suffix; negative values count from the end, so the default `-1` normalizes 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 of `inv_std_var`; this package implements the legacy float32 value `1`.",
-    "keep_dims": "Legacy compatibility attribute accepted and ignored by ONNX Runtime. The supported value `1` records the runtime's fixed behavior: optional statistics retain the input rank with normalized dimensions set to one."
   },
   "attributeConstraints": { "stash_type": { "values": [1] }, "keep_dims": { "values": [1] } },
   "typeConstraints": { "T": ["float32", "float16"], "V": ["float32", "float16"], "U": ["float32"] },
-  "args": {
-    "x": { "kind": "tensor", "semantic": "X", "role": "input" },
-    "scale": { "kind": "tensor", "semantic": "scale", "role": "input" },
-    "y": { "kind": "tensor", "semantic": "Y", "role": "output" },
-    "invStdVar": { "kind": "tensor", "semantic": "inv_std_var", "role": "output", "required": false }
-  },
   "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)",
@@ -61,187 +35,44 @@
     "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",
-    "normRows": "outer(shapes.X, axisNorm)",
-    "normHidden": "dim(shapes.X, axisNorm) * inner(shapes.X, axisNorm)",
-    "normRowStride": "max(1, min(normRows, device.limits.maxComputeWorkgroupsPerDimension))",
     "rowWg": "min(normMaxWorkgroup, pow2ceil(max(1, normHidden)))",
-    "baseOk": "ranks.X >= 1 and sameShape(shapes.Y, shapes.X) and ranks.scale >= 0 and ranks.scale <= ranks.X and broadcastable(shapes.scale, shapes.X) and attrs.axis + ranks.X >= 0 and attrs.axis < ranks.X and normHidden > 0 and attrs.stash_type == onnxDtypeCode(\"float32\") and f16Ok(dtypes.T) and f16Ok(dtypes.V)",
-    "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)",
     "noStats": "not present.invStdVar",
-    "statsOk": "present.invStdVar and ranks.inv_std_var == ranks.X and sameShape(prefix(shapes.inv_std_var, axisNorm), prefix(shapes.X, axisNorm)) and numel(suffix(shapes.inv_std_var, axisNorm)) == 1",
     "sameDtype": "dtypes.T == dtypes.V",
     "splitCount": "min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(normHidden, tunables.SPLIT_TARGET_ELEMENTS)))",
     "splitScratchBytes": "normRows * splitCount * 4",
-    "splitFits": "normRows <= tunables.SPLIT_MAX_ROWS and splitCount <= device.limits.maxComputeWorkgroupsPerDimension and splitScratchBytes <= device.limits.maxStorageBufferBindingSize and splitScratchBytes <= device.limits.maxBufferSize"
   },
-  "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": "normRows" },
-            { "name": "rowStride", "type": "u32", "value": "normRowStride" }
-          ]
-        }
-      }
-    ],
-    "rowsStats": [
-      {
-        "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": "inv_std_out",
-        "arg": "invStdVar",
-        "semantic": "inv_std_var",
-        "buffer": { "type": "storage" },
-        "elementType": "f32"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "normRows" },
-            { "name": "rowStride", "type": "u32", "value": "normRowStride" }
-          ]
-        }
-      }
-    ],
-    "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": "normRows" },
-            { "name": "rowStride", "type": "u32", "value": "normRowStride" }
-          ]
-        }
-      }
-    ],
-    "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": "normRows" },
-            { "name": "rowStride", "type": "u32", "value": "normRowStride" }
-          ]
-        }
-      }
-    ],
-    "splitNormalizeStats": [
-      {
-        "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": "inv_std_out",
-        "arg": "invStdVar",
-        "semantic": "inv_std_var",
-        "buffer": { "type": "storage" },
-        "elementType": "f32"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "normRows" },
-            { "name": "rowStride", "type": "u32", "value": "normRowStride" }
-          ]
-        }
-      }
-    ]
   },
   "variants": [
     {
       "id": "last_axis",
       "priority": 1,
       "when": ["lastAxisOk", "noStats"],
-      "constants": {
         "scalar": "dtypes.V",
         "xElement": "dtypes.T",
         "ioElement": "dtypes.V",
-        "usesF16": "dtypes.T == \"f16\" or dtypes.V == \"f16\"",
         "hiddenSize": "normHidden",
         "workgroupSize": "rowWg",
         "epsilon": "attrs.epsilon"
@@ -250,19 +81,17 @@
         {
           "id": "main",
           "name": "SimplifiedLayerNormalization.Row",
-          "source": {
-            "shader": "rms-normalization.wgsl.jinja",
-            "inputs": {
-              "xShape": "shapes.X",
-              "scaleShape": "shapes.scale",
-              "xRank": "ranks.X",
-              "scaleRank": "ranks.scale",
-              "writeStats": false,
-              "rmsScaleAfterCast": false
-            }
           },
-          "bindings": "rows",
-          "dispatch": { "workgroups": "normRows" }
         }
       ]
     },
@@ -270,11 +99,10 @@
       "id": "last_axis_stats",
       "priority": 2,
       "when": ["lastAxisOk", "statsOk"],
-      "constants": {
         "scalar": "dtypes.V",
         "xElement": "dtypes.T",
         "ioElement": "dtypes.V",
-        "usesF16": "dtypes.T == \"f16\" or dtypes.V == \"f16\"",
         "hiddenSize": "normHidden",
         "workgroupSize": "rowWg",
         "epsilon": "attrs.epsilon"
@@ -283,19 +111,17 @@
         {
           "id": "main",
           "name": "SimplifiedLayerNormalization.Row",
-          "source": {
-            "shader": "rms-normalization.wgsl.jinja",
-            "inputs": {
-              "xShape": "shapes.X",
-              "scaleShape": "shapes.scale",
-              "xRank": "ranks.X",
-              "scaleRank": "ranks.scale",
-              "writeStats": true,
-              "rmsScaleAfterCast": false
-            }
           },
-          "bindings": "rowsStats",
-          "dispatch": { "workgroups": "normRows" }
         }
       ]
     },
@@ -303,11 +129,10 @@
       "id": "suffix_axis",
       "priority": 10,
       "when": ["suffixAxisOk", "noStats"],
-      "constants": {
         "scalar": "dtypes.V",
         "xElement": "dtypes.T",
         "ioElement": "dtypes.V",
-        "usesF16": "dtypes.T == \"f16\" or dtypes.V == \"f16\"",
         "hiddenSize": "normHidden",
         "workgroupSize": "rowWg",
         "epsilon": "attrs.epsilon"
@@ -316,19 +141,17 @@
         {
           "id": "main",
           "name": "SimplifiedLayerNormalization.Row",
-          "source": {
-            "shader": "rms-normalization.wgsl.jinja",
-            "inputs": {
-              "xShape": "shapes.X",
-              "scaleShape": "shapes.scale",
-              "xRank": "ranks.X",
-              "scaleRank": "ranks.scale",
-              "writeStats": false,
-              "rmsScaleAfterCast": false
-            }
           },
-          "bindings": "rows",
-          "dispatch": { "workgroups": "normRows" }
         }
       ]
     },
@@ -336,11 +159,10 @@
       "id": "suffix_axis_stats",
       "priority": 11,
       "when": ["suffixAxisOk", "statsOk"],
-      "constants": {
         "scalar": "dtypes.V",
         "xElement": "dtypes.T",
         "ioElement": "dtypes.V",
-        "usesF16": "dtypes.T == \"f16\" or dtypes.V == \"f16\"",
         "hiddenSize": "normHidden",
         "workgroupSize": "rowWg",
         "epsilon": "attrs.epsilon"
@@ -349,113 +171,127 @@
         {
           "id": "main",
           "name": "SimplifiedLayerNormalization.Row",
-          "source": {
-            "shader": "rms-normalization.wgsl.jinja",
-            "inputs": {
-              "xShape": "shapes.X",
-              "scaleShape": "shapes.scale",
-              "xRank": "ranks.X",
-              "scaleRank": "ranks.scale",
-              "writeStats": true,
-              "rmsScaleAfterCast": false
-            }
           },
-          "bindings": "rowsStats",
-          "dispatch": { "workgroups": "normRows" }
         }
       ]
     },
     {
       "id": "suffix_axis_splitk",
       "priority": 15,
-      "when": ["baseOk", "ranks.X >= 2", "noStats", "splitFits"],
       "demoteWhen": ["reportedNonWave32Adapter", "normHidden < tunables.SPLIT_MIN_HIDDEN"],
-      "constants": {
         "scalar": "dtypes.V",
         "xElement": "dtypes.T",
         "ioElement": "dtypes.V",
-        "usesF16": "dtypes.T == \"f16\" or dtypes.V == \"f16\"",
         "hiddenSize": "normHidden",
         "workgroupSize": "normMaxWorkgroup",
         "split": "splitCount",
-        "epsilon": "attrs.epsilon"
       },
       "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[normRows * splitCount]" }],
       "passes": [
         {
           "id": "partials",
           "name": "SimplifiedLayerNormalization.SplitKPartials",
-          "source": { "shader": "rms-normalization-splitk-partials.wgsl.jinja", "inputs": {} },
-          "bindings": "splitPartials",
-          "dispatch": { "workgroups": "normRows", "z": "splitCount" }
         },
         {
           "id": "normalize",
           "name": "SimplifiedLayerNormalization.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": false
-            }
           },
-          "bindings": "splitNormalize",
-          "dispatch": { "workgroups": "normRows", "z": "splitCount" }
         }
       ]
     },
     {
       "id": "suffix_axis_splitk_stats",
       "priority": 16,
-      "when": ["baseOk", "ranks.X >= 2", "statsOk", "splitFits"],
       "demoteWhen": ["reportedNonWave32Adapter", "normHidden < tunables.SPLIT_MIN_HIDDEN"],
-      "constants": {
         "scalar": "dtypes.V",
         "xElement": "dtypes.T",
         "ioElement": "dtypes.V",
-        "usesF16": "dtypes.T == \"f16\" or dtypes.V == \"f16\"",
         "hiddenSize": "normHidden",
         "workgroupSize": "normMaxWorkgroup",
         "split": "splitCount",
-        "epsilon": "attrs.epsilon"
       },
       "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[normRows * splitCount]" }],
       "passes": [
         {
           "id": "partials",
           "name": "SimplifiedLayerNormalization.SplitKPartials",
-          "source": { "shader": "rms-normalization-splitk-partials.wgsl.jinja", "inputs": {} },
-          "bindings": "splitPartials",
-          "dispatch": { "workgroups": "normRows", "z": "splitCount" }
         },
         {
           "id": "normalize",
           "name": "SimplifiedLayerNormalization.SplitKNormalize",
-          "source": {
-            "shader": "rms-normalization-splitk-normalize.wgsl.jinja",
-            "inputs": {
-              "xShape": "shapes.X",
-              "scaleShape": "shapes.scale",
-              "xRank": "ranks.X",
-              "scaleRank": "ranks.scale",
-              "writeStats": true,
-              "rmsScaleAfterCast": false
-            }
           },
-          "bindings": "splitNormalizeStats",
-          "dispatch": { "workgroups": "normRows", "z": "splitCount" }
         }
       ]
     },
     {
       "id": "last_axis_row_vec4",
       "priority": 110,
-      "when": ["lastAxisOk", "sameDtype", "noStats", "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": {
         "scalar": "dtypes.T",
         "xElement": "\"vec4<\" ~ dtypes.T ~ \">\"",
         "ioElement": "\"vec4<\" ~ dtypes.T ~ \">\""
@@ -464,66 +300,62 @@
         {
           "id": "main",
           "name": "SimplifiedLayerNormalization.LastAxisRow",
-          "source": {
-            "shader": "norm-row-stats.wgsl.jinja",
-            "inputs": {
-              "mode": "\"rms\"",
-              "vec4": true,
-              "writeStats": false,
-              "rmsScaleAfterCast": false,
-              "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": "normRows" }
         }
       ]
     },
     {
       "id": "last_axis_row",
       "priority": 100,
-      "when": ["lastAxisOk", "sameDtype", "noStats", "ranks.scale >= 1", "numel(shapes.scale) == dim(shapes.X, -1)", "dim(shapes.scale, -1) == dim(shapes.X, -1)"],
-      "constants": { "scalar": "dtypes.T", "xElement": "dtypes.T", "ioElement": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "SimplifiedLayerNormalization.LastAxisRow",
-          "source": {
-            "shader": "norm-row-stats.wgsl.jinja",
-            "inputs": {
-              "mode": "\"rms\"",
-              "vec4": false,
-              "writeStats": false,
-              "rmsScaleAfterCast": false,
-              "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": "normRows" }
         }
       ]
     },
     {
       "id": "last_axis_row_vec4_stats",
       "priority": 112,
-      "when": ["lastAxisOk", "sameDtype", "statsOk", "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": {
         "scalar": "dtypes.T",
         "xElement": "\"vec4<\" ~ dtypes.T ~ \">\"",
         "ioElement": "\"vec4<\" ~ dtypes.T ~ \">\""
@@ -532,58 +364,54 @@
         {
           "id": "main",
           "name": "SimplifiedLayerNormalization.LastAxisRow",
-          "source": {
-            "shader": "norm-row-stats.wgsl.jinja",
-            "inputs": {
-              "mode": "\"rms\"",
-              "vec4": true,
-              "writeStats": true,
-              "rmsScaleAfterCast": false,
-              "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": "rowsStats",
-          "dispatch": { "workgroups": "normRows" }
         }
       ]
     },
     {
       "id": "last_axis_row_stats",
       "priority": 102,
-      "when": ["lastAxisOk", "sameDtype", "statsOk", "ranks.scale >= 1", "numel(shapes.scale) == dim(shapes.X, -1)", "dim(shapes.scale, -1) == dim(shapes.X, -1)"],
-      "constants": { "scalar": "dtypes.T", "xElement": "dtypes.T", "ioElement": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "SimplifiedLayerNormalization.LastAxisRow",
-          "source": {
-            "shader": "norm-row-stats.wgsl.jinja",
-            "inputs": {
-              "mode": "\"rms\"",
-              "vec4": false,
-              "writeStats": true,
-              "rmsScaleAfterCast": false,
-              "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": "rowsStats",
-          "dispatch": { "workgroups": "normRows" }
         }
       ]
     }

   "name": "SimplifiedLayerNormalization",
   "conformance": "legacy-default-domain",
   "sinceVersion": 1,
+  "inputs": { "x": { "onnx": "X", "dtype": "T" }, "scale": { "dtype": "V" } },
+  "outputs": {
+    "y": { "onnx": "Y", "dtype": "V", "rank": "ranks.x", "shape": "shapes.x" },
+    "invStdVar": {
+      "onnx": "inv_std_var",
       "dtype": "U",
+      "rank": "ranks.x",
       "optional": true,
+      "shape": "prefix(shapes.x, axisNorm) + fill(1, ranks.x - axisNorm)"
     }
+  },
+  "attributes": {
+    "axis": { "default": -1 },
+    "epsilon": { "default": 0.00001 },
+    "stash_type": { "default": 1 },
+    "keep_dims": { "default": 1 }
   },
   "attributeConstraints": { "stash_type": { "values": [1] }, "keep_dims": { "values": [1] } },
   "typeConstraints": { "T": ["float32", "float16"], "V": ["float32", "float16"], "U": ["float32"] },
   "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",
+    "normRows": "outer(shapes.x, axisNorm)",
+    "normHidden": "dim(shapes.x, axisNorm) * inner(shapes.x, axisNorm)",
+    "normRowStride": "max(1, min(normRows, min(device.limits.maxComputeWorkgroupsPerDimension, 65535)))",
     "rowWg": "min(normMaxWorkgroup, pow2ceil(max(1, normHidden)))",
+    "baseOk": "ranks.x >= 1 and sameShape(shapes.y, shapes.x) and ranks.scale >= 0 and ranks.scale <= ranks.x and broadcastable(shapes.scale, shapes.x) and attrs.axis + ranks.x >= 0 and attrs.axis < ranks.x and normHidden > 0 and attrs.stash_type == onnxDtypeCode(\"float32\") and f16Ok(dtypes.T) and f16Ok(dtypes.V)",
+    "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)",
     "noStats": "not present.invStdVar",
+    "statsOk": "present.invStdVar and ranks.invStdVar == ranks.x and sameShape(prefix(shapes.invStdVar, axisNorm), prefix(shapes.x, axisNorm)) and numel(suffix(shapes.invStdVar, axisNorm)) == 1",
     "sameDtype": "dtypes.T == dtypes.V",
     "splitCount": "min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(normHidden, tunables.SPLIT_TARGET_ELEMENTS)))",
     "splitScratchBytes": "normRows * splitCount * 4",
+    "splitFits": "normRows <= tunables.SPLIT_MAX_ROWS and splitCount <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and splitScratchBytes <= device.limits.maxStorageBufferBindingSize and splitScratchBytes <= device.limits.maxBufferSize"
   },
+  "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": "normRows" },
+        { "name": "rowStride", "type": "u32", "value": "normRowStride" }
+      ]
+    },
+    "inv_std_out": { "arg": "invStdVar", "buffer": "storage", "elementType": "f32" },
+    "partials_2": { "name": "partials", "buffer": "read-only-storage", "elementType": "f32" }
   },
   "variants": [
     {
       "id": "last_axis",
       "priority": 1,
       "when": ["lastAxisOk", "noStats"],
+      "derive": {
         "scalar": "dtypes.V",
         "xElement": "dtypes.T",
         "ioElement": "dtypes.V",
         "hiddenSize": "normHidden",
         "workgroupSize": "rowWg",
         "epsilon": "attrs.epsilon"
         {
           "id": "main",
           "name": "SimplifiedLayerNormalization.Row",
+          "shader": "rms-normalization.wgsl.jinja",
+          "derive": {
+            "xShape": "shapes.x",
+            "scaleShape": "shapes.scale",
+            "xRank": "ranks.x",
+            "scaleRank": "ranks.scale",
+            "writeStats": false,
+            "rmsScaleAfterCast": false
           },
+          "bindings": ["x", "scale", "y", "params"],
+          "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
         }
       ]
     },
       "id": "last_axis_stats",
       "priority": 2,
       "when": ["lastAxisOk", "statsOk"],
+      "derive": {
         "scalar": "dtypes.V",
         "xElement": "dtypes.T",
         "ioElement": "dtypes.V",
         "hiddenSize": "normHidden",
         "workgroupSize": "rowWg",
         "epsilon": "attrs.epsilon"
         {
           "id": "main",
           "name": "SimplifiedLayerNormalization.Row",
+          "shader": "rms-normalization.wgsl.jinja",
+          "derive": {
+            "xShape": "shapes.x",
+            "scaleShape": "shapes.scale",
+            "xRank": "ranks.x",
+            "scaleRank": "ranks.scale",
+            "writeStats": true,
+            "rmsScaleAfterCast": false
           },
+          "bindings": ["x", "scale", "y", "inv_std_out", "params"],
+          "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
         }
       ]
     },
       "id": "suffix_axis",
       "priority": 10,
       "when": ["suffixAxisOk", "noStats"],
+      "derive": {
         "scalar": "dtypes.V",
         "xElement": "dtypes.T",
         "ioElement": "dtypes.V",
         "hiddenSize": "normHidden",
         "workgroupSize": "rowWg",
         "epsilon": "attrs.epsilon"
         {
           "id": "main",
           "name": "SimplifiedLayerNormalization.Row",
+          "shader": "rms-normalization.wgsl.jinja",
+          "derive": {
+            "xShape": "shapes.x",
+            "scaleShape": "shapes.scale",
+            "xRank": "ranks.x",
+            "scaleRank": "ranks.scale",
+            "writeStats": false,
+            "rmsScaleAfterCast": false
           },
+          "bindings": ["x", "scale", "y", "params"],
+          "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
         }
       ]
     },
       "id": "suffix_axis_stats",
       "priority": 11,
       "when": ["suffixAxisOk", "statsOk"],
+      "derive": {
         "scalar": "dtypes.V",
         "xElement": "dtypes.T",
         "ioElement": "dtypes.V",
         "hiddenSize": "normHidden",
         "workgroupSize": "rowWg",
         "epsilon": "attrs.epsilon"
         {
           "id": "main",
           "name": "SimplifiedLayerNormalization.Row",
+          "shader": "rms-normalization.wgsl.jinja",
+          "derive": {
+            "xShape": "shapes.x",
+            "scaleShape": "shapes.scale",
+            "xRank": "ranks.x",
+            "scaleRank": "ranks.scale",
+            "writeStats": true,
+            "rmsScaleAfterCast": false
           },
+          "bindings": ["x", "scale", "y", "inv_std_out", "params"],
+          "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
         }
       ]
     },
     {
       "id": "suffix_axis_splitk",
       "priority": 15,
+      "when": ["baseOk", "ranks.x >= 2", "noStats", "splitFits"],
       "demoteWhen": ["reportedNonWave32Adapter", "normHidden < tunables.SPLIT_MIN_HIDDEN"],
+      "derive": {
         "scalar": "dtypes.V",
         "xElement": "dtypes.T",
         "ioElement": "dtypes.V",
         "hiddenSize": "normHidden",
         "workgroupSize": "normMaxWorkgroup",
         "split": "splitCount",
+        "epsilon": "attrs.epsilon",
+        "normalizeRows": "normRows"
       },
       "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[normRows * splitCount]" }],
       "passes": [
         {
           "id": "partials",
           "name": "SimplifiedLayerNormalization.SplitKPartials",
+          "shader": "rms-normalization-splitk-partials.wgsl.jinja",
+          "bindings": ["x", { "name": "partials", "buffer": "storage", "elementType": "f32" }, "params"],
+          "dispatch": {
+            "x": "min(normRows, DISPATCH_FOLD_WIDTH)",
+            "y": "ceilDiv(normRows, DISPATCH_FOLD_WIDTH)",
+            "z": "splitCount"
+          }
         },
         {
           "id": "normalize",
           "name": "SimplifiedLayerNormalization.SplitKNormalize",
+          "shader": "rms-normalization-splitk-normalize.wgsl.jinja",
+          "derive": {
+            "xShape": "shapes.x",
+            "scaleShape": "shapes.scale",
+            "xRank": "ranks.x",
+            "scaleRank": "ranks.scale",
+            "writeStats": false,
+            "rmsScaleAfterCast": false,
+            "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", "partials_2", "y", "params"],
+          "dispatch": {
+            "x": "min(normRows, DISPATCH_FOLD_WIDTH)",
+            "y": "ceilDiv(normRows, DISPATCH_FOLD_WIDTH)",
+            "z": "normalizeBlocks"
+          }
         }
       ]
     },
     {
       "id": "suffix_axis_splitk_stats",
       "priority": 16,
+      "when": ["baseOk", "ranks.x >= 2", "statsOk", "splitFits"],
       "demoteWhen": ["reportedNonWave32Adapter", "normHidden < tunables.SPLIT_MIN_HIDDEN"],
+      "derive": {
         "scalar": "dtypes.V",
         "xElement": "dtypes.T",
         "ioElement": "dtypes.V",
         "hiddenSize": "normHidden",
         "workgroupSize": "normMaxWorkgroup",
         "split": "splitCount",
+        "epsilon": "attrs.epsilon",
+        "normalizeRows": "normRows"
       },
       "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[normRows * splitCount]" }],
       "passes": [
         {
           "id": "partials",
           "name": "SimplifiedLayerNormalization.SplitKPartials",
+          "shader": "rms-normalization-splitk-partials.wgsl.jinja",
+          "bindings": ["x", { "name": "partials", "buffer": "storage", "elementType": "f32" }, "params"],
+          "dispatch": {
+            "x": "min(normRows, DISPATCH_FOLD_WIDTH)",
+            "y": "ceilDiv(normRows, DISPATCH_FOLD_WIDTH)",
+            "z": "splitCount"
+          }
         },
         {
           "id": "normalize",
           "name": "SimplifiedLayerNormalization.SplitKNormalize",
+          "shader": "rms-normalization-splitk-normalize.wgsl.jinja",
+          "derive": {
+            "xShape": "shapes.x",
+            "scaleShape": "shapes.scale",
+            "xRank": "ranks.x",
+            "scaleRank": "ranks.scale",
+            "writeStats": true,
+            "rmsScaleAfterCast": false,
+            "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", "partials_2", "y", "inv_std_out", "params"],
+          "dispatch": {
+            "x": "min(normRows, DISPATCH_FOLD_WIDTH)",
+            "y": "ceilDiv(normRows, DISPATCH_FOLD_WIDTH)",
+            "z": "normalizeBlocks"
+          }
         }
       ]
     },
     {
       "id": "last_axis_row_vec4",
       "priority": 110,
+      "when": ["lastAxisOk", "sameDtype", "noStats", "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": {
         "scalar": "dtypes.T",
         "xElement": "\"vec4<\" ~ dtypes.T ~ \">\"",
         "ioElement": "\"vec4<\" ~ dtypes.T ~ \">\""
         {
           "id": "main",
           "name": "SimplifiedLayerNormalization.LastAxisRow",
+          "shader": "norm-row-stats.wgsl.jinja",
+          "derive": {
+            "modeSpec": "\"rms\"",
+            "vec4": true,
+            "writeStats": false,
+            "rmsScaleAfterCast": false,
+            "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(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "last_axis_row",
       "priority": 100,
+      "when": ["lastAxisOk", "sameDtype", "noStats", "ranks.scale >= 1", "numel(shapes.scale) == dim(shapes.x, -1)", "dim(shapes.scale, -1) == dim(shapes.x, -1)"],
+      "derive": { "scalar": "dtypes.T", "xElement": "dtypes.T", "ioElement": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "SimplifiedLayerNormalization.LastAxisRow",
+          "shader": "norm-row-stats.wgsl.jinja",
+          "derive": {
+            "modeSpec": "\"rms\"",
+            "vec4": false,
+            "writeStats": false,
+            "rmsScaleAfterCast": false,
+            "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(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "last_axis_row_vec4_stats",
       "priority": 112,
+      "when": ["lastAxisOk", "sameDtype", "statsOk", "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": {
         "scalar": "dtypes.T",
         "xElement": "\"vec4<\" ~ dtypes.T ~ \">\"",
         "ioElement": "\"vec4<\" ~ dtypes.T ~ \">\""
         {
           "id": "main",
           "name": "SimplifiedLayerNormalization.LastAxisRow",
+          "shader": "norm-row-stats.wgsl.jinja",
+          "derive": {
+            "modeSpec": "\"rms\"",
+            "vec4": true,
+            "writeStats": true,
+            "rmsScaleAfterCast": false,
+            "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", "inv_std_out", "params"],
+          "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "last_axis_row_stats",
       "priority": 102,
+      "when": ["lastAxisOk", "sameDtype", "statsOk", "ranks.scale >= 1", "numel(shapes.scale) == dim(shapes.x, -1)", "dim(shapes.scale, -1) == dim(shapes.x, -1)"],
+      "derive": { "scalar": "dtypes.T", "xElement": "dtypes.T", "ioElement": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "SimplifiedLayerNormalization.LastAxisRow",
+          "shader": "norm-row-stats.wgsl.jinja",
+          "derive": {
+            "modeSpec": "\"rms\"",
+            "vec4": false,
+            "writeStats": true,
+            "rmsScaleAfterCast": false,
+            "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", "inv_std_out", "params"],
+          "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,21 +1,35 @@
 {
   "name": "ai.onnx.SimplifiedLayerNormalization",
-  "id": "_ai_onnx_simplifiedlayernormalization_webgpu_e7aa570",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "pvZ/eyxjfpTLTo48nnQhDaKG61UPA6L0EPXu4QnH7MY=",
-      "manifest.json": "hfKp3YN8J74f7KzVbFXsOBq5wNbZm2A8QMJ3ww5MAdw=",
-      "norm-row-stats.wgsl.jinja": "GopNOsand0HoXzNtQnzGdvJKmHbCv96No/7heLcS9VA=",
-      "rms-normalization-splitk-normalize.wgsl.jinja": "YklXSmW/XGI1M5KVzfayKp0EQ+3iAYHeOnTAMA4UW4c=",
-      "rms-normalization-splitk-partials.wgsl.jinja": "GbYu7Cj/LLO9y2Rvhw30RJAUAduI0xfN5YS3H/4Z1k8=",
-      "rms-normalization.wgsl.jinja": "wNOsrtTIc7DmBifozpAm2fZ3iWznPJ04EhFMkOq6GqA=",
-      "test.json": "ycI7oQ1MKS87yo6QXLI36xLz840eYS69Y/anRQHQkUY="
     }
   },
-  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
-  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.SimplifiedLayerNormalization" }
 }

 {
   "name": "ai.onnx.SimplifiedLayerNormalization",
+  "id": "_ai_onnx_simplifiedlayernormalization_webgpu_b0bbd51",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "oL4NDZsKfggbFJ8moDoDY5cYbkZb/dnRCxB7mA3YuW0=",
+      "manifest.json": "M3zOOpjNYbDBPlApZytNSUFrdixB+kGxE4X9xKY7yzk=",
+      "norm-row-stats.wgsl.jinja": "eRBO50QnNhvyqRW/Wdw6rzfJRgVlRT0P6oWVqcDPf5w=",
+      "rms-normalization-splitk-normalize.wgsl.jinja": "vqx3jngJ7c+zkmcD3Nc7mt4Fm77XXDmtrQt9ULWgHKE=",
+      "rms-normalization-splitk-partials.wgsl.jinja": "Vs4HNsa9ZbLPsW/ZOunRg64qFmbg7WfJ/uE6gqALmSQ=",
+      "rms-normalization.wgsl.jinja": "tRDNHNbkHoidOuTQwnH3Jx0msfP74Px5ukRNF9Hh8fw=",
+      "test.json": "PYhuNCaFTpBU8GJFS/LzpdoBKLbT0praPXzHhs7rg50="
     }
   },
+  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.0",
+    "variants": {
+      "last_axis": ["rms-normalization.wgsl.jinja"],
+      "last_axis_stats": ["rms-normalization.wgsl.jinja"],
+      "suffix_axis": ["rms-normalization.wgsl.jinja"],
+      "suffix_axis_stats": ["rms-normalization.wgsl.jinja"],
+      "suffix_axis_splitk": ["rms-normalization-splitk-normalize.wgsl.jinja", "rms-normalization-splitk-partials.wgsl.jinja"],
+      "suffix_axis_splitk_stats": ["rms-normalization-splitk-normalize.wgsl.jinja", "rms-normalization-splitk-partials.wgsl.jinja"],
+      "last_axis_row_vec4": ["norm-row-stats.wgsl.jinja"],
+      "last_axis_row": ["norm-row-stats.wgsl.jinja"],
+      "last_axis_row_vec4_stats": ["norm-row-stats.wgsl.jinja"],
+      "last_axis_row_stats": ["norm-row-stats.wgsl.jinja"]
+    }
+  }
 }

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

@@ -1,11 +1,25 @@
-{% if source.usesF16 %}
 enable f16;
 {% endif %}
-{% set combineSubgroups = source.combineSubgroups %}
-{% set scalarIo = source.scalarIo if source.scalarIo is defined else false %}
-{% set writeStats = source.writeStats if source.writeStats 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"
@@ -24,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>;
@@ -85,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;
@@ -93,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 %}
@@ -114,17 +190,64 @@ fn main(
   }
 {% endif %}
-{% 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 * {{ 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 * {{ 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 %}
   }
 {% 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);
-{% if source.writeStats %}
   // Every split workgroup folds the same partials, so one designated
   // workgroup writes the row statistic.
   if (k == 0u && tid == 0u) {
@@ -81,7 +83,7 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
   }
 {% endif %}
-  let chunk = (HIDDEN + SPLIT - 1u) / SPLIT;
   let start = k * chunk;
   var end = start + chunk;
   if (end > HIDDEN) { end = HIDDEN; }
@@ -91,7 +93,7 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
   loop {
     if (d >= end) { break; }
     let index = base + d;
-    let value = f32(x[index]) * inv * 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;
+{% if writeStats %}
   // Every split workgroup folds the same partials, so one designated
   // workgroup writes the row statistic.
   if (k == 0u && tid == 0u) {
   }
 {% endif %}
+  let chunk = {{ normalizeChunk }}u;
   let start = k * chunk;
   var end = start + chunk;
   if (end > HIDDEN) { end = HIDDEN; }
   loop {
     if (d >= end) { break; }
     let index = base + d;
+    let value = f32(x[index]) * inv * 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.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;
@@ -133,7 +130,7 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
     local_sq = local_sq + value * value;
   }
   let inv = inverseSqrt(reduce_sum(local_sq, tid) / f32(HIDDEN) + EPSILON);
-{% if source.writeStats %}
   if (tid == 0u) {
     inv_std_out[row] = inv;
   }
@@ -141,7 +138,7 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
 {% endif %}
   for (var d = tid; d < HIDDEN; d = d + WG) {
     let index = base + d;
-    let value = f32(x[index]) * inv * 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;
     local_sq = local_sq + value * value;
   }
   let inv = inverseSqrt(reduce_sum(local_sq, tid) / f32(HIDDEN) + EPSILON);
+{% if writeStats %}
   if (tid == 0u) {
     inv_std_out[row] = inv;
   }
 {% endif %}
   for (var d = tid; d < HIDDEN; d = d + WG) {
     let index = base + d;
+    let value = f32(x[index]) * inv * 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.SimplifiedLayerNormalization",
   "fixtureArrays": {
     "f16_scalar_cast_x": [-1.1103515625, 2.982421875, 1.248046875, -1.8544921875],
     "f16_scalar_cast_scale": [2.015625],
@@ -259,7 +258,7 @@
       "provenance": {
         "source": "onnxruntime/core/providers/cpu/nn/layer_norm_impl.cc",
         "test": "SimplifiedLayerNormalization scalar-scale cast boundary",
-        "notes": "A scalar scale forces the generic shared template and pins the same legacy cast ordering independently of the optimized row path."
       },
       "requires": { "features": ["shader-f16"] },
       "attrs": { "epsilon": 0.00001, "axis": -1 },
@@ -383,6 +382,718 @@
         "y": { "dtype": "float32", "shape": [0, 4], "tolerance": 0.000002 },
         "invStdVar": { "dtype": "float32", "shape": [0, 1], "tolerance": 0.00001 }
       }
     }
   ]
 }

 {
   "fixtureArrays": {
     "f16_scalar_cast_x": [-1.1103515625, 2.982421875, 1.248046875, -1.8544921875],
     "f16_scalar_cast_scale": [2.015625],
       "provenance": {
         "source": "onnxruntime/core/providers/cpu/nn/layer_norm_impl.cc",
         "test": "SimplifiedLayerNormalization scalar-scale cast boundary",
+        "notes": "A scalar scale selects the generic path and requires multiplication before the final float16 cast, independently of optimized row handling."
       },
       "requires": { "features": ["shader-f16"] },
       "attrs": { "epsilon": 0.00001, "axis": -1 },
         "y": { "dtype": "float32", "shape": [0, 4], "tolerance": 0.000002 },
         "invStdVar": { "dtype": "float32", "shape": [0, 1], "tolerance": 0.00001 }
       }
+    },
+    {
+      "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 },
+        "invStdVar": { "shape": [1, 1], "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 },
+        "invStdVar": { "shape": [1, 1], "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 },
+        "invStdVar": { "shape": [1, 1], "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 },
+        "invStdVar": { "shape": [1, 1], "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 },
+        "invStdVar": { "shape": [1, 1], "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 },
+        "invStdVar": { "shape": [1, 1], "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 },
+        "invStdVar": { "shape": [1, 1], "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 },
+        "invStdVar": { "shape": [1, 1], "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 },
+        "invStdVar": { "shape": [1, 1], "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 },
+        "invStdVar": { "shape": [2, 1], "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 },
+        "invStdVar": { "shape": [3, 1], "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 },
+        "invStdVar": { "shape": [4, 1], "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 },
+        "invStdVar": { "shape": [8, 1], "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 },
+        "invStdVar": { "shape": [16, 1], "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 },
+        "invStdVar": { "shape": [32, 1], "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 },
+        "invStdVar": { "shape": [128, 1], "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 },
+        "invStdVar": { "shape": [1, 1], "dtype": "float32", "tolerance": 0.00001 }
+      }
+    },
+    {
+      "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 },
+        "invStdVar": { "shape": [1, 1], "dtype": "float32", "tolerance": 0.00001 }
+      }
+    },
+    {
+      "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 },
+        "invStdVar": { "shape": [1, 1], "dtype": "float32", "tolerance": 0.00001 }
+      }
+    },
+    {
+      "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 },
+        "invStdVar": { "shape": [1, 1], "dtype": "float32", "tolerance": 0.00001 }
+      }
+    },
+    {
+      "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 },
+        "invStdVar": { "shape": [3, 1], "dtype": "float32", "tolerance": 0.00001 }
+      }
+    },
+    {
+      "name": "split_f16-16x65536",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "tunables": {},
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [16, 65536],
+          "data": { "kind": "fillFloat32", "sinStep": 0.000021, "cosStep": 0.000037, "scale": 0.5 }
+        },
+        "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 },
+        "invStdVar": { "shape": [16, 1], "dtype": "float32", "tolerance": 0.00001 }
+      }
+    },
+    {
+      "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 },
+        "invStdVar": { "shape": [2, 1], "dtype": "float32", "tolerance": 0.00001 }
+      }
+    },
+    {
+      "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 },
+        "invStdVar": { "shape": [2, 1], "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 },
+        "invStdVar": { "shape": [1, 1], "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 },
+        "invStdVar": { "shape": [1, 1], "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 },
+        "invStdVar": { "shape": [2, 1], "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 },
+        "invStdVar": { "shape": [2, 1], "dtype": "float32", "tolerance": 0.00001 }
+      }
+    },
+    {
+      "name": "split_f32-1x524288-no-stats",
+      "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-1x2097152-no-stats",
+      "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_f16-1x524288-no-stats",
+      "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-no-stats",
+      "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_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 },
+        "invStdVar": { "dtype": "float32", "shape": [3, 1, 1], "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 },
+        "invStdVar": { "dtype": "float32", "shape": [3, 1, 1], "tolerance": 0.00001 }
+      }
+    },
+    {
+      "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 },
+        "invStdVar": { "dtype": "float32", "shape": [0, 1, 1], "tolerance": 0 }
+      }
     }
   ]
 }