sync 91d990483a17

Browse files

Files changed (7) hide show

README.md +15 -11
build/webgpu/bench.json +3 -4
build/webgpu/manifest.json +237 -671
build/webgpu/metadata.json +26 -10
build/webgpu/norm-skip-row-vec4.wgsl.jinja +19 -12
build/webgpu/norm-skip-row.wgsl.jinja +4 -140
build/webgpu/test.json +31 -7

README.md CHANGED Viewed

@@ -18,19 +18,19 @@ See the [ONNX Runtime `SkipSimplifiedLayerNormalization` contrib-operator spec](
 ## Inputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `input` | `inputT` | `T` | — | — | Input tensor of shape `(token_count, hidden_size)` or `(batch, sequence, hidden_size)`, normalized over the last axis. | required |
-| `skip` | `skipT` | `T` | — | — | Residual tensor of the same shape as `input`, added before normalization. | required |
-| `gamma` | `gammaT` | `T` | `1` | — | 1-D scale tensor with shape `(hidden_size)` applied after normalization. | required |
-| `bias` | `biasT` | `T` | `1` | — | Optional 1-D bias tensor with shape `(hidden_size)` added to the `input + skip` sum. | optional |
 ## Outputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `output` | `outputT` | `T` | same as `input` | same as `input` | Normalized output tensor with the same shape as `input`. | required |
-| `input_skip_bias_sum` | `residualT` | `T` | same as `input` | same as `input` | Sum of `input`, `skip`, and optional `bias` before normalization, with the same shape as `input`. | optional |
 ## Attributes
@@ -52,7 +52,7 @@ Some implementation variants require `shader-f16`. These are route-specific capa
 ## 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
@@ -61,10 +61,14 @@ Some implementation variants require `shader-f16`. These are route-specific capa
 ## 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.

 ## Inputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `inputT` | `input` | `T` | — | — | Input tensor of shape `(token_count, hidden_size)` or `(batch, sequence, hidden_size)`, normalized over the last axis. | required |
+| `skipT` | `skip` | `T` | — | — | Residual tensor of the same shape as `input`, added before normalization. | required |
+| `gammaT` | `gamma` | `T` | `1` | — | 1-D scale tensor with shape `(hidden_size)` applied after normalization. | required |
+| `biasT` | `bias` | `T` | `1` | — | Optional 1-D bias tensor with shape `(hidden_size)` added to the `input + skip` sum. | optional |
 ## Outputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `outputT` | `output` | `T` | same as `inputT` | same as `inputT` | Normalized output tensor with the same shape as `input`. | required |
+| `residualT` | `input_skip_bias_sum` | `T` | same as `inputT` | same as `inputT` | Sum of `input`, `skip`, and optional `bias` before normalization, with the same shape as `input`. | optional |
 ## Attributes
 ## Files
+- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
 - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 ## Use with `@huggingface/kernels`
+```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.2
+```
+Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

build/webgpu/bench.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "com.microsoft.SkipSimplifiedLayerNormalization",
   "cases": [
     {
       "name": "skip-rmsnorm-f32-256x128",
@@ -46,7 +45,7 @@
       "name": "kimi-linear-prefill-f32-64x2304",
       "preset": "model",
       "provenance": {
-        "notes": "Original Kimi Linear hidden_size=2304 at the 64-token recurrent chunk boundary; this is a future vectorized-prefill target because the current graph replays T=1."
       },
       "vars": { "rows": 64, "hidden": 2304 },
       "attrs": { "epsilon": 0.00001 },
@@ -65,7 +64,7 @@
       "name": "kimi-linear-prefill-f32-512x2304",
       "preset": "model",
       "provenance": {
-        "notes": "Original Kimi Linear hidden_size=2304 representative future vectorized-prefill shape; the current graph replays T=1."
       },
       "vars": { "rows": 512, "hidden": 2304 },
       "attrs": { "epsilon": 0.00001 },
@@ -84,7 +83,7 @@
       "name": "kimi-linear-prefill-f32-1024x2304",
       "preset": "model",
       "provenance": {
-        "notes": "Original Kimi Linear hidden_size=2304 representative future vectorized-prefill shape; the current graph replays T=1."
       },
       "vars": { "rows": 1024, "hidden": 2304 },
       "attrs": { "epsilon": 0.00001 },

 {
   "cases": [
     {
       "name": "skip-rmsnorm-f32-256x128",
       "name": "kimi-linear-prefill-f32-64x2304",
       "preset": "model",
       "provenance": {
+        "notes": "Kimi Linear hidden size 2304 at the 64-token recurrent chunk boundary, representing a vectorized-prefill workload."
       },
       "vars": { "rows": 64, "hidden": 2304 },
       "attrs": { "epsilon": 0.00001 },
       "name": "kimi-linear-prefill-f32-512x2304",
       "preset": "model",
       "provenance": {
+        "notes": "Kimi Linear hidden size 2304 over 512 tokens, representing a vectorized-prefill workload."
       },
       "vars": { "rows": 512, "hidden": 2304 },
       "attrs": { "epsilon": 0.00001 },
       "name": "kimi-linear-prefill-f32-1024x2304",
       "preset": "model",
       "provenance": {
+        "notes": "Kimi Linear hidden size 2304 over 1,024 tokens, representing a vectorized-prefill workload."
       },
       "vars": { "rows": 1024, "hidden": 2304 },
       "attrs": { "epsilon": 0.00001 },

build/webgpu/manifest.json CHANGED Viewed

@@ -2,66 +2,31 @@
   "domain": "com.microsoft",
   "name": "SkipSimplifiedLayerNormalization",
   "sinceVersion": 1,
-  "description": "Adds `input` and `skip` (plus optional `bias`), then applies RMS normalization scaled by `gamma`. The optional second output exposes the pre-normalization sum. The schema's training-only mean and inverse-standard-deviation outputs are not implemented.",
-  "inputs": [
-    {
-      "role": "input",
-      "dtype": "T",
-      "description": "Input tensor of shape `(token_count, hidden_size)` or `(batch, sequence, hidden_size)`, normalized over the last axis."
-    },
-    {
-      "role": "skip",
-      "dtype": "T",
-      "description": "Residual tensor of the same shape as `input`, added before normalization."
-    },
-    {
-      "role": "gamma",
-      "dtype": "T",
-      "rank": 1,
-      "description": "1-D scale tensor with shape `(hidden_size)` applied after normalization."
-    },
-    {
-      "role": "bias",
-      "dtype": "T",
-      "rank": 1,
-      "optional": true,
-      "description": "Optional 1-D bias tensor with shape `(hidden_size)` added to the `input + skip` sum."
-    }
-  ],
-  "outputs": [
-    {
-      "role": "output",
       "dtype": "T",
       "rank": "ranks.inputT",
-      "shape": "shapes.inputT",
-      "description": "Normalized output tensor with the same shape as `input`."
-    },
-    {
-      "role": "input_skip_bias_sum",
-      "dtype": "T",
       "optional": true,
-      "rank": "ranks.inputT",
-      "shape": "shapes.inputT",
-      "description": "Sum of `input`, `skip`, and optional `bias` before normalization, with the same shape as `input`."
     }
-  ],
-  "attributes": { "epsilon": 9.999999960041972e-13 },
-  "attributeDescriptions": { "epsilon": "Non-negative epsilon added to the mean square before taking the square root." },
-  "args": {
-    "inputT": { "kind": "tensor", "semantic": "input", "role": "input" },
-    "skipT": { "kind": "tensor", "semantic": "skip", "role": "input" },
-    "gammaT": { "kind": "tensor", "semantic": "gamma", "role": "input" },
-    "biasT": { "kind": "tensor", "semantic": "bias", "role": "input", "required": false },
-    "outputT": { "kind": "tensor", "semantic": "output", "role": "output" },
-    "residualT": { "kind": "tensor", "semantic": "input_skip_bias_sum", "role": "output", "required": false }
   },
   "typeConstraints": { "T": ["float32", "float16"] },
   "derive": {
     "rowCount": "numel(shapes.inputT) / max(1, dim(shapes.inputT, -1))",
     "hiddenSize": "dim(shapes.inputT, -1)",
     "skipWg": "max(1, min(tunables.MAX_WORKGROUP_SIZE, device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX, pow2ceil(hiddenSize)))",
     "skipWgVec4": "max(1, min(tunables.MAX_WORKGROUP_SIZE, device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX, pow2ceil(ceilDiv(hiddenSize, 4))))",
-    "rowDispatchFits": "rowCount <= device.limits.maxComputeWorkgroupsPerDimension * device.limits.maxComputeWorkgroupsPerDimension",
     "normResourcesFit": "skipWg * 8 <= device.limits.maxComputeWorkgroupStorageSize and skipWgVec4 * 8 <= device.limits.maxComputeWorkgroupStorageSize",
     "epsilonOk": "attrs.epsilon >= 0",
     "coreContract": "epsilonOk and (ranks.inputT == 2 or ranks.inputT == 3) and ranks.skipT == ranks.inputT and ranks.gammaT == 1 and ranks.outputT == ranks.inputT and sameShape(shapes.inputT, shapes.skipT) and sameShape(shapes.outputT, shapes.inputT) and dim(shapes.inputT, -1) > 0 and dim(shapes.gammaT, 0) == dim(shapes.inputT, -1)",
@@ -86,439 +51,64 @@
     "f16_no_bias_output_contract": "hasF16 and coreContract and outputOnlyContract and f16MainDtypes and no_bias_contract",
     "f16_bias_output_contract": "hasF16 and coreContract and outputOnlyContract and f16MainDtypes and f16_bias_contract"
   },
-  "bindingSets": {
-    "vec4_no_bias_residual": [
-      {
-        "name": "input",
-        "arg": "inputT",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar"
-      },
-      {
-        "name": "skip",
-        "arg": "skipT",
-        "semantic": "skip",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar"
-      },
-      {
-        "name": "gamma",
-        "arg": "gammaT",
-        "semantic": "gamma",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar",
-        "length": "$HIDDEN_LEN"
-      },
-      {
-        "name": "output",
-        "arg": "outputT",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$vectorScalar"
-      },
-      {
-        "name": "input_skip_bias_sum",
-        "arg": "residualT",
-        "semantic": "input_skip_bias_sum",
-        "buffer": { "type": "storage" },
-        "elementType": "$vectorScalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "rowCount" },
-            {
-              "name": "rowStride",
-              "type": "u32",
-              "value": "max(1, min(rowCount, device.limits.maxComputeWorkgroupsPerDimension))"
-            },
-            { "name": "epsilon", "type": "f32", "value": "attrs.epsilon" }
-          ]
-        }
-      }
-    ],
-    "vec4_bias_residual": [
-      {
-        "name": "input",
-        "arg": "inputT",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar"
-      },
-      {
-        "name": "skip",
-        "arg": "skipT",
-        "semantic": "skip",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar"
-      },
-      {
-        "name": "gamma",
-        "arg": "gammaT",
-        "semantic": "gamma",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar",
-        "length": "$HIDDEN_LEN"
-      },
-      {
-        "name": "bias",
-        "arg": "biasT",
-        "semantic": "bias",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar",
-        "length": "$HIDDEN_LEN"
-      },
-      {
-        "name": "output",
-        "arg": "outputT",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$vectorScalar"
-      },
-      {
-        "name": "input_skip_bias_sum",
-        "arg": "residualT",
-        "semantic": "input_skip_bias_sum",
-        "buffer": { "type": "storage" },
-        "elementType": "$vectorScalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "rowCount" },
-            {
-              "name": "rowStride",
-              "type": "u32",
-              "value": "max(1, min(rowCount, device.limits.maxComputeWorkgroupsPerDimension))"
-            },
-            { "name": "epsilon", "type": "f32", "value": "attrs.epsilon" }
-          ]
-        }
-      }
-    ],
-    "vec4_no_bias_output_only": [
-      {
-        "name": "input",
-        "arg": "inputT",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar"
-      },
-      {
-        "name": "skip",
-        "arg": "skipT",
-        "semantic": "skip",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar"
-      },
-      {
-        "name": "gamma",
-        "arg": "gammaT",
-        "semantic": "gamma",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar",
-        "length": "$HIDDEN_LEN"
-      },
-      {
-        "name": "output",
-        "arg": "outputT",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$vectorScalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "rowCount" },
-            {
-              "name": "rowStride",
-              "type": "u32",
-              "value": "max(1, min(rowCount, device.limits.maxComputeWorkgroupsPerDimension))"
-            },
-            { "name": "epsilon", "type": "f32", "value": "attrs.epsilon" }
-          ]
-        }
-      }
-    ],
-    "vec4_bias_output_only": [
-      {
-        "name": "input",
-        "arg": "inputT",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar"
-      },
-      {
-        "name": "skip",
-        "arg": "skipT",
-        "semantic": "skip",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar"
-      },
-      {
-        "name": "gamma",
-        "arg": "gammaT",
-        "semantic": "gamma",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar",
-        "length": "$HIDDEN_LEN"
-      },
-      {
-        "name": "bias",
-        "arg": "biasT",
-        "semantic": "bias",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar",
-        "length": "$HIDDEN_LEN"
-      },
-      {
-        "name": "output",
-        "arg": "outputT",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$vectorScalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "rowCount" },
-            {
-              "name": "rowStride",
-              "type": "u32",
-              "value": "max(1, min(rowCount, device.limits.maxComputeWorkgroupsPerDimension))"
-            },
-            { "name": "epsilon", "type": "f32", "value": "attrs.epsilon" }
-          ]
-        }
-      }
-    ],
-    "scalar_no_bias_residual": [
-      {
-        "name": "input",
-        "arg": "inputT",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "skip",
-        "arg": "skipT",
-        "semantic": "skip",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "gamma",
-        "arg": "gammaT",
-        "semantic": "gamma",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar",
-        "length": "$HIDDEN_LEN"
-      },
-      {
-        "name": "output",
-        "arg": "outputT",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "input_skip_bias_sum",
-        "arg": "residualT",
-        "semantic": "input_skip_bias_sum",
-        "buffer": { "type": "storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "rowCount" },
-            { "name": "epsilon", "type": "f32", "value": "attrs.epsilon" }
-          ]
-        }
-      }
-    ],
-    "scalar_bias_residual": [
-      {
-        "name": "input",
-        "arg": "inputT",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "skip",
-        "arg": "skipT",
-        "semantic": "skip",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "gamma",
-        "arg": "gammaT",
-        "semantic": "gamma",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar",
-        "length": "$HIDDEN_LEN"
-      },
-      {
-        "name": "bias",
-        "arg": "biasT",
-        "semantic": "bias",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar",
-        "length": "$HIDDEN_LEN"
-      },
-      {
-        "name": "output",
-        "arg": "outputT",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "input_skip_bias_sum",
-        "arg": "residualT",
-        "semantic": "input_skip_bias_sum",
-        "buffer": { "type": "storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "rowCount" },
-            { "name": "epsilon", "type": "f32", "value": "attrs.epsilon" }
-          ]
-        }
-      }
-    ],
-    "scalar_no_bias_output_only": [
-      {
-        "name": "input",
-        "arg": "inputT",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "skip",
-        "arg": "skipT",
-        "semantic": "skip",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "gamma",
-        "arg": "gammaT",
-        "semantic": "gamma",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar",
-        "length": "$HIDDEN_LEN"
-      },
-      {
-        "name": "output",
-        "arg": "outputT",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "rowCount" },
-            { "name": "epsilon", "type": "f32", "value": "attrs.epsilon" }
-          ]
-        }
-      }
-    ],
-    "scalar_bias_output_only": [
-      {
-        "name": "input",
-        "arg": "inputT",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "skip",
-        "arg": "skipT",
-        "semantic": "skip",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "gamma",
-        "arg": "gammaT",
-        "semantic": "gamma",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar",
-        "length": "$HIDDEN_LEN"
-      },
-      {
-        "name": "bias",
-        "arg": "biasT",
-        "semantic": "bias",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar",
-        "length": "$HIDDEN_LEN"
-      },
-      {
-        "name": "output",
-        "arg": "outputT",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "rowCount" },
-            { "name": "epsilon", "type": "f32", "value": "attrs.epsilon" }
-          ]
-        }
-      }
-    ]
   },
-  "tunables": { "MAX_WORKGROUP_SIZE": 256 },
   "variants": [
     {
       "id": "no_bias_vec4_f16",
       "priority": 21,
-      "when": ["f16_no_bias_residual_contract", "vec4Aligned", "normResourcesFit", "rowDispatchFits"],
-      "constants": {
         "scalar": "\"f16\"",
         "vectorScalar": "\"vec4<f16>\"",
         "hasBias": "\"no_bias\" == \"bias\"",
@@ -528,32 +118,30 @@
         {
           "id": "main",
           "name": "SkipSimplifiedLayerNormalization.Vec4",
-          "source": {
-            "shader": "norm-skip-row-vec4.wgsl.jinja",
-            "inputs": {
-              "simplified": true,
-              "hasBias": "\"no_bias\" == \"bias\"",
-              "hasBeta": false,
-              "writeResidualSum": true,
-              "usesF16": true,
-              "hidden": "hiddenSize",
-              "hiddenVec": "hiddenSize / 4",
-              "wg": "skipWgVec4",
-              "vecType": "\"vec4<f16>\"",
-              "useSubgroups": "hasSubgroups"
-            }
           },
-          "subgroupCollectivesWidth": "portable",
-          "bindings": "vec4_no_bias_residual",
-          "dispatch": { "workgroups": "rowCount" }
         }
       ]
     },
     {
       "id": "no_bias_vec4",
       "priority": 20,
-      "when": ["f32_no_bias_residual_contract", "vec4Aligned", "normResourcesFit", "rowDispatchFits"],
-      "constants": {
         "scalar": "\"f32\"",
         "vectorScalar": "\"vec4<f32>\"",
         "hasBias": "\"no_bias\" == \"bias\"",
@@ -563,32 +151,30 @@
         {
           "id": "main",
           "name": "SkipSimplifiedLayerNormalization.Vec4",
-          "source": {
-            "shader": "norm-skip-row-vec4.wgsl.jinja",
-            "inputs": {
-              "simplified": true,
-              "hasBias": "\"no_bias\" == \"bias\"",
-              "hasBeta": false,
-              "writeResidualSum": true,
-              "usesF16": false,
-              "hidden": "hiddenSize",
-              "hiddenVec": "hiddenSize / 4",
-              "wg": "skipWgVec4",
-              "vecType": "\"vec4<f32>\"",
-              "useSubgroups": "hasSubgroups"
-            }
           },
-          "subgroupCollectivesWidth": "portable",
-          "bindings": "vec4_no_bias_residual",
-          "dispatch": { "workgroups": "rowCount" }
         }
       ]
     },
     {
       "id": "no_bias_output_only_vec4",
       "priority": 20,
-      "when": ["f32_no_bias_output_contract", "vec4Aligned", "normResourcesFit", "rowDispatchFits"],
-      "constants": {
         "scalar": "\"f32\"",
         "vectorScalar": "\"vec4<f32>\"",
         "hasBias": "\"no_bias\" == \"bias\"",
@@ -598,32 +184,30 @@
         {
           "id": "main",
           "name": "SkipSimplifiedLayerNormalization.Vec4OutputOnly",
-          "source": {
-            "shader": "norm-skip-row-vec4.wgsl.jinja",
-            "inputs": {
-              "simplified": true,
-              "hasBias": "\"no_bias\" == \"bias\"",
-              "hasBeta": false,
-              "writeResidualSum": false,
-              "usesF16": false,
-              "hidden": "hiddenSize",
-              "hiddenVec": "hiddenSize / 4",
-              "wg": "skipWgVec4",
-              "vecType": "\"vec4<f32>\"",
-              "useSubgroups": "hasSubgroups"
-            }
           },
-          "subgroupCollectivesWidth": "portable",
-          "bindings": "vec4_no_bias_output_only",
-          "dispatch": { "workgroups": "rowCount" }
         }
       ]
     },
     {
       "id": "no_bias_output_only_vec4_f16",
       "priority": 21,
-      "when": ["f16_no_bias_output_contract", "vec4Aligned", "normResourcesFit", "rowDispatchFits"],
-      "constants": {
         "scalar": "\"f16\"",
         "vectorScalar": "\"vec4<f16>\"",
         "hasBias": "\"no_bias\" == \"bias\"",
@@ -633,39 +217,36 @@
         {
           "id": "main",
           "name": "SkipSimplifiedLayerNormalization.Vec4OutputOnly",
-          "source": {
-            "shader": "norm-skip-row-vec4.wgsl.jinja",
-            "inputs": {
-              "simplified": true,
-              "hasBias": "\"no_bias\" == \"bias\"",
-              "hasBeta": false,
-              "writeResidualSum": false,
-              "usesF16": true,
-              "hidden": "hiddenSize",
-              "hiddenVec": "hiddenSize / 4",
-              "wg": "skipWgVec4",
-              "vecType": "\"vec4<f16>\"",
-              "useSubgroups": "hasSubgroups"
-            }
           },
-          "subgroupCollectivesWidth": "portable",
-          "bindings": "vec4_no_bias_output_only",
-          "dispatch": { "workgroups": "rowCount" }
         }
       ]
     },
     {
       "id": "no_bias",
       "priority": 0,
-      "when": ["f32_no_bias_residual_contract", "normResourcesFit", "rowDispatchFits"],
-      "constants": {
         "simplified": true,
         "useSubgroups": false,
         "hasBeta": false,
         "writeResidualSum": true,
         "hasBias": "\"no_bias\" == \"bias\"",
         "scalar": "\"f32\"",
-        "hiddenSize": "dim(shapes.inputT, -1)",
         "workgroupSize": "skipWg",
         "HIDDEN_LEN": "hiddenSize"
       },
@@ -674,17 +255,17 @@
           "id": "main",
           "name": "SkipSimplifiedLayerNormalization",
           "shader": "norm-skip-row.wgsl.jinja",
-          "bindings": "scalar_no_bias_residual",
-          "dispatch": { "workgroups": "rowCount" }
         }
       ]
     },
     {
       "id": "no_bias_f16",
-      "requires": { "features": ["shader-f16"] },
       "priority": 0,
-      "when": ["f16_no_bias_residual_contract", "normResourcesFit", "rowDispatchFits"],
-      "constants": {
         "simplified": true,
         "useSubgroups": false,
         "hasBeta": false,
@@ -692,7 +273,6 @@
         "hasBias": "\"no_bias\" == \"bias\"",
         "scalar": "\"f16\"",
         "usesF16": true,
-        "hiddenSize": "dim(shapes.inputT, -1)",
         "workgroupSize": "skipWg",
         "HIDDEN_LEN": "hiddenSize"
       },
@@ -701,17 +281,17 @@
           "id": "main",
           "name": "SkipSimplifiedLayerNormalization",
           "shader": "norm-skip-row.wgsl.jinja",
-          "bindings": "scalar_no_bias_residual",
-          "dispatch": { "workgroups": "rowCount" }
         }
       ]
     },
     {
       "id": "no_bias_output_only_f16",
-      "requires": { "features": ["shader-f16"] },
       "priority": 0,
-      "when": ["f16_no_bias_output_contract", "normResourcesFit", "rowDispatchFits"],
-      "constants": {
         "simplified": true,
         "useSubgroups": false,
         "hasBeta": false,
@@ -719,7 +299,6 @@
         "hasBias": "\"no_bias\" == \"bias\"",
         "scalar": "\"f16\"",
         "usesF16": true,
-        "hiddenSize": "dim(shapes.inputT, -1)",
         "workgroupSize": "skipWg",
         "HIDDEN_LEN": "hiddenSize"
       },
@@ -728,23 +307,22 @@
           "id": "main",
           "name": "SkipSimplifiedLayerNormalization.OutputOnly",
           "shader": "norm-skip-row.wgsl.jinja",
-          "bindings": "scalar_no_bias_output_only",
-          "dispatch": { "workgroups": "rowCount" }
         }
       ]
     },
     {
       "id": "no_bias_output_only",
       "priority": 0,
-      "when": ["f32_no_bias_output_contract", "normResourcesFit", "rowDispatchFits"],
-      "constants": {
         "simplified": true,
         "useSubgroups": false,
         "hasBeta": false,
         "writeResidualSum": false,
         "hasBias": "\"no_bias\" == \"bias\"",
         "scalar": "\"f32\"",
-        "hiddenSize": "dim(shapes.inputT, -1)",
         "workgroupSize": "skipWg",
         "HIDDEN_LEN": "hiddenSize"
       },
@@ -753,16 +331,16 @@
           "id": "main",
           "name": "SkipSimplifiedLayerNormalization.OutputOnly",
           "shader": "norm-skip-row.wgsl.jinja",
-          "bindings": "scalar_no_bias_output_only",
-          "dispatch": { "workgroups": "rowCount" }
         }
       ]
     },
     {
       "id": "bias_vec4_f16",
       "priority": 21,
-      "when": ["f16_bias_residual_contract", "vec4Aligned", "normResourcesFit", "rowDispatchFits"],
-      "constants": {
         "scalar": "\"f16\"",
         "vectorScalar": "\"vec4<f16>\"",
         "hasBias": "\"bias\" == \"bias\"",
@@ -772,32 +350,30 @@
         {
           "id": "main",
           "name": "SkipSimplifiedLayerNormalization.Vec4",
-          "source": {
-            "shader": "norm-skip-row-vec4.wgsl.jinja",
-            "inputs": {
-              "simplified": true,
-              "hasBias": "\"bias\" == \"bias\"",
-              "hasBeta": false,
-              "writeResidualSum": true,
-              "usesF16": true,
-              "hidden": "hiddenSize",
-              "hiddenVec": "hiddenSize / 4",
-              "wg": "skipWgVec4",
-              "vecType": "\"vec4<f16>\"",
-              "useSubgroups": "hasSubgroups"
-            }
           },
-          "subgroupCollectivesWidth": "portable",
-          "bindings": "vec4_bias_residual",
-          "dispatch": { "workgroups": "rowCount" }
         }
       ]
     },
     {
       "id": "bias_vec4",
       "priority": 20,
-      "when": ["f32_bias_residual_contract", "vec4Aligned", "normResourcesFit", "rowDispatchFits"],
-      "constants": {
         "scalar": "\"f32\"",
         "vectorScalar": "\"vec4<f32>\"",
         "hasBias": "\"bias\" == \"bias\"",
@@ -807,32 +383,30 @@
         {
           "id": "main",
           "name": "SkipSimplifiedLayerNormalization.Vec4",
-          "source": {
-            "shader": "norm-skip-row-vec4.wgsl.jinja",
-            "inputs": {
-              "simplified": true,
-              "hasBias": "\"bias\" == \"bias\"",
-              "hasBeta": false,
-              "writeResidualSum": true,
-              "usesF16": false,
-              "hidden": "hiddenSize",
-              "hiddenVec": "hiddenSize / 4",
-              "wg": "skipWgVec4",
-              "vecType": "\"vec4<f32>\"",
-              "useSubgroups": "hasSubgroups"
-            }
           },
-          "subgroupCollectivesWidth": "portable",
-          "bindings": "vec4_bias_residual",
-          "dispatch": { "workgroups": "rowCount" }
         }
       ]
     },
     {
       "id": "bias_output_only_vec4",
       "priority": 20,
-      "when": ["f32_bias_output_contract", "vec4Aligned", "normResourcesFit", "rowDispatchFits"],
-      "constants": {
         "scalar": "\"f32\"",
         "vectorScalar": "\"vec4<f32>\"",
         "hasBias": "\"bias\" == \"bias\"",
@@ -842,32 +416,30 @@
         {
           "id": "main",
           "name": "SkipSimplifiedLayerNormalization.Vec4OutputOnly",
-          "source": {
-            "shader": "norm-skip-row-vec4.wgsl.jinja",
-            "inputs": {
-              "simplified": true,
-              "hasBias": "\"bias\" == \"bias\"",
-              "hasBeta": false,
-              "writeResidualSum": false,
-              "usesF16": false,
-              "hidden": "hiddenSize",
-              "hiddenVec": "hiddenSize / 4",
-              "wg": "skipWgVec4",
-              "vecType": "\"vec4<f32>\"",
-              "useSubgroups": "hasSubgroups"
-            }
           },
-          "subgroupCollectivesWidth": "portable",
-          "bindings": "vec4_bias_output_only",
-          "dispatch": { "workgroups": "rowCount" }
         }
       ]
     },
     {
       "id": "bias_output_only_vec4_f16",
       "priority": 21,
-      "when": ["f16_bias_output_contract", "vec4Aligned", "normResourcesFit", "rowDispatchFits"],
-      "constants": {
         "scalar": "\"f16\"",
         "vectorScalar": "\"vec4<f16>\"",
         "hasBias": "\"bias\" == \"bias\"",
@@ -877,39 +449,36 @@
         {
           "id": "main",
           "name": "SkipSimplifiedLayerNormalization.Vec4OutputOnly",
-          "source": {
-            "shader": "norm-skip-row-vec4.wgsl.jinja",
-            "inputs": {
-              "simplified": true,
-              "hasBias": "\"bias\" == \"bias\"",
-              "hasBeta": false,
-              "writeResidualSum": false,
-              "usesF16": true,
-              "hidden": "hiddenSize",
-              "hiddenVec": "hiddenSize / 4",
-              "wg": "skipWgVec4",
-              "vecType": "\"vec4<f16>\"",
-              "useSubgroups": "hasSubgroups"
-            }
           },
-          "subgroupCollectivesWidth": "portable",
-          "bindings": "vec4_bias_output_only",
-          "dispatch": { "workgroups": "rowCount" }
         }
       ]
     },
     {
       "id": "bias",
       "priority": 0,
-      "when": ["f32_bias_residual_contract", "normResourcesFit", "rowDispatchFits"],
-      "constants": {
         "simplified": true,
         "useSubgroups": false,
         "hasBeta": false,
         "writeResidualSum": true,
         "hasBias": "\"bias\" == \"bias\"",
         "scalar": "\"f32\"",
-        "hiddenSize": "dim(shapes.inputT, -1)",
         "workgroupSize": "skipWg",
         "HIDDEN_LEN": "hiddenSize"
       },
@@ -918,17 +487,17 @@
           "id": "main",
           "name": "SkipSimplifiedLayerNormalization",
           "shader": "norm-skip-row.wgsl.jinja",
-          "bindings": "scalar_bias_residual",
-          "dispatch": { "workgroups": "rowCount" }
         }
       ]
     },
     {
       "id": "bias_f16",
-      "requires": { "features": ["shader-f16"] },
       "priority": 0,
-      "when": ["f16_bias_residual_contract", "normResourcesFit", "rowDispatchFits"],
-      "constants": {
         "simplified": true,
         "useSubgroups": false,
         "hasBeta": false,
@@ -936,7 +505,6 @@
         "hasBias": "\"bias\" == \"bias\"",
         "scalar": "\"f16\"",
         "usesF16": true,
-        "hiddenSize": "dim(shapes.inputT, -1)",
         "workgroupSize": "skipWg",
         "HIDDEN_LEN": "hiddenSize"
       },
@@ -945,17 +513,17 @@
           "id": "main",
           "name": "SkipSimplifiedLayerNormalization",
           "shader": "norm-skip-row.wgsl.jinja",
-          "bindings": "scalar_bias_residual",
-          "dispatch": { "workgroups": "rowCount" }
         }
       ]
     },
     {
       "id": "bias_output_only_f16",
-      "requires": { "features": ["shader-f16"] },
       "priority": 0,
-      "when": ["f16_bias_output_contract", "normResourcesFit", "rowDispatchFits"],
-      "constants": {
         "simplified": true,
         "useSubgroups": false,
         "hasBeta": false,
@@ -963,7 +531,6 @@
         "hasBias": "\"bias\" == \"bias\"",
         "scalar": "\"f16\"",
         "usesF16": true,
-        "hiddenSize": "dim(shapes.inputT, -1)",
         "workgroupSize": "skipWg",
         "HIDDEN_LEN": "hiddenSize"
       },
@@ -972,23 +539,22 @@
           "id": "main",
           "name": "SkipSimplifiedLayerNormalization.OutputOnly",
           "shader": "norm-skip-row.wgsl.jinja",
-          "bindings": "scalar_bias_output_only",
-          "dispatch": { "workgroups": "rowCount" }
         }
       ]
     },
     {
       "id": "bias_output_only",
       "priority": 0,
-      "when": ["f32_bias_output_contract", "normResourcesFit", "rowDispatchFits"],
-      "constants": {
         "simplified": true,
         "useSubgroups": false,
         "hasBeta": false,
         "writeResidualSum": false,
         "hasBias": "\"bias\" == \"bias\"",
         "scalar": "\"f32\"",
-        "hiddenSize": "dim(shapes.inputT, -1)",
         "workgroupSize": "skipWg",
         "HIDDEN_LEN": "hiddenSize"
       },
@@ -997,8 +563,8 @@
           "id": "main",
           "name": "SkipSimplifiedLayerNormalization.OutputOnly",
           "shader": "norm-skip-row.wgsl.jinja",
-          "bindings": "scalar_bias_output_only",
-          "dispatch": { "workgroups": "rowCount" }
         }
       ]
     }

   "domain": "com.microsoft",
   "name": "SkipSimplifiedLayerNormalization",
   "sinceVersion": 1,
+  "inputs": {
+    "inputT": { "onnx": "input", "dtype": "T" },
+    "skipT": { "onnx": "skip", "dtype": "T" },
+    "gammaT": { "onnx": "gamma", "dtype": "T", "rank": 1 },
+    "biasT": { "onnx": "bias", "dtype": "T", "rank": 1, "optional": true }
+  },
+  "outputs": {
+    "outputT": { "onnx": "output", "dtype": "T", "rank": "ranks.inputT", "shape": "shapes.inputT" },
+    "residualT": {
+      "onnx": "input_skip_bias_sum",
       "dtype": "T",
       "rank": "ranks.inputT",
       "optional": true,
+      "shape": "shapes.inputT"
     }
   },
+  "attributes": { "epsilon": { "default": 9.999999960041972e-13 } },
   "typeConstraints": { "T": ["float32", "float16"] },
+  "tunables": { "MAX_WORKGROUP_SIZE": { "default": 256 } },
   "derive": {
     "rowCount": "numel(shapes.inputT) / max(1, dim(shapes.inputT, -1))",
     "hiddenSize": "dim(shapes.inputT, -1)",
     "skipWg": "max(1, min(tunables.MAX_WORKGROUP_SIZE, device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX, pow2ceil(hiddenSize)))",
     "skipWgVec4": "max(1, min(tunables.MAX_WORKGROUP_SIZE, device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX, pow2ceil(ceilDiv(hiddenSize, 4))))",
+    "rowDispatchFits": "rowCount <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
     "normResourcesFit": "skipWg * 8 <= device.limits.maxComputeWorkgroupStorageSize and skipWgVec4 * 8 <= device.limits.maxComputeWorkgroupStorageSize",
     "epsilonOk": "attrs.epsilon >= 0",
     "coreContract": "epsilonOk and (ranks.inputT == 2 or ranks.inputT == 3) and ranks.skipT == ranks.inputT and ranks.gammaT == 1 and ranks.outputT == ranks.inputT and sameShape(shapes.inputT, shapes.skipT) and sameShape(shapes.outputT, shapes.inputT) and dim(shapes.inputT, -1) > 0 and dim(shapes.gammaT, 0) == dim(shapes.inputT, -1)",
     "f16_no_bias_output_contract": "hasF16 and coreContract and outputOnlyContract and f16MainDtypes and no_bias_contract",
     "f16_bias_output_contract": "hasF16 and coreContract and outputOnlyContract and f16MainDtypes and f16_bias_contract"
   },
+  "when": ["normResourcesFit", "rowDispatchFits"],
+  "bindings": {
+    "input": { "arg": "inputT", "buffer": "read-only-storage", "elementType": "$vectorScalar" },
+    "skip": { "arg": "skipT", "buffer": "read-only-storage", "elementType": "$vectorScalar" },
+    "gamma": { "arg": "gammaT", "buffer": "read-only-storage", "elementType": "$vectorScalar", "length": "$HIDDEN_LEN" },
+    "output": { "arg": "outputT", "buffer": "storage", "elementType": "$vectorScalar" },
+    "input_skip_bias_sum": { "arg": "residualT", "buffer": "storage", "elementType": "$vectorScalar" },
+    "params": {
+      "buffer": "uniform",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "rowCount" },
+        {
+          "name": "rowStride",
+          "type": "u32",
+          "value": "max(1, min(rowCount, min(device.limits.maxComputeWorkgroupsPerDimension, 65535)))"
+        },
+        { "name": "epsilon", "type": "f32", "value": "attrs.epsilon" }
+      ]
+    },
+    "bias": { "arg": "biasT", "buffer": "read-only-storage", "elementType": "$vectorScalar", "length": "$HIDDEN_LEN" },
+    "input_2": { "arg": "inputT", "name": "input", "buffer": "read-only-storage", "elementType": "$scalar" },
+    "skip_2": { "arg": "skipT", "name": "skip", "buffer": "read-only-storage", "elementType": "$scalar" },
+    "gamma_2": {
+      "arg": "gammaT",
+      "name": "gamma",
+      "buffer": "read-only-storage",
+      "elementType": "$scalar",
+      "length": "$HIDDEN_LEN"
+    },
+    "output_2": { "arg": "outputT", "name": "output", "buffer": "storage", "elementType": "$scalar" },
+    "input_skip_bias_sum_2": {
+      "arg": "residualT",
+      "name": "input_skip_bias_sum",
+      "buffer": "storage",
+      "elementType": "$scalar"
+    },
+    "params_2": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "rowCount" },
+        { "name": "epsilon", "type": "f32", "value": "attrs.epsilon" }
+      ]
+    },
+    "bias_2": {
+      "arg": "biasT",
+      "name": "bias",
+      "buffer": "read-only-storage",
+      "elementType": "$scalar",
+      "length": "$HIDDEN_LEN"
+    }
   },
   "variants": [
     {
       "id": "no_bias_vec4_f16",
       "priority": 21,
+      "when": ["f16_no_bias_residual_contract", "vec4Aligned"],
+      "derive": {
         "scalar": "\"f16\"",
         "vectorScalar": "\"vec4<f16>\"",
         "hasBias": "\"no_bias\" == \"bias\"",
         {
           "id": "main",
           "name": "SkipSimplifiedLayerNormalization.Vec4",
+          "shader": "norm-skip-row-vec4.wgsl.jinja",
+          "derive": {
+            "simplified": true,
+            "hasBias": "\"no_bias\" == \"bias\"",
+            "hasBeta": false,
+            "writeResidualSum": true,
+            "usesF16Spec": true,
+            "hidden": "hiddenSize",
+            "hiddenVec": "hiddenSize / 4",
+            "wg": "skipWgVec4",
+            "vecType": "\"vec4<f16>\"",
+            "useSubgroups": "hasSubgroups"
           },
+          "bindings": ["input", "skip", "gamma", "output", "input_skip_bias_sum", "params"],
+          "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "no_bias_vec4",
       "priority": 20,
+      "when": ["f32_no_bias_residual_contract", "vec4Aligned"],
+      "derive": {
         "scalar": "\"f32\"",
         "vectorScalar": "\"vec4<f32>\"",
         "hasBias": "\"no_bias\" == \"bias\"",
         {
           "id": "main",
           "name": "SkipSimplifiedLayerNormalization.Vec4",
+          "shader": "norm-skip-row-vec4.wgsl.jinja",
+          "derive": {
+            "simplified": true,
+            "hasBias": "\"no_bias\" == \"bias\"",
+            "hasBeta": false,
+            "writeResidualSum": true,
+            "usesF16Spec": false,
+            "hidden": "hiddenSize",
+            "hiddenVec": "hiddenSize / 4",
+            "wg": "skipWgVec4",
+            "vecType": "\"vec4<f32>\"",
+            "useSubgroups": "hasSubgroups"
           },
+          "bindings": ["input", "skip", "gamma", "output", "input_skip_bias_sum", "params"],
+          "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "no_bias_output_only_vec4",
       "priority": 20,
+      "when": ["f32_no_bias_output_contract", "vec4Aligned"],
+      "derive": {
         "scalar": "\"f32\"",
         "vectorScalar": "\"vec4<f32>\"",
         "hasBias": "\"no_bias\" == \"bias\"",
         {
           "id": "main",
           "name": "SkipSimplifiedLayerNormalization.Vec4OutputOnly",
+          "shader": "norm-skip-row-vec4.wgsl.jinja",
+          "derive": {
+            "simplified": true,
+            "hasBias": "\"no_bias\" == \"bias\"",
+            "hasBeta": false,
+            "writeResidualSum": false,
+            "usesF16Spec": false,
+            "hidden": "hiddenSize",
+            "hiddenVec": "hiddenSize / 4",
+            "wg": "skipWgVec4",
+            "vecType": "\"vec4<f32>\"",
+            "useSubgroups": "hasSubgroups"
           },
+          "bindings": ["input", "skip", "gamma", "output", "params"],
+          "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "no_bias_output_only_vec4_f16",
       "priority": 21,
+      "when": ["f16_no_bias_output_contract", "vec4Aligned"],
+      "derive": {
         "scalar": "\"f16\"",
         "vectorScalar": "\"vec4<f16>\"",
         "hasBias": "\"no_bias\" == \"bias\"",
         {
           "id": "main",
           "name": "SkipSimplifiedLayerNormalization.Vec4OutputOnly",
+          "shader": "norm-skip-row-vec4.wgsl.jinja",
+          "derive": {
+            "simplified": true,
+            "hasBias": "\"no_bias\" == \"bias\"",
+            "hasBeta": false,
+            "writeResidualSum": false,
+            "usesF16Spec": true,
+            "hidden": "hiddenSize",
+            "hiddenVec": "hiddenSize / 4",
+            "wg": "skipWgVec4",
+            "vecType": "\"vec4<f16>\"",
+            "useSubgroups": "hasSubgroups"
           },
+          "bindings": ["input", "skip", "gamma", "output", "params"],
+          "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "no_bias",
       "priority": 0,
+      "when": ["f32_no_bias_residual_contract"],
+      "derive": {
         "simplified": true,
         "useSubgroups": false,
         "hasBeta": false,
         "writeResidualSum": true,
         "hasBias": "\"no_bias\" == \"bias\"",
         "scalar": "\"f32\"",
         "workgroupSize": "skipWg",
         "HIDDEN_LEN": "hiddenSize"
       },
           "id": "main",
           "name": "SkipSimplifiedLayerNormalization",
           "shader": "norm-skip-row.wgsl.jinja",
+          "bindings": ["input_2", "skip_2", "gamma_2", "output_2", "input_skip_bias_sum_2", "params_2"],
+          "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 }
         }
       ]
     },
     {
       "id": "no_bias_f16",
       "priority": 0,
+      "when": ["f16_no_bias_residual_contract"],
+      "requires": { "features": ["shader-f16"] },
+      "derive": {
         "simplified": true,
         "useSubgroups": false,
         "hasBeta": false,
         "hasBias": "\"no_bias\" == \"bias\"",
         "scalar": "\"f16\"",
         "usesF16": true,
         "workgroupSize": "skipWg",
         "HIDDEN_LEN": "hiddenSize"
       },
           "id": "main",
           "name": "SkipSimplifiedLayerNormalization",
           "shader": "norm-skip-row.wgsl.jinja",
+          "bindings": ["input_2", "skip_2", "gamma_2", "output_2", "input_skip_bias_sum_2", "params_2"],
+          "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 }
         }
       ]
     },
     {
       "id": "no_bias_output_only_f16",
       "priority": 0,
+      "when": ["f16_no_bias_output_contract"],
+      "requires": { "features": ["shader-f16"] },
+      "derive": {
         "simplified": true,
         "useSubgroups": false,
         "hasBeta": false,
         "hasBias": "\"no_bias\" == \"bias\"",
         "scalar": "\"f16\"",
         "usesF16": true,
         "workgroupSize": "skipWg",
         "HIDDEN_LEN": "hiddenSize"
       },
           "id": "main",
           "name": "SkipSimplifiedLayerNormalization.OutputOnly",
           "shader": "norm-skip-row.wgsl.jinja",
+          "bindings": ["input_2", "skip_2", "gamma_2", "output_2", "params_2"],
+          "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 }
         }
       ]
     },
     {
       "id": "no_bias_output_only",
       "priority": 0,
+      "when": ["f32_no_bias_output_contract"],
+      "derive": {
         "simplified": true,
         "useSubgroups": false,
         "hasBeta": false,
         "writeResidualSum": false,
         "hasBias": "\"no_bias\" == \"bias\"",
         "scalar": "\"f32\"",
         "workgroupSize": "skipWg",
         "HIDDEN_LEN": "hiddenSize"
       },
           "id": "main",
           "name": "SkipSimplifiedLayerNormalization.OutputOnly",
           "shader": "norm-skip-row.wgsl.jinja",
+          "bindings": ["input_2", "skip_2", "gamma_2", "output_2", "params_2"],
+          "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 }
         }
       ]
     },
     {
       "id": "bias_vec4_f16",
       "priority": 21,
+      "when": ["f16_bias_residual_contract", "vec4Aligned"],
+      "derive": {
         "scalar": "\"f16\"",
         "vectorScalar": "\"vec4<f16>\"",
         "hasBias": "\"bias\" == \"bias\"",
         {
           "id": "main",
           "name": "SkipSimplifiedLayerNormalization.Vec4",
+          "shader": "norm-skip-row-vec4.wgsl.jinja",
+          "derive": {
+            "simplified": true,
+            "hasBias": "\"bias\" == \"bias\"",
+            "hasBeta": false,
+            "writeResidualSum": true,
+            "usesF16Spec": true,
+            "hidden": "hiddenSize",
+            "hiddenVec": "hiddenSize / 4",
+            "wg": "skipWgVec4",
+            "vecType": "\"vec4<f16>\"",
+            "useSubgroups": "hasSubgroups"
           },
+          "bindings": ["input", "skip", "gamma", "bias", "output", "input_skip_bias_sum", "params"],
+          "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "bias_vec4",
       "priority": 20,
+      "when": ["f32_bias_residual_contract", "vec4Aligned"],
+      "derive": {
         "scalar": "\"f32\"",
         "vectorScalar": "\"vec4<f32>\"",
         "hasBias": "\"bias\" == \"bias\"",
         {
           "id": "main",
           "name": "SkipSimplifiedLayerNormalization.Vec4",
+          "shader": "norm-skip-row-vec4.wgsl.jinja",
+          "derive": {
+            "simplified": true,
+            "hasBias": "\"bias\" == \"bias\"",
+            "hasBeta": false,
+            "writeResidualSum": true,
+            "usesF16Spec": false,
+            "hidden": "hiddenSize",
+            "hiddenVec": "hiddenSize / 4",
+            "wg": "skipWgVec4",
+            "vecType": "\"vec4<f32>\"",
+            "useSubgroups": "hasSubgroups"
           },
+          "bindings": ["input", "skip", "gamma", "bias", "output", "input_skip_bias_sum", "params"],
+          "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "bias_output_only_vec4",
       "priority": 20,
+      "when": ["f32_bias_output_contract", "vec4Aligned"],
+      "derive": {
         "scalar": "\"f32\"",
         "vectorScalar": "\"vec4<f32>\"",
         "hasBias": "\"bias\" == \"bias\"",
         {
           "id": "main",
           "name": "SkipSimplifiedLayerNormalization.Vec4OutputOnly",
+          "shader": "norm-skip-row-vec4.wgsl.jinja",
+          "derive": {
+            "simplified": true,
+            "hasBias": "\"bias\" == \"bias\"",
+            "hasBeta": false,
+            "writeResidualSum": false,
+            "usesF16Spec": false,
+            "hidden": "hiddenSize",
+            "hiddenVec": "hiddenSize / 4",
+            "wg": "skipWgVec4",
+            "vecType": "\"vec4<f32>\"",
+            "useSubgroups": "hasSubgroups"
           },
+          "bindings": ["input", "skip", "gamma", "bias", "output", "params"],
+          "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "bias_output_only_vec4_f16",
       "priority": 21,
+      "when": ["f16_bias_output_contract", "vec4Aligned"],
+      "derive": {
         "scalar": "\"f16\"",
         "vectorScalar": "\"vec4<f16>\"",
         "hasBias": "\"bias\" == \"bias\"",
         {
           "id": "main",
           "name": "SkipSimplifiedLayerNormalization.Vec4OutputOnly",
+          "shader": "norm-skip-row-vec4.wgsl.jinja",
+          "derive": {
+            "simplified": true,
+            "hasBias": "\"bias\" == \"bias\"",
+            "hasBeta": false,
+            "writeResidualSum": false,
+            "usesF16Spec": true,
+            "hidden": "hiddenSize",
+            "hiddenVec": "hiddenSize / 4",
+            "wg": "skipWgVec4",
+            "vecType": "\"vec4<f16>\"",
+            "useSubgroups": "hasSubgroups"
           },
+          "bindings": ["input", "skip", "gamma", "bias", "output", "params"],
+          "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "bias",
       "priority": 0,
+      "when": ["f32_bias_residual_contract"],
+      "derive": {
         "simplified": true,
         "useSubgroups": false,
         "hasBeta": false,
         "writeResidualSum": true,
         "hasBias": "\"bias\" == \"bias\"",
         "scalar": "\"f32\"",
         "workgroupSize": "skipWg",
         "HIDDEN_LEN": "hiddenSize"
       },
           "id": "main",
           "name": "SkipSimplifiedLayerNormalization",
           "shader": "norm-skip-row.wgsl.jinja",
+          "bindings": ["input_2", "skip_2", "gamma_2", "bias_2", "output_2", "input_skip_bias_sum_2", "params_2"],
+          "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 }
         }
       ]
     },
     {
       "id": "bias_f16",
       "priority": 0,
+      "when": ["f16_bias_residual_contract"],
+      "requires": { "features": ["shader-f16"] },
+      "derive": {
         "simplified": true,
         "useSubgroups": false,
         "hasBeta": false,
         "hasBias": "\"bias\" == \"bias\"",
         "scalar": "\"f16\"",
         "usesF16": true,
         "workgroupSize": "skipWg",
         "HIDDEN_LEN": "hiddenSize"
       },
           "id": "main",
           "name": "SkipSimplifiedLayerNormalization",
           "shader": "norm-skip-row.wgsl.jinja",
+          "bindings": ["input_2", "skip_2", "gamma_2", "bias_2", "output_2", "input_skip_bias_sum_2", "params_2"],
+          "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 }
         }
       ]
     },
     {
       "id": "bias_output_only_f16",
       "priority": 0,
+      "when": ["f16_bias_output_contract"],
+      "requires": { "features": ["shader-f16"] },
+      "derive": {
         "simplified": true,
         "useSubgroups": false,
         "hasBeta": false,
         "hasBias": "\"bias\" == \"bias\"",
         "scalar": "\"f16\"",
         "usesF16": true,
         "workgroupSize": "skipWg",
         "HIDDEN_LEN": "hiddenSize"
       },
           "id": "main",
           "name": "SkipSimplifiedLayerNormalization.OutputOnly",
           "shader": "norm-skip-row.wgsl.jinja",
+          "bindings": ["input_2", "skip_2", "gamma_2", "bias_2", "output_2", "params_2"],
+          "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 }
         }
       ]
     },
     {
       "id": "bias_output_only",
       "priority": 0,
+      "when": ["f32_bias_output_contract"],
+      "derive": {
         "simplified": true,
         "useSubgroups": false,
         "hasBeta": false,
         "writeResidualSum": false,
         "hasBias": "\"bias\" == \"bias\"",
         "scalar": "\"f32\"",
         "workgroupSize": "skipWg",
         "HIDDEN_LEN": "hiddenSize"
       },
           "id": "main",
           "name": "SkipSimplifiedLayerNormalization.OutputOnly",
           "shader": "norm-skip-row.wgsl.jinja",
+          "bindings": ["input_2", "skip_2", "gamma_2", "bias_2", "output_2", "params_2"],
+          "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 }
         }
       ]
     }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,23 +1,39 @@
 {
   "name": "com.microsoft.SkipSimplifiedLayerNormalization",
-  "id": "_com_microsoft_skipsimplifiedlayernormalization_webgpu_6026e4e",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "Qjd8vGH/qoXU3rXEQhkz+xTQK2qSM5Kl4BHnWY/nzeQ=",
-      "manifest.json": "0QZ+D26SsY6uIHFcZV326mQMPoEYAClEUEifNI/9jms=",
-      "norm-skip-row-vec4.wgsl.jinja": "L7sH/FpRWVXbyaGm4QTjBAaLMacnNmwb0dbt2pKie5s=",
-      "norm-skip-row.wgsl.jinja": "bo/2mKHvO3Xkb7rbXnHkw/csvMsxID6bhdF8Wv7efKM=",
-      "test.json": "oQayex/aOYAhZXp9ESTPlwVGivSDeuEumSnU6kNNUR0="
     }
   },
-  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
   "webgpu": {
-    "manifestSpec": "1.0",
-    "specialized": true,
-    "opPath": "ops/com.microsoft.SkipSimplifiedLayerNormalization"
   }
 }

 {
   "name": "com.microsoft.SkipSimplifiedLayerNormalization",
+  "id": "_com_microsoft_skipsimplifiedlayernormalization_webgpu_b8ecf51",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "w8dZhYzUtAD/LB4ByqSaHZ6t6a1LKgH6Fb7sFyVZLY8=",
+      "manifest.json": "Z3UvZaHGuT5XUJRtn5Cb07il2Jc1HM60BXm8GXxeWyY=",
+      "norm-skip-row-vec4.wgsl.jinja": "gTtEoczLSil0/2T/beWedsZzghnBsunoUp5zcKKAfpM=",
+      "norm-skip-row.wgsl.jinja": "HwkrJ+4YZjuFEQ87AvYOJX89/G3JeFc/CfOiU0BcPcU=",
+      "test.json": "E137wUKjS95iD915YA7yUZ4fppQlh8c/s3b5V7spWas="
     }
   },
+  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
   "webgpu": {
+    "manifestSpec": "2.0",
+    "variants": {
+      "no_bias_vec4_f16": ["norm-skip-row-vec4.wgsl.jinja"],
+      "no_bias_vec4": ["norm-skip-row-vec4.wgsl.jinja"],
+      "no_bias_output_only_vec4": ["norm-skip-row-vec4.wgsl.jinja"],
+      "no_bias_output_only_vec4_f16": ["norm-skip-row-vec4.wgsl.jinja"],
+      "no_bias": ["norm-skip-row.wgsl.jinja"],
+      "no_bias_f16": ["norm-skip-row.wgsl.jinja"],
+      "no_bias_output_only_f16": ["norm-skip-row.wgsl.jinja"],
+      "no_bias_output_only": ["norm-skip-row.wgsl.jinja"],
+      "bias_vec4_f16": ["norm-skip-row-vec4.wgsl.jinja"],
+      "bias_vec4": ["norm-skip-row-vec4.wgsl.jinja"],
+      "bias_output_only_vec4": ["norm-skip-row-vec4.wgsl.jinja"],
+      "bias_output_only_vec4_f16": ["norm-skip-row-vec4.wgsl.jinja"],
+      "bias": ["norm-skip-row.wgsl.jinja"],
+      "bias_f16": ["norm-skip-row.wgsl.jinja"],
+      "bias_output_only_f16": ["norm-skip-row.wgsl.jinja"],
+      "bias_output_only": ["norm-skip-row.wgsl.jinja"]
+    }
   }
 }

build/webgpu/norm-skip-row-vec4.wgsl.jinja CHANGED Viewed

@@ -42,8 +42,9 @@
     {{ svar }} = {{ svar }} / 2u;
 {% endif %}
   }
-{%- endmacro %}{% set useSubgroups = source.useSubgroups %}
-{% if source.usesF16 %}
 enable f16;
 {% endif %}
 {% if useSubgroups %}
@@ -51,9 +52,9 @@ enable subgroups;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
-const HIDDEN: u32 = {{ source.hidden }}u;
-const HIDDEN_V: u32 = {{ source.hiddenVec }}u;
-const WG: u32 = {{ source.wg }}u;
 var<workgroup> sg_partials: array<f32, WG>;
@@ -84,9 +85,9 @@ fn reduce_scalar(value: f32{% if useSubgroups %}, sg_lane: u32, sg_id: u32, num_
 // 4 contiguous residual elements (input[idx] + skip[skip_idx] [+ bias]) at vec4
 // index `vi`. skip_idx == idx for the normal (non-broadcast) path; for a skip
 // that broadcasts across the leading/batch dim uses a folded index.
-fn residual_value(idx: u32, skip_idx: u32{% if source.hasBias %}, vi: u32{% endif %}) -> vec4<f32> {
   var value = vec4<f32>(input[idx]) + vec4<f32>(skip[skip_idx]);
-{% if source.hasBias %}
   value = value + vec4<f32>(bias[vi]);
 {% endif %}
   return value;
@@ -106,12 +107,18 @@ fn main(
   }
   let tid = lid.x;
   let base = row * HIDDEN_V;
   let skip_base = base;
   var acc = 0.0;
   for (var i = tid; i < HIDDEN_V; i = i + WG) {
-    let v = residual_value(base + i, skip_base + i{% if source.hasBias %}, i{% endif %});
     acc = acc + dot(v, v);
   }
@@ -120,10 +127,10 @@ fn main(
   for (var i = tid; i < HIDDEN_V; i = i + WG) {
     let idx = base + i;
-    let residual = residual_value(idx, skip_base + i{% if source.hasBias %}, i{% endif %});
-{% if source.writeResidualSum %}
-    input_skip_bias_sum[idx] = {{ source.vecType }}(residual);
 {% endif %}
-    output[idx] = {{ source.vecType }}(residual * row_inv * vec4<f32>(gamma[i]));
   }
 }

     {{ svar }} = {{ svar }} / 2u;
 {% endif %}
   }
+{%- endmacro %}{% set broadcastSkip = broadcastSkip is defined and broadcastSkip %}
+{% set useSubgroups = useSubgroups %}
+{% if usesF16Spec %}
 enable f16;
 {% endif %}
 {% if useSubgroups %}
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
+const HIDDEN: u32 = {{ hidden }}u;
+const HIDDEN_V: u32 = {{ hiddenVec }}u;
+const WG: u32 = {{ wg }}u;
 var<workgroup> sg_partials: array<f32, WG>;
 // 4 contiguous residual elements (input[idx] + skip[skip_idx] [+ bias]) at vec4
 // index `vi`. skip_idx == idx for the normal (non-broadcast) path; for a skip
 // that broadcasts across the leading/batch dim uses a folded index.
+fn residual_value(idx: u32, skip_idx: u32{% if hasBias %}, vi: u32{% endif %}) -> vec4<f32> {
   var value = vec4<f32>(input[idx]) + vec4<f32>(skip[skip_idx]);
+{% if hasBias %}
   value = value + vec4<f32>(bias[vi]);
 {% endif %}
   return value;
   }
   let tid = lid.x;
   let base = row * HIDDEN_V;
+{% if broadcastSkip %}
+  // skip broadcasts across the batch dim: fold row into [0, skipRows) so every
+  // batch reuses the same skip row (skipRows == params.rows ⇒ identity).
+  let skip_base = (row % params.skipRows) * HIDDEN_V;
+{% else %}
   let skip_base = base;
+{% endif %}
   var acc = 0.0;
   for (var i = tid; i < HIDDEN_V; i = i + WG) {
+    let v = residual_value(base + i, skip_base + i{% if hasBias %}, i{% endif %});
     acc = acc + dot(v, v);
   }
   for (var i = tid; i < HIDDEN_V; i = i + WG) {
     let idx = base + i;
+    let residual = residual_value(idx, skip_base + i{% if hasBias %}, i{% endif %});
+{% if writeResidualSum %}
+    input_skip_bias_sum[idx] = {{ vecType }}(residual);
 {% endif %}
+    output[idx] = {{ vecType }}(residual * row_inv * vec4<f32>(gamma[i]));
   }
 }

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

@@ -1,65 +1,13 @@
-{% macro wgsl_tree_fold_stmt(a, op, idx, svar) %}
-{% if op == "max" %}
-{{ a }}[{{ idx }}] = max({{ a }}[{{ idx }}], {{ a }}[{{ idx }} + {{ svar }}]);
-{%- else %}
-{{ a }}[{{ idx }}] = {{ a }}[{{ idx }}] + {{ a }}[{{ idx }} + {{ svar }}];
-{%- endif %}
-{% endmacro %}
-{% macro wgsl_tree_fold(arrays, op="add", idx="lid", wg="WORKGROUP_SIZE", svar="stride", typed=false, form="tail", breakInline=false, bodyInline=false, barrierFirst=false) %}
-  var {{ svar }}{{ ": u32 " if typed else " " }}= {{ wg }} / 2u;
-  loop {
-{% if form == "head" %}
-{% if breakInline %}
-    if ({{ svar }} == 0u) { break; }
-{% else %}
-    if ({{ svar }} == 0u) {
-      break;
-    }
-{% endif %}
-{% endif %}
-{% if bodyInline %}
-    if ({{ idx }} < {{ svar }}) { {{ wgsl_tree_fold_stmt(arrays[0], op, idx, svar) }} }
-{% else %}
-    if ({{ idx }} < {{ svar }}) {
-{% for a in arrays %}
-      {{ wgsl_tree_fold_stmt(a, op, idx, svar) }}
-{% endfor %}
-    }
-{% endif %}
-{% if form == "head" %}
-{% if barrierFirst %}
-    workgroupBarrier();
-    {{ svar }} = {{ svar }} / 2u;
-{% else %}
-    {{ svar }} = {{ svar }} / 2u;
-    workgroupBarrier();
-{% endif %}
-{% else %}
-    workgroupBarrier();
-    if ({{ svar }} == 1u) {
-      break;
-    }
-    {{ svar }} = {{ svar }} / 2u;
-{% endif %}
-  }
-{%- endmacro %}
 /* One workgroup normalizes each row of residual = input + skip, with an
  * optional bias. */
-{% set degenerateRow = (not simplified) and hiddenSize == 1 %}
 {% if usesF16 %}
 enable f16;
 {% endif %}
-{% if useSubgroups and not degenerateRow %}
-enable subgroups;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
-{% if not degenerateRow or writeResidualSum %}
 const HIDDEN: u32 = {{ hiddenSize }}u;
-{% endif %}
 const WG: u32 = {{ workgroupSize }}u;
-{% if simplified %}
 var<workgroup> partial: array<f32, WG>;
 {% macro wgsl_tree_reduce_f32(name, mode, buffer="partial", wg="WG", trailingBarrier=true) %}
@@ -98,39 +46,7 @@ fn {{ name }}(value: f32, tid: u32) -> f32 {
 {{ wgsl_tree_reduce_f32("reduce_sum", "add", "partial", "WG") }}
 var<workgroup> row_inv: f32;
-{% else %}
-{% if not degenerateRow %}
-var<workgroup> pair_partial: array<vec2<f32>, WG>;
-{% if useSubgroups %}
-fn reduce_pair(value: vec2<f32>, sg_lane: u32, sg_id: u32, num_sg: u32) -> vec2<f32> {
-  let s = vec2<f32>(subgroupAdd(value.x), subgroupAdd(value.y));
-  if (num_sg == 1u) {
-    return s;
-  }
-  if (sg_lane == 0u) {
-    pair_partial[sg_id] = s;
-  }
-  workgroupBarrier();
-  var total = vec2<f32>(0.0, 0.0);
-  for (var i = 0u; i < num_sg; i = i + 1u) {
-    total = total + pair_partial[i];
-  }
-  return total;
-}
-{% else %}
-fn reduce_pair(value: vec2<f32>, tid: u32) -> vec2<f32> {
-  pair_partial[tid] = value;
-  workgroupBarrier();
-{{ wgsl_tree_fold(["pair_partial"], idx="tid", wg="WG", form="head") }}
-  return pair_partial[0];
-}
-{% endif %}
-{% endif %}
-{% endif %}
-{% if not degenerateRow or writeResidualSum %}
 fn residual_value(row: u32, d: u32) -> f32 {
   let index = row * HIDDEN + d;
   var value = f32(input[index]) + f32(skip[index]);
@@ -139,28 +55,19 @@ fn residual_value(row: u32, d: u32) -> f32 {
 {% endif %}
   return value;
 }
-{% endif %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(
   @builtin(workgroup_id) wg: vec3<u32>,
-  @builtin(num_workgroups) nwg: vec3<u32>{% if not degenerateRow %},
-  @builtin(local_invocation_id) lid: vec3<u32>{% endif %}{% if useSubgroups and not degenerateRow %},
-  @builtin(subgroup_invocation_id) sg_lane: u32,
-  @builtin(subgroup_id) sg_id: u32,
-  @builtin(num_subgroups) num_sg: u32{% endif %}
-) {
-  // 2D-folded row index: wg.y carries the high bits past the maxComputeWorkgroupsPerDimension
-  // workgroup-per-dimension dispatch limit. Reduces to wg.x when nwg.y == 1;
   // the row >= params.rows guard drops the over-dispatched tail.
-  let row = wg.x + wg.y * nwg.x;
   if (row >= params.rows) {
     return;
   }
-{% if not degenerateRow %}
   let tid = lid.x;
-{% endif %}
-{% if simplified %}
   // RMS normalization uses one sum-of-squares sweep, without a mean or beta.
@@ -183,47 +90,4 @@ fn main(
 {% endif %}
     output[index] = {{ scalar }}(residual * row_inv * f32(gamma[d]));
   }
-{% elif degenerateRow %}
-  // HIDDEN == 1: the row's mean is its only element, so the centered value and
-  // the variance are exactly zero and the output reduces to beta. The closed
-  // form avoids computing that zero by subtracting two equal rounded values.
-  let row_inv = inverseSqrt(params.epsilon);
-{% if writeResidualSum %}
-  let residual = residual_value(row, 0u);
-  input_skip_bias_sum[row] = {{ scalar }}(residual);
-{% endif %}
-  // 0.0 * row_inv keeps the IEEE result when epsilon == 0 makes row_inv +Inf.
-  output[row] = {{ scalar }}(0.0 * row_inv * f32(gamma[0]){% if hasBeta %} + f32(beta[0]){% endif %});
-{% else %}
-  // Shifted moments: accumulating (x - x[0], (x - x[0])^2) keeps the sums
-  // small for rows with a large common offset; every thread reconstructs the
-  // row mean and variance from the merged pair.
-  let shift = residual_value(row, 0u);
-  var acc = vec2<f32>(0.0, 0.0);
-  for (var d = tid; d < HIDDEN; d = d + WG) {
-    let centered = residual_value(row, d) - shift;
-    acc.x = acc.x + centered;
-    acc.y = acc.y + centered * centered;
-  }
-{% if useSubgroups %}
-  let totals = reduce_pair(acc, sg_lane, sg_id, num_sg);
-{% else %}
-  let totals = reduce_pair(acc, tid);
-{% endif %}
-  let mean_delta = totals.x / f32(HIDDEN);
-  let row_mean = shift + mean_delta;
-  let variance = max(totals.y / f32(HIDDEN) - mean_delta * mean_delta, 0.0);
-  let row_inv = inverseSqrt(variance + params.epsilon);
-  for (var d = tid; d < HIDDEN; d = d + WG) {
-    let index = row * HIDDEN + d;
-    let residual = residual_value(row, d);
-{% if writeResidualSum %}
-    input_skip_bias_sum[index] = {{ scalar }}(residual);
-{% endif %}
-    output[index] = {{ scalar }}((residual - row_mean) * row_inv * f32(gamma[d]){% if hasBeta %} + f32(beta[d]){% endif %});
-  }
-{% endif %}
 }

 /* One workgroup normalizes each row of residual = input + skip, with an
  * optional bias. */
 {% if usesF16 %}
 enable f16;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
 const HIDDEN: u32 = {{ hiddenSize }}u;
 const WG: u32 = {{ workgroupSize }}u;
 var<workgroup> partial: array<f32, WG>;
 {% macro wgsl_tree_reduce_f32(name, mode, buffer="partial", wg="WG", trailingBarrier=true) %}
 {{ wgsl_tree_reduce_f32("reduce_sum", "add", "partial", "WG") }}
 var<workgroup> row_inv: f32;
 fn residual_value(row: u32, d: u32) -> f32 {
   let index = row * HIDDEN + d;
   var value = f32(input[index]) + f32(skip[index]);
 {% endif %}
   return value;
 }
 @compute @workgroup_size(WG, 1, 1)
 fn main(
   @builtin(workgroup_id) wg: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>) {
+  // 2D-folded row index: wg.y carries the high bits past the per-axis dispatch fold width.
+  // Reduces to wg.x when the dispatch does not fold;
   // the row >= params.rows guard drops the over-dispatched tail.
+  let row = wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u;
   if (row >= params.rows) {
     return;
   }
   let tid = lid.x;
   // RMS normalization uses one sum-of-squares sweep, without a mean or beta.
 {% endif %}
     output[index] = {{ scalar }}(residual * row_inv * f32(gamma[d]));
   }
 }

build/webgpu/test.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "com.microsoft.SkipSimplifiedLayerNormalization",
   "cases": [
     {
       "name": "no_bias",
@@ -93,7 +92,7 @@
         "inputT": {
           "dtype": "float32",
           "shape": [2, 4],
-          "data": { "kind": "values", "values": [5.0, 5.0, 5.0, 5.0, -3.0, -3.0, -3.0, -3.0] }
         },
         "skipT": {
           "dtype": "float32",
@@ -133,7 +132,7 @@
         "skipT": {
           "dtype": "float32",
           "shape": [1, 4],
-          "data": { "kind": "values", "values": [-39999.0, -40000.0, -40001.0, -40002.0] }
         },
         "gammaT": { "dtype": "float32", "shape": [4], "data": { "kind": "values", "values": [1.0, 2.0, -1.0, 0.5] } }
       },
@@ -147,7 +146,7 @@
       "provenance": {
         "source": "onnxruntime/test/contrib_ops/skiplayernorm_op_test.cc",
         "test": "SkipLayerNormTest.SkipLayerNormNullInput",
-        "notes": "Same zero-token lowered shape as ORT's SkipLayerNormalization null-input case, applied to simplified RMS normalization."
       },
       "attrs": { "epsilon": 1e-12 },
       "inputs": {
@@ -275,7 +274,7 @@
       "provenance": {
         "source": "onnxruntime/test/contrib_ops/skiplayernorm_op_test.cc",
         "test": "SkipLayerNormTest.SkipSimplifiedLayerNormBatch1_Float16",
-        "notes": "ORT fp16 shape [1, 2, 4] is represented as float32 [2, 4] tokens by this lowered kernel. Epsilon is omitted to exercise the schema default of 1e-12."
       },
       "inputs": {
         "inputT": {
@@ -300,7 +299,7 @@
       "provenance": {
         "source": "onnxruntime/test/contrib_ops/skiplayernorm_op_test.cc",
         "test": "SkipLayerNormTest.SkipSimplifiedLayerNormBatch1_Bias_Float16",
-        "notes": "ORT fp16 shape [1, 1, 8] is represented as float32 [1, 8] tokens by this lowered kernel."
       },
       "attrs": { "epsilon": 0.00001 },
       "inputs": {
@@ -373,7 +372,7 @@
     {
       "name": "no_bias_hidden128_compact",
       "provenance": {
-        "notes": "Compact hidden-size-128 companion for the output-only bench family; keeps the same row width under reference-computed full outputs."
       },
       "attrs": { "epsilon": 0.00001 },
       "inputs": {
@@ -832,6 +831,31 @@
         }
       },
       "outputs": { "outputT": { "dtype": "float16", "shape": [2, 6], "tolerance": 0.01 } }
     }
   ]
 }

 {
   "cases": [
     {
       "name": "no_bias",
         "inputT": {
           "dtype": "float32",
           "shape": [2, 4],
+          "data": { "kind": "values", "values": [5.0, 5.0, 5.0, 5.0, -3.0, -2.0, -1.0, 0.0] }
         },
         "skipT": {
           "dtype": "float32",
         "skipT": {
           "dtype": "float32",
           "shape": [1, 4],
+          "data": { "kind": "values", "values": [-39999.0, -39999.5, -40000.0, -40000.5] }
         },
         "gammaT": { "dtype": "float32", "shape": [4], "data": { "kind": "values", "values": [1.0, 2.0, -1.0, 0.5] } }
       },
       "provenance": {
         "source": "onnxruntime/test/contrib_ops/skiplayernorm_op_test.cc",
         "test": "SkipLayerNormTest.SkipLayerNormNullInput",
+        "notes": "A zero-token ORT shape is flattened to an empty token-row input for simplified RMS normalization."
       },
       "attrs": { "epsilon": 1e-12 },
       "inputs": {
       "provenance": {
         "source": "onnxruntime/test/contrib_ops/skiplayernorm_op_test.cc",
         "test": "SkipLayerNormTest.SkipSimplifiedLayerNormBatch1_Float16",
+        "notes": "This package represents ORT's float16 shape [1, 2, 4] as two float32 token rows of width 4. Epsilon is omitted to exercise the schema default of 1e-12."
       },
       "inputs": {
         "inputT": {
       "provenance": {
         "source": "onnxruntime/test/contrib_ops/skiplayernorm_op_test.cc",
         "test": "SkipLayerNormTest.SkipSimplifiedLayerNormBatch1_Bias_Float16",
+        "notes": "This package represents ORT's float16 shape [1, 1, 8] as one float32 token row of width 8."
       },
       "attrs": { "epsilon": 0.00001 },
       "inputs": {
     {
       "name": "no_bias_hidden128_compact",
       "provenance": {
+        "notes": "A compact hidden-size-128 case exercises full outputs at the row width used by the output-only benchmark family."
       },
       "attrs": { "epsilon": 0.00001 },
       "inputs": {
         }
       },
       "outputs": { "outputT": { "dtype": "float16", "shape": [2, 6], "tolerance": 0.01 } }
+    },
+    {
+      "name": "zero_row_epsilon_zero",
+      "provenance": {
+        "notes": "An all-zero padded row with epsilon 0 computes `inverseSqrt(0) = +Inf` and then `0 * Inf = NaN`; the simplified RMS contract has no finite degenerate-row special case."
+      },
+      "attrs": { "epsilon": 0 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [1, 4, 64],
+          "data": { "kind": "cycle", "values": [0.5, -0.25, 0.75, 0.0, 0.0, 0.0, 0.0, 0.0] }
+        },
+        "skipT": { "dtype": "float32", "shape": [1, 4, 64], "data": { "kind": "constant", "value": 0.0 } },
+        "gammaT": { "dtype": "float32", "shape": [64], "data": { "kind": "constant", "value": 1.0 } }
+      },
+      "outputs": {
+        "outputT": {
+          "dtype": "float32",
+          "shape": [1, 4, 64],
+          "tolerance": 0.00002,
+          "relTolerance": 0.00002,
+          "allowNaN": true
+        }
+      }
     }
   ]
 }