sync 91d990483a17

Browse files

Files changed (7) hide show

README.md +21 -13
build/webgpu/bench.json +3 -4
build/webgpu/manifest.json +418 -861
build/webgpu/metadata.json +30 -8
build/webgpu/norm-skip-row-vec4.wgsl.jinja +16 -15
build/webgpu/norm-skip-row.wgsl.jinja +8 -70
build/webgpu/test.json +251 -13

README.md CHANGED Viewed

@@ -12,26 +12,26 @@ tags:
 ## Description
-Fuses skip addition with layer normalization. The rank-3 standard surface currently supports float32, required `beta`, no `bias` or residual output, hidden sizes divisible by four, and exact or documented broadcast skip shapes. The provider's rank-2 extension supports float32 output-only with optional `beta`, or `beta` with optional `bias` when emitting the residual; its float16 path requires `beta`, `bias`, a residual output, and four-wide hidden size. Other combinations, bfloat16, and training statistics are not implemented.
 See the [ONNX Runtime `SkipLayerNormalization` contrib-operator spec](https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.SkipLayerNormalization) for the reference semantics.
 ## Inputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `input` | `inputT` | `T` | — | — | Primary input normalized over the final hidden-size axis. Rank 3 is the public shape; rank 2 is an ONNX Runtime provider extension supported by this package. | required |
-| `skip` | `skipT` | `T` | — | — | Residual tensor. For rank-3 input it is exact shape, `(1, sequence_length, hidden_size)`, or `(sequence_length, hidden_size)`; rank-2 input requires exact shape. | required |
-| `gamma` | `gammaT` | `T` | `1` | — | Layer-norm scale weights of shape `(hidden_size)`. | required |
-| `beta` | `betaT` | `T` | `1` | — | Layer-norm bias weights of shape `(hidden_size)`. | optional |
-| `bias` | `biasT` | `T` | `1` | — | Optional additive bias of shape `(hidden_size)` added to `input + skip` before normalization. | 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 `bias` (when present) before normalization, with the same shape as `input`. | optional |
 ## Attributes
@@ -47,9 +47,13 @@ Default values (overridable per request):
 | --- | --- |
 | `T` | `float32`, `float16` |
 ## 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
@@ -58,10 +62,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
+Fuses skip addition with layer normalization for rank-2 or rank-3 input and a non-empty hidden axis. With exact-shape skip, float32 and float16 output-only paths support optional `beta`; adding `bias` requires `beta`. Returning the residual sum requires `beta`: float32 supports optional `bias` and arbitrary hidden sizes, while float16 requires `bias` and a hidden size divisible by four. Broadcast skip is supported for rank-3 float32 input, required `beta`, no `bias` or residual output, and a hidden size divisible by four. Bfloat16 and training statistics are not implemented.
 See the [ONNX Runtime `SkipLayerNormalization` contrib-operator spec](https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.SkipLayerNormalization) for the reference semantics.
 ## Inputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `inputT` | `input` | `T` | — | — | Primary input normalized over the final hidden-size axis. Rank 3 is the standard shape; rank 2 is a supported extension. | required |
+| `skipT` | `skip` | `T` | — | — | Residual tensor. For rank-3 input it is exact shape, `(1, sequence_length, hidden_size)`, or `(sequence_length, hidden_size)`; rank-2 input requires exact shape. | required |
+| `gammaT` | `gamma` | `T` | `1` | — | Layer-norm scale weights of shape `(hidden_size)`. | required |
+| `betaT` | `beta` | `T` | `1` | — | Layer-norm bias weights of shape `(hidden_size)`. | optional |
+| `biasT` | `bias` | `T` | `1` | — | Optional additive bias of shape `(hidden_size)` added to `input + skip` before normalization. | 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 `bias` (when present) before normalization, with the same shape as `input`. | optional |
 ## Attributes
 | --- | --- |
 | `T` | `float32`, `float16` |
+## Device requirements
+Some implementation variants require `shader-f16`. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.
 ## 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.SkipLayerNormalization",
   "cases": [
     {
       "name": "skip-layernorm-f32-256x128",
@@ -107,7 +106,7 @@
       "name": "skip-layernorm-f32-bias-output-only-4096x768",
       "preset": "model",
       "provenance": {
-        "notes": "Ordinary inference shape with bias present and all three optional outputs omitted. It uses six storage buffers and writes only the primary output, isolating the output-only path from the full-output sibling at the same geometry."
       },
       "vars": { "rows": 4096, "hidden": 768 },
       "attrs": { "epsilon": 0.00001 },
@@ -139,7 +138,7 @@
       "bench": { "metrics": [{ "type": "bandwidth", "value": "4 * (4 * args.rows * args.hidden + 2 * args.hidden)" }] }
     },
     {
-      "name": "skip-layernorm-f32-4096x768-hidden768-aligned-sibling",
       "preset": "model",
       "vars": { "rows": 4096, "hidden": 768 },
       "attrs": { "epsilon": 0.00001 },
@@ -173,7 +172,7 @@
       "bench": { "metrics": [{ "type": "bandwidth", "value": "4 * (4 * args.rows * args.hidden + 2 * args.hidden)" }] }
     },
     {
-      "name": "skip-layernorm-f32-rows4096-hidden4096-prefill-sibling",
       "preset": "model",
       "vars": { "rows": 4096, "hidden": 4096 },
       "attrs": { "epsilon": 0.00001 },

 {
   "cases": [
     {
       "name": "skip-layernorm-f32-256x128",
       "name": "skip-layernorm-f32-bias-output-only-4096x768",
       "preset": "model",
       "provenance": {
+        "notes": "An ordinary inference shape with bias omits all three optional outputs, uses six storage buffers, and writes only the primary output."
       },
       "vars": { "rows": 4096, "hidden": 768 },
       "attrs": { "epsilon": 0.00001 },
       "bench": { "metrics": [{ "type": "bandwidth", "value": "4 * (4 * args.rows * args.hidden + 2 * args.hidden)" }] }
     },
     {
+      "name": "skip-layernorm-f32-4096x768-hidden768-aligned",
       "preset": "model",
       "vars": { "rows": 4096, "hidden": 768 },
       "attrs": { "epsilon": 0.00001 },
       "bench": { "metrics": [{ "type": "bandwidth", "value": "4 * (4 * args.rows * args.hidden + 2 * args.hidden)" }] }
     },
     {
+      "name": "skip-layernorm-f32-rows4096-hidden4096-prefill",
       "preset": "model",
       "vars": { "rows": 4096, "hidden": 4096 },
       "attrs": { "epsilon": 0.00001 },

build/webgpu/manifest.json CHANGED Viewed

@@ -2,765 +2,170 @@
   "domain": "com.microsoft",
   "name": "SkipLayerNormalization",
   "sinceVersion": 1,
-  "description": "Fuses skip addition with layer normalization. The rank-3 standard surface currently supports float32, required `beta`, no `bias` or residual output, hidden sizes divisible by four, and exact or documented broadcast skip shapes. The provider's rank-2 extension supports float32 output-only with optional `beta`, or `beta` with optional `bias` when emitting the residual; its float16 path requires `beta`, `bias`, a residual output, and four-wide hidden size. Other combinations, bfloat16, and training statistics are not implemented.",
-  "inputs": [
-    {
-      "role": "input",
-      "dtype": "T",
-      "description": "Primary input normalized over the final hidden-size axis. Rank 3 is the public shape; rank 2 is an ONNX Runtime provider extension supported by this package."
-    },
-    {
-      "role": "skip",
-      "dtype": "T",
-      "description": "Residual tensor. For rank-3 input it is exact shape, `(1, sequence_length, hidden_size)`, or `(sequence_length, hidden_size)`; rank-2 input requires exact shape."
-    },
-    { "role": "gamma", "dtype": "T", "rank": 1, "description": "Layer-norm scale weights of shape `(hidden_size)`." },
-    {
-      "role": "beta",
-      "dtype": "T",
-      "rank": 1,
-      "optional": true,
-      "description": "Layer-norm bias weights of shape `(hidden_size)`."
-    },
-    {
-      "role": "bias",
-      "dtype": "T",
-      "rank": 1,
-      "optional": true,
-      "description": "Optional additive bias of shape `(hidden_size)` added to `input + skip` before normalization."
-    }
-  ],
-  "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",
       "rank": "ranks.inputT",
       "optional": true,
-      "shape": "shapes.inputT",
-      "description": "Sum of `input`, `skip`, and `bias` (when present) before normalization, with the same shape as `input`."
     }
-  ],
-  "attributes": { "epsilon": 9.999999960041972e-13 },
-  "attributeDescriptions": { "epsilon": "Non-negative epsilon added to the variance 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" },
-    "betaT": { "kind": "tensor", "semantic": "beta", "role": "input", "required": false },
-    "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": "dim(shapes.inputT, 0)",
-    "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))))",
     "portableWideExecution": "not has(device.adapterInfo, \"subgroupMinSize\") or device.adapterInfo.subgroupMinSize >= 32",
     "broadcastRows": "dim(shapes.inputT, 0) * dim(shapes.inputT, 1)",
     "broadcastHiddenSize": "dim(shapes.inputT, 2)",
     "broadcastSkipWgVec4": "max(1, min(tunables.MAX_WORKGROUP_SIZE, device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX, pow2ceil(ceilDiv(broadcastHiddenSize, 4))))",
-    "rowDispatchFits": "rowCount <= device.limits.maxComputeWorkgroupsPerDimension * device.limits.maxComputeWorkgroupsPerDimension",
-    "broadcastDispatchFits": "broadcastRows <= device.limits.maxComputeWorkgroupsPerDimension * device.limits.maxComputeWorkgroupsPerDimension",
     "normResourcesFit": "skipWg * 8 <= device.limits.maxComputeWorkgroupStorageSize and skipWgVec4 * 8 <= device.limits.maxComputeWorkgroupStorageSize",
     "broadcastResourcesFit": "broadcastSkipWgVec4 * 8 <= device.limits.maxComputeWorkgroupStorageSize",
     "epsilonOk": "attrs.epsilon >= 0",
-    "rank2CoreContract": "epsilonOk and ranks.inputT == 2 and ranks.skipT == 2 and ranks.gammaT == 1 and ranks.outputT == 2 and dim(shapes.inputT, 0) == dim(shapes.skipT, 0) and dim(shapes.inputT, 1) == dim(shapes.skipT, 1) and dim(shapes.inputT, 1) > 0 and dim(shapes.gammaT, 0) == dim(shapes.inputT, 1) and dim(shapes.outputT, 0) == dim(shapes.inputT, 0) and dim(shapes.outputT, 1) == dim(shapes.inputT, 1)",
-    "residualOutputContract": "present.residualT and ranks.residualT == 2 and dim(shapes.residualT, 0) == dim(shapes.inputT, 0) and dim(shapes.residualT, 1) == dim(shapes.inputT, 1)",
     "outputOnlyContract": "not present.residualT",
-    "betaContract": "false if not present.betaT else (ranks.betaT == 1 and dim(shapes.betaT, 0) == dim(shapes.inputT, 1))",
     "noBetaContract": "not present.betaT",
     "f32MainDtypes": "tensorDtypes.inputT == \"float32\" and tensorDtypes.skipT == \"float32\" and tensorDtypes.gammaT == \"float32\" and tensorDtypes.outputT == \"float32\"",
     "f16MainDtypes": "tensorDtypes.inputT == \"float16\" and tensorDtypes.skipT == \"float16\" and tensorDtypes.gammaT == \"float16\" and tensorDtypes.outputT == \"float16\"",
     "f32ResidualDtypes": "f32MainDtypes and tensorDtypes.residualT == \"float32\" if present.residualT else false",
     "f16ResidualDtypes": "f16MainDtypes and tensorDtypes.residualT == \"float16\" if present.residualT else false",
-    "vec4Aligned": "dim(shapes.inputT, 1) % 4 == 0",
     "broadcastSkipShapeOk": "(ranks.skipT == 2 and dim(shapes.skipT, 0) == dim(shapes.inputT, 1) and dim(shapes.skipT, 1) == dim(shapes.inputT, 2)) or (ranks.skipT == 3 and ((dim(shapes.skipT, 0) == 1 and dim(shapes.skipT, 1) == dim(shapes.inputT, 1) and dim(shapes.skipT, 2) == dim(shapes.inputT, 2)) or sameShape(shapes.skipT, shapes.inputT)))",
     "broadcastOutputOnlyContract": "false if ranks.inputT != 3 or not present.betaT else (epsilonOk and not present.biasT and not present.residualT and dim(shapes.inputT, 2) % 4 == 0 and broadcastSkipShapeOk and ranks.gammaT == 1 and ranks.betaT == 1 and ranks.outputT == 3 and tensorDtypes.inputT == \"float32\" and tensorDtypes.skipT == \"float32\" and tensorDtypes.gammaT == \"float32\" and tensorDtypes.betaT == \"float32\" and tensorDtypes.outputT == \"float32\" and dim(shapes.inputT, 2) > 0 and dim(shapes.gammaT, 0) == dim(shapes.inputT, 2) and dim(shapes.betaT, 0) == dim(shapes.inputT, 2) and sameShape(shapes.outputT, shapes.inputT))",
     "hasSubgroups": "device.features.has(\"subgroups\") and device.wgslLanguageFeatures.has(\"subgroup_id\")",
     "hasF16": "device.features.has(\"shader-f16\")",
-    "f32_beta_no_bias_residual_contract": "rank2CoreContract and residualOutputContract and betaContract and f32ResidualDtypes and not present.biasT and tensorDtypes.betaT == \"float32\"",
-    "f32_beta_bias_residual_contract": "false if not present.biasT else (rank2CoreContract and residualOutputContract and betaContract and f32ResidualDtypes and ranks.biasT == 1 and tensorDtypes.betaT == \"float32\" and tensorDtypes.biasT == \"float32\" and dim(shapes.biasT, 0) == hiddenSize)",
-    "f16_beta_bias_residual_contract": "false if not present.biasT else (hasF16 and rank2CoreContract and residualOutputContract and betaContract and f16ResidualDtypes and ranks.biasT == 1 and tensorDtypes.betaT == \"float16\" and tensorDtypes.biasT == \"float16\" and dim(shapes.biasT, 0) == hiddenSize)",
-    "f32_no_beta_output_contract": "rank2CoreContract and outputOnlyContract and noBetaContract and f32MainDtypes and not present.biasT",
-    "f32_beta_no_bias_output_only_contract": "rank2CoreContract and outputOnlyContract and betaContract and f32MainDtypes and not present.biasT and tensorDtypes.betaT == \"float32\"",
-    "f32_beta_bias_output_only_contract": "false if not present.biasT else (rank2CoreContract and outputOnlyContract and betaContract and f32MainDtypes and ranks.biasT == 1 and tensorDtypes.betaT == \"float32\" and tensorDtypes.biasT == \"float32\" and dim(shapes.biasT, 0) == hiddenSize)"
   },
-  "bindingSets": {
-    "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": "beta",
-        "arg": "betaT",
-        "semantic": "beta",
-        "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_output_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": "bias",
-        "arg": "biasT",
-        "semantic": "bias",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar",
-        "length": "$HIDDEN_LEN"
-      },
-      {
-        "name": "gamma",
-        "arg": "gammaT",
-        "semantic": "gamma",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar",
-        "length": "$HIDDEN_LEN"
-      },
-      {
-        "name": "beta",
-        "arg": "betaT",
-        "semantic": "beta",
-        "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" }
-          ]
-        }
-      }
-    ],
-    "vec4_bias_output_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": "bias",
-        "arg": "biasT",
-        "semantic": "bias",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar",
-        "length": "$HIDDEN_LEN"
-      },
-      {
-        "name": "gamma",
-        "arg": "gammaT",
-        "semantic": "gamma",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar",
-        "length": "$HIDDEN_LEN"
-      },
-      {
-        "name": "beta",
-        "arg": "betaT",
-        "semantic": "beta",
-        "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" }
-          ]
-        }
-      }
-    ],
-    "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": "beta",
-        "arg": "betaT",
-        "semantic": "beta",
-        "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": "bias",
-        "arg": "biasT",
-        "semantic": "bias",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar",
-        "length": "$HIDDEN_LEN"
-      },
-      {
-        "name": "gamma",
-        "arg": "gammaT",
-        "semantic": "gamma",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar",
-        "length": "$HIDDEN_LEN"
-      },
-      {
-        "name": "beta",
-        "arg": "betaT",
-        "semantic": "beta",
-        "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_no_beta_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" }
-          ]
-        }
-      }
-    ],
-    "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": "beta",
-        "arg": "betaT",
-        "semantic": "beta",
-        "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": "beta",
-        "arg": "betaT",
-        "semantic": "beta",
-        "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": "beta",
-        "arg": "betaT",
-        "semantic": "beta",
-        "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" }
-          ]
-        }
-      }
-    ],
-    "vec4_no_beta_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_beta_broadcast_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": "beta",
-        "arg": "betaT",
-        "semantic": "beta",
-        "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": "dim(shapes.inputT, 0) * dim(shapes.inputT, 1)" },
-            {
-              "name": "rowStride",
-              "type": "u32",
-              "value": "max(1, min(broadcastRows, device.limits.maxComputeWorkgroupsPerDimension))"
-            },
-            { "name": "epsilon", "type": "f32", "value": "attrs.epsilon" },
-            { "name": "skipRows", "type": "u32", "value": "numel(shapes.skipT) / broadcastHiddenSize" }
-          ]
-        }
-      }
-    ]
   },
   "variants": [
     {
       "id": "beta_output_only_vec4_broadcast",
       "priority": 19,
       "when": ["broadcastOutputOnlyContract", "broadcastResourcesFit", "broadcastDispatchFits"],
-      "constants": { "scalar": "\"f32\"", "vectorScalar": "\"vec4<f32>\"", "HIDDEN_LEN": "broadcastHiddenSize / 4" },
       "passes": [
         {
           "id": "main",
           "name": "SkipLayerNormalization.BroadcastSkip",
-          "source": {
-            "shader": "norm-skip-row-vec4.wgsl.jinja",
-            "inputs": {
-              "simplified": false,
-              "hasBias": false,
-              "hasBeta": true,
-              "writeResidualSum": false,
-              "usesF16": false,
-              "broadcastSkip": true,
-              "hidden": "broadcastHiddenSize",
-              "hiddenVec": "broadcastHiddenSize / 4",
-              "wg": "broadcastSkipWgVec4",
-              "vecType": "\"vec4<f32>\"",
-              "useSubgroups": "hasSubgroups"
-            }
           },
-          "subgroupCollectivesWidth": "portable",
-          "bindings": "vec4_beta_broadcast_output_only",
-          "dispatch": { "workgroups": "broadcastRows" }
         }
       ]
     },
@@ -768,11 +173,10 @@
       "id": "beta_bias_vec4",
       "priority": 15,
       "when": ["f32_beta_bias_residual_contract", "vec4Aligned", "hasSubgroups or \"bias\" == \"bias\" or portableWideExecution", "normResourcesFit", "rowDispatchFits"],
-      "constants": {
         "scalar": "\"f32\"",
         "vectorScalar": "\"vec4<f32>\"",
         "hasBias": "\"bias\" == \"bias\"",
-        "hiddenSize": "hiddenSize",
         "workgroupSize": "skipWg",
         "HIDDEN_LEN": "hiddenSize / 4"
       },
@@ -780,24 +184,22 @@
         {
           "id": "normalize",
           "name": "SkipLayerNormalization.Vec4.Normalize",
-          "source": {
-            "shader": "norm-skip-row-vec4.wgsl.jinja",
-            "inputs": {
-              "simplified": false,
-              "hasBias": "\"bias\" == \"bias\"",
-              "hasBeta": true,
-              "writeResidualSum": true,
-              "usesF16": false,
-              "hidden": "hiddenSize",
-              "hiddenVec": "hiddenSize / 4",
-              "wg": "skipWgVec4",
-              "vecType": "\"vec4<f32>\"",
-              "useSubgroups": "hasSubgroups"
-            }
           },
-          "subgroupCollectivesWidth": "portable",
-          "bindings": "vec4_bias_output_residual",
-          "dispatch": { "workgroups": "rowCount" }
         }
       ]
     },
@@ -805,14 +207,13 @@
       "id": "beta_bias_row",
       "priority": 5,
       "when": ["f32_beta_bias_residual_contract", "normResourcesFit", "rowDispatchFits"],
-      "constants": {
         "simplified": false,
         "useSubgroups": "hasSubgroups",
         "hasBeta": true,
         "writeResidualSum": true,
         "hasBias": "\"bias\" == \"bias\"",
         "scalar": "\"f32\"",
-        "hiddenSize": "hiddenSize",
         "workgroupSize": "skipWg",
         "HIDDEN_LEN": "hiddenSize"
       },
@@ -821,10 +222,10 @@
           "id": "normalize",
           "name": "SkipLayerNormalization.Row.Normalize",
           "shader": "norm-skip-row.wgsl.jinja",
-          "subgroupCollectivesWidth": "portable",
-          "bindings": "scalar_bias_output_residual",
-          "dispatch": { "workgroups": "rowCount" },
-          "constants": { "writeResidualSum": true }
         }
       ]
     },
@@ -832,29 +233,27 @@
       "id": "beta_bias_vec4_f16",
       "priority": 21,
       "when": ["f16_beta_bias_residual_contract", "vec4Aligned", "normResourcesFit", "rowDispatchFits"],
-      "constants": { "scalar": "\"f16\"", "vectorScalar": "\"vec4<f16>\"", "HIDDEN_LEN": "hiddenSize / 4" },
       "passes": [
         {
           "id": "main",
           "name": "SkipLayerNormalization.Vec4",
-          "source": {
-            "shader": "norm-skip-row-vec4.wgsl.jinja",
-            "inputs": {
-              "simplified": false,
-              "hasBias": true,
-              "hasBeta": true,
-              "writeResidualSum": true,
-              "usesF16": true,
-              "hidden": "hiddenSize",
-              "hiddenVec": "hiddenSize / 4",
-              "wg": "skipWgVec4",
-              "vecType": "\"vec4<f16>\"",
-              "useSubgroups": "hasSubgroups"
-            }
           },
-          "subgroupCollectivesWidth": "portable",
-          "bindings": "vec4_bias_output_residual",
-          "dispatch": { "workgroups": "rowCount" }
         }
       ]
     },
@@ -862,29 +261,27 @@
       "id": "no_beta_output_only_vec4",
       "priority": 20,
       "when": ["f32_no_beta_output_contract", "vec4Aligned", "normResourcesFit", "rowDispatchFits"],
-      "constants": { "scalar": "\"f32\"", "vectorScalar": "\"vec4<f32>\"", "HIDDEN_LEN": "hiddenSize / 4" },
       "passes": [
         {
           "id": "main",
           "name": "SkipLayerNormalization.NoBetaOutputOnly.Vec4",
-          "source": {
-            "shader": "norm-skip-row-vec4.wgsl.jinja",
-            "inputs": {
-              "simplified": false,
-              "hasBias": false,
-              "hasBeta": false,
-              "writeResidualSum": false,
-              "usesF16": false,
-              "hidden": "hiddenSize",
-              "hiddenVec": "hiddenSize / 4",
-              "wg": "skipWgVec4",
-              "vecType": "\"vec4<f32>\"",
-              "useSubgroups": "hasSubgroups"
-            }
           },
-          "subgroupCollectivesWidth": "portable",
-          "bindings": "vec4_no_beta_output_only",
-          "dispatch": { "workgroups": "rowCount" }
         }
       ]
     },
@@ -892,14 +289,13 @@
       "id": "no_beta_output_only_row",
       "priority": 10,
       "when": ["f32_no_beta_output_contract", "normResourcesFit", "rowDispatchFits"],
-      "constants": {
         "simplified": false,
         "hasBias": false,
         "hasBeta": false,
         "writeResidualSum": false,
         "useSubgroups": "hasSubgroups",
         "scalar": "\"f32\"",
-        "hiddenSize": "hiddenSize",
         "workgroupSize": "skipWg",
         "HIDDEN_LEN": "hiddenSize"
       },
@@ -908,9 +304,63 @@
           "id": "main",
           "name": "SkipLayerNormalization.NoBetaOutputOnly.Row",
           "shader": "norm-skip-row.wgsl.jinja",
-          "subgroupCollectivesWidth": "portable",
-          "bindings": "scalar_no_beta_output_only",
-          "dispatch": { "workgroups": "rowCount" }
         }
       ]
     },
@@ -918,11 +368,10 @@
       "id": "beta_no_bias_vec4",
       "priority": 20,
       "when": ["f32_beta_no_bias_residual_contract", "vec4Aligned", "hasSubgroups or portableWideExecution", "normResourcesFit", "rowDispatchFits"],
-      "constants": {
         "scalar": "\"f32\"",
         "vectorScalar": "\"vec4<f32>\"",
         "hasBias": "\"no_bias\" == \"bias\"",
-        "hiddenSize": "hiddenSize",
         "workgroupSize": "skipWg",
         "HIDDEN_LEN": "hiddenSize / 4"
       },
@@ -930,24 +379,22 @@
         {
           "id": "main",
           "name": "SkipLayerNormalization.Vec4",
-          "source": {
-            "shader": "norm-skip-row-vec4.wgsl.jinja",
-            "inputs": {
-              "simplified": false,
-              "hasBias": "\"no_bias\" == \"bias\"",
-              "hasBeta": true,
-              "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" }
         }
       ]
     },
@@ -955,14 +402,13 @@
       "id": "beta_no_bias_row",
       "priority": 10,
       "when": ["f32_beta_no_bias_residual_contract", "normResourcesFit", "rowDispatchFits"],
-      "constants": {
         "simplified": false,
         "useSubgroups": "hasSubgroups",
         "hasBeta": true,
         "writeResidualSum": true,
         "hasBias": "\"no_bias\" == \"bias\"",
         "scalar": "\"f32\"",
-        "hiddenSize": "hiddenSize",
         "workgroupSize": "skipWg",
         "HIDDEN_LEN": "hiddenSize"
       },
@@ -971,9 +417,9 @@
           "id": "main",
           "name": "SkipLayerNormalization.Row",
           "shader": "norm-skip-row.wgsl.jinja",
-          "subgroupCollectivesWidth": "portable",
-          "bindings": "scalar_no_bias_residual",
-          "dispatch": { "workgroups": "rowCount" }
         }
       ]
     },
@@ -981,11 +427,10 @@
       "id": "beta_no_bias_output_only_vec4",
       "priority": 20,
       "when": ["f32_beta_no_bias_output_only_contract", "vec4Aligned", "hasSubgroups or \"no_bias\" == \"bias\" or portableWideExecution", "normResourcesFit", "rowDispatchFits"],
-      "constants": {
         "scalar": "\"f32\"",
         "vectorScalar": "\"vec4<f32>\"",
         "hasBias": "\"no_bias\" == \"bias\"",
-        "hiddenSize": "hiddenSize",
         "workgroupSize": "skipWg",
         "HIDDEN_LEN": "hiddenSize / 4"
       },
@@ -993,24 +438,22 @@
         {
           "id": "main",
           "name": "SkipLayerNormalization.Vec4",
-          "source": {
-            "shader": "norm-skip-row-vec4.wgsl.jinja",
-            "inputs": {
-              "simplified": false,
-              "hasBias": "\"no_bias\" == \"bias\"",
-              "hasBeta": true,
-              "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" }
         }
       ]
     },
@@ -1018,14 +461,13 @@
       "id": "beta_no_bias_output_only_row",
       "priority": 10,
       "when": ["f32_beta_no_bias_output_only_contract", "normResourcesFit", "rowDispatchFits"],
-      "constants": {
         "simplified": false,
         "useSubgroups": "hasSubgroups",
         "hasBeta": true,
         "writeResidualSum": false,
         "hasBias": "\"no_bias\" == \"bias\"",
         "scalar": "\"f32\"",
-        "hiddenSize": "hiddenSize",
         "workgroupSize": "skipWg",
         "HIDDEN_LEN": "hiddenSize"
       },
@@ -1034,9 +476,9 @@
           "id": "main",
           "name": "SkipLayerNormalization.Row",
           "shader": "norm-skip-row.wgsl.jinja",
-          "subgroupCollectivesWidth": "portable",
-          "bindings": "scalar_no_bias_output_only",
-          "dispatch": { "workgroups": "rowCount" }
         }
       ]
     },
@@ -1044,11 +486,10 @@
       "id": "beta_bias_output_only_vec4",
       "priority": 20,
       "when": ["f32_beta_bias_output_only_contract", "vec4Aligned", "hasSubgroups or \"bias\" == \"bias\" or portableWideExecution", "normResourcesFit", "rowDispatchFits"],
-      "constants": {
         "scalar": "\"f32\"",
         "vectorScalar": "\"vec4<f32>\"",
         "hasBias": "\"bias\" == \"bias\"",
-        "hiddenSize": "hiddenSize",
         "workgroupSize": "skipWg",
         "HIDDEN_LEN": "hiddenSize / 4"
       },
@@ -1056,24 +497,22 @@
         {
           "id": "main",
           "name": "SkipLayerNormalization.Vec4",
-          "source": {
-            "shader": "norm-skip-row-vec4.wgsl.jinja",
-            "inputs": {
-              "simplified": false,
-              "hasBias": "\"bias\" == \"bias\"",
-              "hasBeta": true,
-              "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" }
         }
       ]
     },
@@ -1081,14 +520,13 @@
       "id": "beta_bias_output_only_row",
       "priority": 10,
       "when": ["f32_beta_bias_output_only_contract", "normResourcesFit", "rowDispatchFits"],
-      "constants": {
         "simplified": false,
         "useSubgroups": "hasSubgroups",
         "hasBeta": true,
         "writeResidualSum": false,
         "hasBias": "\"bias\" == \"bias\"",
         "scalar": "\"f32\"",
-        "hiddenSize": "hiddenSize",
         "workgroupSize": "skipWg",
         "HIDDEN_LEN": "hiddenSize"
       },
@@ -1097,12 +535,131 @@
           "id": "main",
           "name": "SkipLayerNormalization.Row",
           "shader": "norm-skip-row.wgsl.jinja",
-          "subgroupCollectivesWidth": "portable",
-          "bindings": "scalar_bias_output_only",
-          "dispatch": { "workgroups": "rowCount" }
         }
       ]
     }
-  ],
-  "tunables": { "MAX_WORKGROUP_SIZE": 256 }
 }

   "domain": "com.microsoft",
   "name": "SkipLayerNormalization",
   "sinceVersion": 1,
+  "inputs": {
+    "inputT": { "onnx": "input", "dtype": "T" },
+    "skipT": { "onnx": "skip", "dtype": "T" },
+    "gammaT": { "onnx": "gamma", "dtype": "T", "rank": 1 },
+    "betaT": { "onnx": "beta", "dtype": "T", "rank": 1, "optional": true },
+    "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))))",
     "portableWideExecution": "not has(device.adapterInfo, \"subgroupMinSize\") or device.adapterInfo.subgroupMinSize >= 32",
     "broadcastRows": "dim(shapes.inputT, 0) * dim(shapes.inputT, 1)",
     "broadcastHiddenSize": "dim(shapes.inputT, 2)",
     "broadcastSkipWgVec4": "max(1, min(tunables.MAX_WORKGROUP_SIZE, device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX, pow2ceil(ceilDiv(broadcastHiddenSize, 4))))",
+    "rowDispatchFits": "rowCount <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
+    "broadcastDispatchFits": "broadcastRows <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
     "normResourcesFit": "skipWg * 8 <= device.limits.maxComputeWorkgroupStorageSize and skipWgVec4 * 8 <= device.limits.maxComputeWorkgroupStorageSize",
     "broadcastResourcesFit": "broadcastSkipWgVec4 * 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)",
+    "residualOutputContract": "present.residualT and sameShape(shapes.residualT, shapes.inputT)",
     "outputOnlyContract": "not present.residualT",
+    "betaContract": "false if not present.betaT else (ranks.betaT == 1 and dim(shapes.betaT, 0) == dim(shapes.inputT, -1))",
     "noBetaContract": "not present.betaT",
     "f32MainDtypes": "tensorDtypes.inputT == \"float32\" and tensorDtypes.skipT == \"float32\" and tensorDtypes.gammaT == \"float32\" and tensorDtypes.outputT == \"float32\"",
     "f16MainDtypes": "tensorDtypes.inputT == \"float16\" and tensorDtypes.skipT == \"float16\" and tensorDtypes.gammaT == \"float16\" and tensorDtypes.outputT == \"float16\"",
     "f32ResidualDtypes": "f32MainDtypes and tensorDtypes.residualT == \"float32\" if present.residualT else false",
     "f16ResidualDtypes": "f16MainDtypes and tensorDtypes.residualT == \"float16\" if present.residualT else false",
+    "vec4Aligned": "dim(shapes.inputT, -1) % 4 == 0",
     "broadcastSkipShapeOk": "(ranks.skipT == 2 and dim(shapes.skipT, 0) == dim(shapes.inputT, 1) and dim(shapes.skipT, 1) == dim(shapes.inputT, 2)) or (ranks.skipT == 3 and ((dim(shapes.skipT, 0) == 1 and dim(shapes.skipT, 1) == dim(shapes.inputT, 1) and dim(shapes.skipT, 2) == dim(shapes.inputT, 2)) or sameShape(shapes.skipT, shapes.inputT)))",
     "broadcastOutputOnlyContract": "false if ranks.inputT != 3 or not present.betaT else (epsilonOk and not present.biasT and not present.residualT and dim(shapes.inputT, 2) % 4 == 0 and broadcastSkipShapeOk and ranks.gammaT == 1 and ranks.betaT == 1 and ranks.outputT == 3 and tensorDtypes.inputT == \"float32\" and tensorDtypes.skipT == \"float32\" and tensorDtypes.gammaT == \"float32\" and tensorDtypes.betaT == \"float32\" and tensorDtypes.outputT == \"float32\" and dim(shapes.inputT, 2) > 0 and dim(shapes.gammaT, 0) == dim(shapes.inputT, 2) and dim(shapes.betaT, 0) == dim(shapes.inputT, 2) and sameShape(shapes.outputT, shapes.inputT))",
     "hasSubgroups": "device.features.has(\"subgroups\") and device.wgslLanguageFeatures.has(\"subgroup_id\")",
     "hasF16": "device.features.has(\"shader-f16\")",
+    "f32_beta_no_bias_residual_contract": "coreContract and residualOutputContract and betaContract and f32ResidualDtypes and not present.biasT and tensorDtypes.betaT == \"float32\"",
+    "f32_beta_bias_residual_contract": "false if not present.biasT else (coreContract and residualOutputContract and betaContract and f32ResidualDtypes and ranks.biasT == 1 and tensorDtypes.betaT == \"float32\" and tensorDtypes.biasT == \"float32\" and dim(shapes.biasT, 0) == hiddenSize)",
+    "f16_beta_bias_residual_contract": "false if not present.biasT else (hasF16 and coreContract and residualOutputContract and betaContract and f16ResidualDtypes and ranks.biasT == 1 and tensorDtypes.betaT == \"float16\" and tensorDtypes.biasT == \"float16\" and dim(shapes.biasT, 0) == hiddenSize)",
+    "f32_no_beta_output_contract": "coreContract and outputOnlyContract and noBetaContract and f32MainDtypes and not present.biasT",
+    "f32_beta_no_bias_output_only_contract": "coreContract and outputOnlyContract and betaContract and f32MainDtypes and not present.biasT and tensorDtypes.betaT == \"float32\"",
+    "f32_beta_bias_output_only_contract": "false if not present.biasT else (coreContract and outputOnlyContract and betaContract and f32MainDtypes and ranks.biasT == 1 and tensorDtypes.betaT == \"float32\" and tensorDtypes.biasT == \"float32\" and dim(shapes.biasT, 0) == hiddenSize)",
+    "f16_no_beta_output_contract": "hasF16 and coreContract and outputOnlyContract and noBetaContract and f16MainDtypes and not present.biasT",
+    "f16_beta_no_bias_output_only_contract": "hasF16 and coreContract and outputOnlyContract and betaContract and f16MainDtypes and not present.biasT and tensorDtypes.betaT == \"float16\"",
+    "f16_beta_bias_output_only_contract": "false if not present.biasT else (hasF16 and coreContract and outputOnlyContract and betaContract and f16MainDtypes and ranks.biasT == 1 and tensorDtypes.betaT == \"float16\" and tensorDtypes.biasT == \"float16\" and dim(shapes.biasT, 0) == hiddenSize)"
   },
+  "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" },
+    "beta": { "arg": "betaT", "buffer": "read-only-storage", "elementType": "$vectorScalar", "length": "$HIDDEN_LEN" },
+    "output": { "arg": "outputT", "buffer": "storage", "elementType": "$vectorScalar" },
+    "bias": { "arg": "biasT", "buffer": "read-only-storage", "elementType": "$vectorScalar", "length": "$HIDDEN_LEN" },
+    "input_skip_bias_sum": { "arg": "residualT", "buffer": "storage", "elementType": "$vectorScalar" },
+    "params_2": {
+      "name": "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" }
+      ]
+    },
+    "input_2": { "arg": "inputT", "name": "input", "buffer": "read-only-storage", "elementType": "$scalar" },
+    "skip_2": { "arg": "skipT", "name": "skip", "buffer": "read-only-storage", "elementType": "$scalar" },
+    "bias_2": {
+      "arg": "biasT",
+      "name": "bias",
+      "buffer": "read-only-storage",
+      "elementType": "$scalar",
+      "length": "$HIDDEN_LEN"
+    },
+    "gamma_2": {
+      "arg": "gammaT",
+      "name": "gamma",
+      "buffer": "read-only-storage",
+      "elementType": "$scalar",
+      "length": "$HIDDEN_LEN"
+    },
+    "beta_2": {
+      "arg": "betaT",
+      "name": "beta",
+      "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_3": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "rowCount" },
+        { "name": "epsilon", "type": "f32", "value": "attrs.epsilon" }
+      ]
+    }
   },
   "variants": [
     {
       "id": "beta_output_only_vec4_broadcast",
       "priority": 19,
       "when": ["broadcastOutputOnlyContract", "broadcastResourcesFit", "broadcastDispatchFits"],
+      "derive": { "scalar": "\"f32\"", "vectorScalar": "\"vec4<f32>\"", "HIDDEN_LEN": "broadcastHiddenSize / 4" },
       "passes": [
         {
           "id": "main",
           "name": "SkipLayerNormalization.BroadcastSkip",
+          "shader": "norm-skip-row-vec4.wgsl.jinja",
+          "derive": {
+            "simplified": false,
+            "hasBias": false,
+            "hasBeta": true,
+            "writeResidualSum": false,
+            "usesF16Spec": false,
+            "broadcastSkip": true,
+            "hidden": "broadcastHiddenSize",
+            "hiddenVec": "broadcastHiddenSize / 4",
+            "wg": "broadcastSkipWgVec4",
+            "vecType": "\"vec4<f32>\"",
+            "useSubgroups": "hasSubgroups"
           },
+          "bindings": [
+            "input",
+            "skip",
+            "gamma",
+            "beta",
+            "output",
+            {
+              "name": "params",
+              "struct": [
+                { "name": "rows", "type": "u32", "value": "dim(shapes.inputT, 0) * dim(shapes.inputT, 1)" },
+                {
+                  "name": "rowStride",
+                  "type": "u32",
+                  "value": "max(1, min(broadcastRows, min(device.limits.maxComputeWorkgroupsPerDimension, 65535)))"
+                },
+                { "name": "epsilon", "type": "f32", "value": "attrs.epsilon" },
+                { "name": "skipRows", "type": "u32", "value": "numel(shapes.skipT) / broadcastHiddenSize" }
+              ]
+            }
+          ],
+          "dispatch": { "x": "min(broadcastRows, 65535)", "y": "ceilDiv(broadcastRows, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
       "id": "beta_bias_vec4",
       "priority": 15,
       "when": ["f32_beta_bias_residual_contract", "vec4Aligned", "hasSubgroups or \"bias\" == \"bias\" or portableWideExecution", "normResourcesFit", "rowDispatchFits"],
+      "derive": {
         "scalar": "\"f32\"",
         "vectorScalar": "\"vec4<f32>\"",
         "hasBias": "\"bias\" == \"bias\"",
         "workgroupSize": "skipWg",
         "HIDDEN_LEN": "hiddenSize / 4"
       },
         {
           "id": "normalize",
           "name": "SkipLayerNormalization.Vec4.Normalize",
+          "shader": "norm-skip-row-vec4.wgsl.jinja",
+          "derive": {
+            "simplified": false,
+            "hasBias": "\"bias\" == \"bias\"",
+            "hasBeta": true,
+            "writeResidualSum": true,
+            "usesF16Spec": false,
+            "hidden": "hiddenSize",
+            "hiddenVec": "hiddenSize / 4",
+            "wg": "skipWgVec4",
+            "vecType": "\"vec4<f32>\"",
+            "useSubgroups": "hasSubgroups"
           },
+          "bindings": ["input", "skip", "bias", "gamma", "beta", "output", "input_skip_bias_sum", "params_2"],
+          "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
       "id": "beta_bias_row",
       "priority": 5,
       "when": ["f32_beta_bias_residual_contract", "normResourcesFit", "rowDispatchFits"],
+      "derive": {
         "simplified": false,
         "useSubgroups": "hasSubgroups",
         "hasBeta": true,
         "writeResidualSum": true,
         "hasBias": "\"bias\" == \"bias\"",
         "scalar": "\"f32\"",
         "workgroupSize": "skipWg",
         "HIDDEN_LEN": "hiddenSize"
       },
           "id": "normalize",
           "name": "SkipLayerNormalization.Row.Normalize",
           "shader": "norm-skip-row.wgsl.jinja",
+          "derive": { "writeResidualSum": true },
+          "bindings": ["input_2", "skip_2", "bias_2", "gamma_2", "beta_2", "output_2", "input_skip_bias_sum_2", "params_3"],
+          "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
       "id": "beta_bias_vec4_f16",
       "priority": 21,
       "when": ["f16_beta_bias_residual_contract", "vec4Aligned", "normResourcesFit", "rowDispatchFits"],
+      "derive": { "scalar": "\"f16\"", "vectorScalar": "\"vec4<f16>\"", "HIDDEN_LEN": "hiddenSize / 4" },
       "passes": [
         {
           "id": "main",
           "name": "SkipLayerNormalization.Vec4",
+          "shader": "norm-skip-row-vec4.wgsl.jinja",
+          "derive": {
+            "simplified": false,
+            "hasBias": true,
+            "hasBeta": true,
+            "writeResidualSum": true,
+            "usesF16Spec": true,
+            "hidden": "hiddenSize",
+            "hiddenVec": "hiddenSize / 4",
+            "wg": "skipWgVec4",
+            "vecType": "\"vec4<f16>\"",
+            "useSubgroups": "hasSubgroups"
           },
+          "bindings": ["input", "skip", "bias", "gamma", "beta", "output", "input_skip_bias_sum", "params_2"],
+          "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
       "id": "no_beta_output_only_vec4",
       "priority": 20,
       "when": ["f32_no_beta_output_contract", "vec4Aligned", "normResourcesFit", "rowDispatchFits"],
+      "derive": { "scalar": "\"f32\"", "vectorScalar": "\"vec4<f32>\"", "HIDDEN_LEN": "hiddenSize / 4" },
       "passes": [
         {
           "id": "main",
           "name": "SkipLayerNormalization.NoBetaOutputOnly.Vec4",
+          "shader": "norm-skip-row-vec4.wgsl.jinja",
+          "derive": {
+            "simplified": false,
+            "hasBias": false,
+            "hasBeta": false,
+            "writeResidualSum": false,
+            "usesF16Spec": false,
+            "hidden": "hiddenSize",
+            "hiddenVec": "hiddenSize / 4",
+            "wg": "skipWgVec4",
+            "vecType": "\"vec4<f32>\"",
+            "useSubgroups": "hasSubgroups"
           },
+          "bindings": ["input", "skip", "gamma", "output", "params_2"],
+          "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
       "id": "no_beta_output_only_row",
       "priority": 10,
       "when": ["f32_no_beta_output_contract", "normResourcesFit", "rowDispatchFits"],
+      "derive": {
         "simplified": false,
         "hasBias": false,
         "hasBeta": false,
         "writeResidualSum": false,
         "useSubgroups": "hasSubgroups",
         "scalar": "\"f32\"",
         "workgroupSize": "skipWg",
         "HIDDEN_LEN": "hiddenSize"
       },
           "id": "main",
           "name": "SkipLayerNormalization.NoBetaOutputOnly.Row",
           "shader": "norm-skip-row.wgsl.jinja",
+          "bindings": ["input_2", "skip_2", "gamma_2", "output_2", "params_3"],
+          "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
+        }
+      ]
+    },
+    {
+      "id": "no_beta_output_only_vec4_f16",
+      "priority": 20,
+      "when": ["f16_no_beta_output_contract", "vec4Aligned", "normResourcesFit", "rowDispatchFits"],
+      "derive": { "scalar": "\"f16\"", "vectorScalar": "\"vec4<f16>\"", "HIDDEN_LEN": "hiddenSize / 4" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "SkipLayerNormalization.NoBetaOutputOnly.Vec4.F16",
+          "shader": "norm-skip-row-vec4.wgsl.jinja",
+          "derive": {
+            "simplified": false,
+            "hasBias": false,
+            "hasBeta": false,
+            "writeResidualSum": false,
+            "usesF16Spec": true,
+            "hidden": "hiddenSize",
+            "hiddenVec": "hiddenSize / 4",
+            "wg": "skipWgVec4",
+            "vecType": "\"vec4<f16>\"",
+            "useSubgroups": "hasSubgroups"
+          },
+          "bindings": ["input", "skip", "gamma", "output", "params_2"],
+          "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
+        }
+      ]
+    },
+    {
+      "id": "no_beta_output_only_row_f16",
+      "priority": 10,
+      "when": ["f16_no_beta_output_contract", "normResourcesFit", "rowDispatchFits"],
+      "derive": {
+        "simplified": false,
+        "hasBias": false,
+        "hasBeta": false,
+        "writeResidualSum": false,
+        "useSubgroups": "hasSubgroups",
+        "scalar": "\"f16\"",
+        "usesF16": true,
+        "workgroupSize": "skipWg",
+        "HIDDEN_LEN": "hiddenSize"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "SkipLayerNormalization.NoBetaOutputOnly.Row.F16",
+          "shader": "norm-skip-row.wgsl.jinja",
+          "bindings": ["input_2", "skip_2", "gamma_2", "output_2", "params_3"],
+          "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
       "id": "beta_no_bias_vec4",
       "priority": 20,
       "when": ["f32_beta_no_bias_residual_contract", "vec4Aligned", "hasSubgroups or portableWideExecution", "normResourcesFit", "rowDispatchFits"],
+      "derive": {
         "scalar": "\"f32\"",
         "vectorScalar": "\"vec4<f32>\"",
         "hasBias": "\"no_bias\" == \"bias\"",
         "workgroupSize": "skipWg",
         "HIDDEN_LEN": "hiddenSize / 4"
       },
         {
           "id": "main",
           "name": "SkipLayerNormalization.Vec4",
+          "shader": "norm-skip-row-vec4.wgsl.jinja",
+          "derive": {
+            "simplified": false,
+            "hasBias": "\"no_bias\" == \"bias\"",
+            "hasBeta": true,
+            "writeResidualSum": true,
+            "usesF16Spec": false,
+            "hidden": "hiddenSize",
+            "hiddenVec": "hiddenSize / 4",
+            "wg": "skipWgVec4",
+            "vecType": "\"vec4<f32>\"",
+            "useSubgroups": "hasSubgroups"
           },
+          "bindings": ["input", "skip", "gamma", "beta", "output", "input_skip_bias_sum", "params_2"],
+          "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
       "id": "beta_no_bias_row",
       "priority": 10,
       "when": ["f32_beta_no_bias_residual_contract", "normResourcesFit", "rowDispatchFits"],
+      "derive": {
         "simplified": false,
         "useSubgroups": "hasSubgroups",
         "hasBeta": true,
         "writeResidualSum": true,
         "hasBias": "\"no_bias\" == \"bias\"",
         "scalar": "\"f32\"",
         "workgroupSize": "skipWg",
         "HIDDEN_LEN": "hiddenSize"
       },
           "id": "main",
           "name": "SkipLayerNormalization.Row",
           "shader": "norm-skip-row.wgsl.jinja",
+          "bindings": ["input_2", "skip_2", "gamma_2", "beta_2", "output_2", "input_skip_bias_sum_2", "params_3"],
+          "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
       "id": "beta_no_bias_output_only_vec4",
       "priority": 20,
       "when": ["f32_beta_no_bias_output_only_contract", "vec4Aligned", "hasSubgroups or \"no_bias\" == \"bias\" or portableWideExecution", "normResourcesFit", "rowDispatchFits"],
+      "derive": {
         "scalar": "\"f32\"",
         "vectorScalar": "\"vec4<f32>\"",
         "hasBias": "\"no_bias\" == \"bias\"",
         "workgroupSize": "skipWg",
         "HIDDEN_LEN": "hiddenSize / 4"
       },
         {
           "id": "main",
           "name": "SkipLayerNormalization.Vec4",
+          "shader": "norm-skip-row-vec4.wgsl.jinja",
+          "derive": {
+            "simplified": false,
+            "hasBias": "\"no_bias\" == \"bias\"",
+            "hasBeta": true,
+            "writeResidualSum": false,
+            "usesF16Spec": false,
+            "hidden": "hiddenSize",
+            "hiddenVec": "hiddenSize / 4",
+            "wg": "skipWgVec4",
+            "vecType": "\"vec4<f32>\"",
+            "useSubgroups": "hasSubgroups"
           },
+          "bindings": ["input", "skip", "gamma", "beta", "output", "params_2"],
+          "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
       "id": "beta_no_bias_output_only_row",
       "priority": 10,
       "when": ["f32_beta_no_bias_output_only_contract", "normResourcesFit", "rowDispatchFits"],
+      "derive": {
         "simplified": false,
         "useSubgroups": "hasSubgroups",
         "hasBeta": true,
         "writeResidualSum": false,
         "hasBias": "\"no_bias\" == \"bias\"",
         "scalar": "\"f32\"",
         "workgroupSize": "skipWg",
         "HIDDEN_LEN": "hiddenSize"
       },
           "id": "main",
           "name": "SkipLayerNormalization.Row",
           "shader": "norm-skip-row.wgsl.jinja",
+          "bindings": ["input_2", "skip_2", "gamma_2", "beta_2", "output_2", "params_3"],
+          "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
       "id": "beta_bias_output_only_vec4",
       "priority": 20,
       "when": ["f32_beta_bias_output_only_contract", "vec4Aligned", "hasSubgroups or \"bias\" == \"bias\" or portableWideExecution", "normResourcesFit", "rowDispatchFits"],
+      "derive": {
         "scalar": "\"f32\"",
         "vectorScalar": "\"vec4<f32>\"",
         "hasBias": "\"bias\" == \"bias\"",
         "workgroupSize": "skipWg",
         "HIDDEN_LEN": "hiddenSize / 4"
       },
         {
           "id": "main",
           "name": "SkipLayerNormalization.Vec4",
+          "shader": "norm-skip-row-vec4.wgsl.jinja",
+          "derive": {
+            "simplified": false,
+            "hasBias": "\"bias\" == \"bias\"",
+            "hasBeta": true,
+            "writeResidualSum": false,
+            "usesF16Spec": false,
+            "hidden": "hiddenSize",
+            "hiddenVec": "hiddenSize / 4",
+            "wg": "skipWgVec4",
+            "vecType": "\"vec4<f32>\"",
+            "useSubgroups": "hasSubgroups"
           },
+          "bindings": ["input", "skip", "gamma", "beta", "bias", "output", "params_2"],
+          "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
       "id": "beta_bias_output_only_row",
       "priority": 10,
       "when": ["f32_beta_bias_output_only_contract", "normResourcesFit", "rowDispatchFits"],
+      "derive": {
         "simplified": false,
         "useSubgroups": "hasSubgroups",
         "hasBeta": true,
         "writeResidualSum": false,
         "hasBias": "\"bias\" == \"bias\"",
         "scalar": "\"f32\"",
         "workgroupSize": "skipWg",
         "HIDDEN_LEN": "hiddenSize"
       },
           "id": "main",
           "name": "SkipLayerNormalization.Row",
           "shader": "norm-skip-row.wgsl.jinja",
+          "bindings": ["input_2", "skip_2", "bias_2", "gamma_2", "beta_2", "output_2", "params_3"],
+          "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
+        }
+      ]
+    },
+    {
+      "id": "beta_no_bias_output_only_vec4_f16",
+      "priority": 20,
+      "when": ["f16_beta_no_bias_output_only_contract", "vec4Aligned", "hasSubgroups or \"no_bias\" == \"bias\" or portableWideExecution", "normResourcesFit", "rowDispatchFits"],
+      "derive": {
+        "scalar": "\"f16\"",
+        "vectorScalar": "\"vec4<f16>\"",
+        "hasBias": "\"no_bias\" == \"bias\"",
+        "workgroupSize": "skipWg",
+        "HIDDEN_LEN": "hiddenSize / 4"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "SkipLayerNormalization.Vec4.F16",
+          "shader": "norm-skip-row-vec4.wgsl.jinja",
+          "derive": {
+            "simplified": false,
+            "hasBias": "\"no_bias\" == \"bias\"",
+            "hasBeta": true,
+            "writeResidualSum": false,
+            "usesF16Spec": true,
+            "hidden": "hiddenSize",
+            "hiddenVec": "hiddenSize / 4",
+            "wg": "skipWgVec4",
+            "vecType": "\"vec4<f16>\"",
+            "useSubgroups": "hasSubgroups"
+          },
+          "bindings": ["input", "skip", "gamma", "beta", "output", "params_2"],
+          "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
+        }
+      ]
+    },
+    {
+      "id": "beta_no_bias_output_only_row_f16",
+      "priority": 10,
+      "when": ["f16_beta_no_bias_output_only_contract", "normResourcesFit", "rowDispatchFits"],
+      "derive": {
+        "simplified": false,
+        "useSubgroups": "hasSubgroups",
+        "hasBeta": true,
+        "writeResidualSum": false,
+        "hasBias": "\"no_bias\" == \"bias\"",
+        "scalar": "\"f16\"",
+        "usesF16": true,
+        "workgroupSize": "skipWg",
+        "HIDDEN_LEN": "hiddenSize"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "SkipLayerNormalization.Row.F16",
+          "shader": "norm-skip-row.wgsl.jinja",
+          "bindings": ["input_2", "skip_2", "gamma_2", "beta_2", "output_2", "params_3"],
+          "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
+        }
+      ]
+    },
+    {
+      "id": "beta_bias_output_only_vec4_f16",
+      "priority": 20,
+      "when": ["f16_beta_bias_output_only_contract", "vec4Aligned", "hasSubgroups or \"bias\" == \"bias\" or portableWideExecution", "normResourcesFit", "rowDispatchFits"],
+      "derive": {
+        "scalar": "\"f16\"",
+        "vectorScalar": "\"vec4<f16>\"",
+        "hasBias": "\"bias\" == \"bias\"",
+        "workgroupSize": "skipWg",
+        "HIDDEN_LEN": "hiddenSize / 4"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "SkipLayerNormalization.Vec4.F16",
+          "shader": "norm-skip-row-vec4.wgsl.jinja",
+          "derive": {
+            "simplified": false,
+            "hasBias": "\"bias\" == \"bias\"",
+            "hasBeta": true,
+            "writeResidualSum": false,
+            "usesF16Spec": true,
+            "hidden": "hiddenSize",
+            "hiddenVec": "hiddenSize / 4",
+            "wg": "skipWgVec4",
+            "vecType": "\"vec4<f16>\"",
+            "useSubgroups": "hasSubgroups"
+          },
+          "bindings": ["input", "skip", "gamma", "beta", "bias", "output", "params_2"],
+          "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
+        }
+      ]
+    },
+    {
+      "id": "beta_bias_output_only_row_f16",
+      "priority": 10,
+      "when": ["f16_beta_bias_output_only_contract", "normResourcesFit", "rowDispatchFits"],
+      "derive": {
+        "simplified": false,
+        "useSubgroups": "hasSubgroups",
+        "hasBeta": true,
+        "writeResidualSum": false,
+        "hasBias": "\"bias\" == \"bias\"",
+        "scalar": "\"f16\"",
+        "usesF16": true,
+        "workgroupSize": "skipWg",
+        "HIDDEN_LEN": "hiddenSize"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "SkipLayerNormalization.Row.F16",
+          "shader": "norm-skip-row.wgsl.jinja",
+          "bindings": ["input_2", "skip_2", "bias_2", "gamma_2", "beta_2", "output_2", "params_3"],
+          "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     }
+  ]
 }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,19 +1,41 @@
 {
   "name": "com.microsoft.SkipLayerNormalization",
-  "id": "_com_microsoft_skiplayernormalization_webgpu_1767bd7",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "STz1miuskWFMJijUblWNS3pG/DHJsRwQo0M+up/ng3s=",
-      "manifest.json": "MM5RYZ3xnC2rbelhnBH9kdS+qwkWaQmvEVXhBmag8jg=",
-      "norm-skip-row-vec4.wgsl.jinja": "qHoa8QjoRD1rgxTGGuJpfVb16Jou8PEg2CF1GjDt78s=",
-      "norm-skip-row.wgsl.jinja": "ck3s/G7WRph2YZRJs4LTft+R0IDjMPeE5Oy4UXDsqQs=",
-      "test.json": "Dkphsy4eO4Z8f2ZY1wfBJVmLcKNVWcOupJHOighE3RU="
     }
   },
-  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
-  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.SkipLayerNormalization" }
 }

 {
   "name": "com.microsoft.SkipLayerNormalization",
+  "id": "_com_microsoft_skiplayernormalization_webgpu_28a934b",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "dZADuGnHvG3kWGPS0BGgmikuwdi16tD/Sk4HbGC9UV0=",
+      "manifest.json": "G6ZofqPdSuhkAxEuuKIlW4fQ/fCsDwxEzafKwVqGEY8=",
+      "norm-skip-row-vec4.wgsl.jinja": "bH8L9BYJ/3XQXE2zP8kvZc4wlhAbcAolI4/p0xXskzM=",
+      "norm-skip-row.wgsl.jinja": "aERQfDDXx6lwKuyahwcemgFllRSq46egf2++7f6bNtA=",
+      "test.json": "K+KtXhWTSxbRlFwFfszdd1YKz9JnAOuDZ9ebLg7Y8w0="
     }
   },
+  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.0",
+    "variants": {
+      "beta_output_only_vec4_broadcast": ["norm-skip-row-vec4.wgsl.jinja"],
+      "beta_bias_vec4": ["norm-skip-row-vec4.wgsl.jinja"],
+      "beta_bias_row": ["norm-skip-row.wgsl.jinja"],
+      "beta_bias_vec4_f16": ["norm-skip-row-vec4.wgsl.jinja"],
+      "no_beta_output_only_vec4": ["norm-skip-row-vec4.wgsl.jinja"],
+      "no_beta_output_only_row": ["norm-skip-row.wgsl.jinja"],
+      "no_beta_output_only_vec4_f16": ["norm-skip-row-vec4.wgsl.jinja"],
+      "no_beta_output_only_row_f16": ["norm-skip-row.wgsl.jinja"],
+      "beta_no_bias_vec4": ["norm-skip-row-vec4.wgsl.jinja"],
+      "beta_no_bias_row": ["norm-skip-row.wgsl.jinja"],
+      "beta_no_bias_output_only_vec4": ["norm-skip-row-vec4.wgsl.jinja"],
+      "beta_no_bias_output_only_row": ["norm-skip-row.wgsl.jinja"],
+      "beta_bias_output_only_vec4": ["norm-skip-row-vec4.wgsl.jinja"],
+      "beta_bias_output_only_row": ["norm-skip-row.wgsl.jinja"],
+      "beta_no_bias_output_only_vec4_f16": ["norm-skip-row-vec4.wgsl.jinja"],
+      "beta_no_bias_output_only_row_f16": ["norm-skip-row.wgsl.jinja"],
+      "beta_bias_output_only_vec4_f16": ["norm-skip-row-vec4.wgsl.jinja"],
+      "beta_bias_output_only_row_f16": ["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<vec2<f32>, WG>;
@@ -84,9 +85,9 @@ fn reduce_pair(value: vec2<f32>{% if useSubgroups %}, sg_lane: u32, sg_id: u32,
 // 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,7 +107,7 @@ fn main(
   }
   let tid = lid.x;
   let base = row * HIDDEN_V;
-{% if source.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;
@@ -114,11 +115,11 @@ fn main(
   let skip_base = base;
 {% endif %}
-  let shift = residual_value(base, skip_base{% if source.hasBias %}, 0u{% endif %}).x;
   var acc = vec2<f32>(0.0, 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 %});
     let d = v - vec4<f32>(shift);
     acc.x = acc.x + d.x + d.y + d.z + d.w;
     acc.y = acc.y + dot(d, d);
@@ -132,14 +133,14 @@ 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 %}
     var value = (residual - vec4<f32>(row_mean)) * row_inv * vec4<f32>(gamma[i]);
-{% if source.hasBeta %}
     value = value + vec4<f32>(beta[i]);
 {% endif %}
-    output[idx] = {{ source.vecType }}(value);
   }
 }

     {{ 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<vec2<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;
   let skip_base = base;
 {% endif %}
+  let shift = residual_value(base, skip_base{% if hasBias %}, 0u{% endif %}).x;
   var acc = vec2<f32>(0.0, 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 %});
     let d = v - vec4<f32>(shift);
     acc.x = acc.x + d.x + d.y + d.z + d.w;
     acc.y = acc.y + dot(d, d);
   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 %}
     var value = (residual - vec4<f32>(row_mean)) * row_inv * vec4<f32>(gamma[i]);
+{% if hasBeta %}
     value = value + vec4<f32>(beta[i]);
 {% endif %}
+    output[idx] = {{ vecType }}(value);
   }
 }

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

@@ -47,6 +47,9 @@
 /* One workgroup normalizes each row of residual = input + skip, with an
  * optional bias. */
 {% set degenerateRow = (not simplified) and hiddenSize == 1 %}
 {% if useSubgroups and not degenerateRow %}
 enable subgroups;
 {% endif %}
@@ -56,46 +59,6 @@ enable subgroups;
 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) %}
-fn {{ name }}(value: f32, tid: u32) -> f32 {
-  {{ buffer }}[tid] = value;
-  workgroupBarrier();
-  // Ceil-halving keeps every lane when the workgroup size is not a power of
-  // two. For even n this matches the power-of-two tree order; for odd n, lanes
-  // [0, n-half) fold the upper tail while the middle lane carries forward.
-  var n: u32 = {{ wg }};
-  loop {
-    let half = (n + 1u) / 2u;
-    if (tid < n - half) {
-{% if mode == "max" %}
-      {{ buffer }}[tid] = max({{ buffer }}[tid], {{ buffer }}[tid + half]);
-{% else %}
-      {{ buffer }}[tid] = {{ buffer }}[tid] + {{ buffer }}[tid + half];
-{% endif %}
-    }
-    workgroupBarrier();
-    n = half;
-    if (n == 1u) {
-      break;
-    }
-  }
-  // The default trailing barrier makes this helper safe for back-to-back calls: every lane reads
-  // slot 0 here, so the next call's first store must not run until all lanes have read it.
-  // `trailingBarrier=false` is safe only when the buffer is never written again before kernel exit.
-  let reduced = {{ buffer }}[0];
-{% if trailingBarrier %}
-  workgroupBarrier();
-{% endif %}
-  return reduced;
-}
-{% endmacro %}
-{{ 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>;
@@ -125,7 +88,6 @@ fn reduce_pair(value: vec2<f32>, tid: u32) -> vec2<f32> {
 }
 {% endif %}
 {% endif %}
-{% endif %}
 {% if not degenerateRow or writeResidualSum %}
 fn residual_value(row: u32, d: u32) -> f32 {
@@ -140,47 +102,23 @@ fn residual_value(row: u32, d: u32) -> f32 {
 @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.
-  var local_sq = 0.0;
-  for (var d: u32 = tid; d < HIDDEN; d = d + WG) {
-    let value = residual_value(row, d);
-    local_sq = local_sq + value * value;
-  }
-  let sq = reduce_sum(local_sq, tid);
-  if (tid == 0u) {
-    row_inv = inverseSqrt(sq / f32(HIDDEN) + params.epsilon);
-  }
-  workgroupBarrier();
-  for (var d: u32 = 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_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

 /* 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 %}
 const HIDDEN: u32 = {{ hiddenSize }}u;
 {% endif %}
 const WG: u32 = {{ workgroupSize }}u;
 {% if not degenerateRow %}
 var<workgroup> pair_partial: array<vec2<f32>, WG>;
 }
 {% endif %}
 {% endif %}
 {% if not degenerateRow or writeResidualSum %}
 fn residual_value(row: u32, d: u32) -> f32 {
 @compute @workgroup_size(WG, 1, 1)
 fn main(
+  @builtin(workgroup_id) wg: 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 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;
   }
 {% if not degenerateRow %}
   let tid = lid.x;
 {% endif %}
+{% if 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

build/webgpu/test.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "com.microsoft.SkipLayerNormalization",
   "fixtureArrays": {
     "ort_batch2_bias_flattened_tokens_input_skipT": [0.1, -0.2, 0.3, 1, 0.5, 0.1, 0.4, 1.6, 1.8, -0.3, 0, 1, -0.5, 0.4, 0.8, -0.6],
     "ort_batch2_flattened_tokens_input_inputT": [0.8, -0.5, 0, 1, 0.5, 0.2, 0.3, -0.6, 0.8, -0.5, 0, 1, 0.5, 0.2, 0.3, -0.6],
@@ -189,7 +188,7 @@
     {
       "name": "hidden_size_one_bias_rows65535_dispatch_edge",
       "provenance": {
-        "notes": "Correctness sibling for the low-occupancy hidden=1 benchmark at rows=65535, the maximum single-dimension workgroup count edge."
       },
       "attrs": { "epsilon": 0.00001 },
       "inputs": {
@@ -224,7 +223,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] } },
         "betaT": { "dtype": "float32", "shape": [4], "data": { "kind": "values", "values": [0.0, 0.5, -0.25, 1.0] } }
@@ -258,7 +257,7 @@
       "provenance": {
         "source": "onnxruntime/test/contrib_ops/skiplayernorm_op_test.cc",
         "test": "SkipLayerNormTest.SkipLayerNormBatch1",
-        "notes": "ORT shape [1, 2, 4] is represented as [2, 4] tokens by this lowered kernel. Epsilon is omitted to exercise the schema default of 1e-12."
       },
       "inputs": {
         "inputT": {
@@ -284,7 +283,7 @@
       "provenance": {
         "source": "onnxruntime/test/contrib_ops/skiplayernorm_op_test.cc",
         "test": "SkipLayerNormTest.SkipLayerNormBatch2_Bias",
-        "notes": "ORT shape [2, 2, 4] is represented as [4, 4] tokens by this lowered kernel."
       },
       "attrs": { "epsilon": 1e-12 },
       "inputs": {
@@ -318,7 +317,7 @@
       "provenance": {
         "source": "onnxruntime/test/contrib_ops/skiplayernorm_op_test.cc",
         "test": "SkipLayerNormTest.SkipLayerNormBatch2",
-        "notes": "ORT shape [2, 2, 4] is represented as [4, 4] tokens by this lowered kernel."
       },
       "attrs": { "epsilon": 1e-12 },
       "inputs": {
@@ -477,7 +476,7 @@
     {
       "name": "no_beta_output_only_hidden6_unaligned_row",
       "provenance": {
-        "notes": "Only coverage for the scalar no_beta_output_only_row variant: beta omitted with output as the sole output picks the no-beta output-only pair, and hidden=6 is not a multiple of 4, so the vec4 sibling's vec4Aligned guard is false and the row kernel is selected on every tier."
       },
       "attrs": { "epsilon": 0.00001 },
       "inputs": {
@@ -532,7 +531,7 @@
     {
       "name": "beta_no_bias_output_only_hidden6_unaligned_row",
       "provenance": {
-        "notes": "Row-kernel arm of beta_no_bias_output_only: hidden=6 fails vec4Aligned, so the scalar sibling is selected on every tier."
       },
       "attrs": { "epsilon": 0.00001 },
       "inputs": {
@@ -701,7 +700,7 @@
     {
       "name": "f32_hidden768_no_bias_residual",
       "provenance": {
-        "notes": "Compact companion for bench cases skip-layernorm-f32-4096x768 and skip-layernorm-f32-4096x768-hidden768-aligned-sibling, where reduced tiers expose a subgroup fallback gap."
       },
       "attrs": { "epsilon": 0.00001 },
       "inputs": {
@@ -769,7 +768,7 @@
     {
       "name": "f32_hidden1025_bias_residual",
       "provenance": {
-        "notes": "Compact companion for bench case skip-layernorm-f32-bias-8192x1025-unaligned-twopass-fallback, preserving the odd hidden-size bias path."
       },
       "attrs": { "epsilon": 0.00001 },
       "inputs": {
@@ -911,9 +910,7 @@
     },
     {
       "name": "hidden_size_one_bias_many_rows",
-      "provenance": {
-        "notes": "Compact correctness sibling for the hidden-size-one low-occupancy benchmark; many rows with hidden=1 keep the variance-zero bias path honest without using benchmark-scale dimensions."
-      },
       "attrs": { "epsilon": 0.00001 },
       "inputs": {
         "inputT": {
@@ -934,6 +931,247 @@
         "outputT": { "dtype": "float32", "shape": [257, 1], "tolerance": 0.000001 },
         "residualT": { "dtype": "float32", "shape": [257, 1], "tolerance": 0.000001 }
       }
     }
   ]
 }

 {
   "fixtureArrays": {
     "ort_batch2_bias_flattened_tokens_input_skipT": [0.1, -0.2, 0.3, 1, 0.5, 0.1, 0.4, 1.6, 1.8, -0.3, 0, 1, -0.5, 0.4, 0.8, -0.6],
     "ort_batch2_flattened_tokens_input_inputT": [0.8, -0.5, 0, 1, 0.5, 0.2, 0.3, -0.6, 0.8, -0.5, 0, 1, 0.5, 0.2, 0.3, -0.6],
     {
       "name": "hidden_size_one_bias_rows65535_dispatch_edge",
       "provenance": {
+        "notes": "Hidden size 1 across 65,535 rows exercises the maximum single-dimension workgroup count with one value per normalization row."
       },
       "attrs": { "epsilon": 0.00001 },
       "inputs": {
         "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] } },
         "betaT": { "dtype": "float32", "shape": [4], "data": { "kind": "values", "values": [0.0, 0.5, -0.25, 1.0] } }
       "provenance": {
         "source": "onnxruntime/test/contrib_ops/skiplayernorm_op_test.cc",
         "test": "SkipLayerNormTest.SkipLayerNormBatch1",
+        "notes": "This package flattens ORT shape [1, 2, 4] to two four-wide token rows. 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.SkipLayerNormBatch2_Bias",
+        "notes": "This package flattens ORT shape [2, 2, 4] to four four-wide token rows."
       },
       "attrs": { "epsilon": 1e-12 },
       "inputs": {
       "provenance": {
         "source": "onnxruntime/test/contrib_ops/skiplayernorm_op_test.cc",
         "test": "SkipLayerNormTest.SkipLayerNormBatch2",
+        "notes": "This package flattens ORT shape [2, 2, 4] to four four-wide token rows."
       },
       "attrs": { "epsilon": 1e-12 },
       "inputs": {
     {
       "name": "no_beta_output_only_hidden6_unaligned_row",
       "provenance": {
+        "notes": "Beta and optional outputs are omitted, and hidden size 6 is not divisible by four, selecting the scalar no-beta output-only row path on every tier."
       },
       "attrs": { "epsilon": 0.00001 },
       "inputs": {
     {
       "name": "beta_no_bias_output_only_hidden6_unaligned_row",
       "provenance": {
+        "notes": "Beta is present, bias and optional outputs are omitted, and hidden size 6 selects the scalar output-only row path on every tier."
       },
       "attrs": { "epsilon": 0.00001 },
       "inputs": {
     {
       "name": "f32_hidden768_no_bias_residual",
       "provenance": {
+        "notes": "A compact hidden-size-768 residual normalization exercises the subgroup route and its reduced-tier fallback without bias."
       },
       "attrs": { "epsilon": 0.00001 },
       "inputs": {
     {
       "name": "f32_hidden1025_bias_residual",
       "provenance": {
+        "notes": "A compact hidden-size-1025 normalization exercises the unaligned two-pass path with bias and residual output."
       },
       "attrs": { "epsilon": 0.00001 },
       "inputs": {
     },
     {
       "name": "hidden_size_one_bias_many_rows",
+      "provenance": { "notes": "Many hidden-size-one rows exercise the variance-zero bias path at a compact scale." },
       "attrs": { "epsilon": 0.00001 },
       "inputs": {
         "inputT": {
         "outputT": { "dtype": "float32", "shape": [257, 1], "tolerance": 0.000001 },
         "residualT": { "dtype": "float32", "shape": [257, 1], "tolerance": 0.000001 }
       }
+    },
+    {
+      "name": "rank3_no_bias_residual",
+      "provenance": {
+        "notes": "Rank-3 activation shape carrying the residual output, the form ONNX Runtime's transformer fusion emits when the pre-normalization sum feeds the next block."
+      },
+      "attrs": { "epsilon": 0.00001 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [1, 4, 768],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.017, "cosStep": 0.031 }
+        },
+        "skipT": {
+          "dtype": "float32",
+          "shape": [1, 4, 768],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.023 }
+        },
+        "gammaT": {
+          "dtype": "float32",
+          "shape": [768],
+          "data": { "kind": "fillFloat32", "scale": 0.1, "offset": 1.0, "sinStep": 0.007, "cosStep": 0.041 }
+        },
+        "betaT": {
+          "dtype": "float32",
+          "shape": [768],
+          "data": { "kind": "fillFloat32", "scale": 0.1, "sinStep": 0.019, "cosStep": 0.013 }
+        }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 4, 768], "tolerance": 0.0002 },
+        "residualT": { "dtype": "float32", "shape": [1, 4, 768], "tolerance": 0.000002 }
+      }
+    },
+    {
+      "name": "rank3_beta_no_bias_output_only_hidden6_unaligned_row",
+      "provenance": {
+        "notes": "Rank-3 at a hidden size that is not a multiple of four, so the scalar row route serves it."
+      },
+      "attrs": { "epsilon": 0.00001 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float32",
+          "shape": [1, 2, 6],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.37 }
+        },
+        "skipT": {
+          "dtype": "float32",
+          "shape": [1, 2, 6],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.43 }
+        },
+        "gammaT": {
+          "dtype": "float32",
+          "shape": [6],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "scale": 0.2 }
+        },
+        "betaT": {
+          "dtype": "float32",
+          "shape": [6],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.1 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [1, 2, 6], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "f16_rank3_beta_bias_output_only",
+      "provenance": {
+        "notes": "Half-precision beta and bias at a rank-3 activation shape, the ordinary on-device inference form."
+      },
+      "requires": { "features": ["shader-f16"] },
+      "attrs": { "epsilon": 0.00001 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float16",
+          "shape": [1, 3, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "skipT": {
+          "dtype": "float16",
+          "shape": [1, 3, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "gammaT": {
+          "dtype": "float16",
+          "shape": [8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "scale": 0.2 }
+        },
+        "betaT": {
+          "dtype": "float16",
+          "shape": [8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.1 }
+        },
+        "biasT": {
+          "dtype": "float16",
+          "shape": [8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.05, "cosStep": 0.47, "scale": 0.3 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float16", "shape": [1, 3, 8], "tolerance": 0.01 } }
+    },
+    {
+      "name": "f16_beta_bias_output_only_hidden6_unaligned_row",
+      "provenance": { "notes": "Half precision at an unaligned hidden size, which the scalar row route serves." },
+      "requires": { "features": ["shader-f16"] },
+      "attrs": { "epsilon": 0.00001 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float16",
+          "shape": [2, 6],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.37 }
+        },
+        "skipT": {
+          "dtype": "float16",
+          "shape": [2, 6],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.43 }
+        },
+        "gammaT": {
+          "dtype": "float16",
+          "shape": [6],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "scale": 0.2 }
+        },
+        "betaT": {
+          "dtype": "float16",
+          "shape": [6],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.1 }
+        },
+        "biasT": {
+          "dtype": "float16",
+          "shape": [6],
+          "data": { "kind": "fillFloat32", "sinStep": 0.05, "cosStep": 0.47, "scale": 0.3 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float16", "shape": [2, 6], "tolerance": 0.01 } }
+    },
+    {
+      "name": "f16_beta_no_bias_output_only",
+      "provenance": { "notes": "Half precision with beta and no bias." },
+      "requires": { "features": ["shader-f16"] },
+      "attrs": { "epsilon": 0.00001 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float16",
+          "shape": [3, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "skipT": {
+          "dtype": "float16",
+          "shape": [3, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "gammaT": {
+          "dtype": "float16",
+          "shape": [8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "scale": 0.2 }
+        },
+        "betaT": {
+          "dtype": "float16",
+          "shape": [8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.1 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float16", "shape": [3, 8], "tolerance": 0.01 } }
+    },
+    {
+      "name": "f16_beta_no_bias_output_only_hidden6_unaligned_row",
+      "provenance": { "notes": "Half precision with beta, no bias, at an unaligned hidden size." },
+      "requires": { "features": ["shader-f16"] },
+      "attrs": { "epsilon": 0.00001 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float16",
+          "shape": [2, 6],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.37 }
+        },
+        "skipT": {
+          "dtype": "float16",
+          "shape": [2, 6],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.43 }
+        },
+        "gammaT": {
+          "dtype": "float16",
+          "shape": [6],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "scale": 0.2 }
+        },
+        "betaT": {
+          "dtype": "float16",
+          "shape": [6],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.1 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float16", "shape": [2, 6], "tolerance": 0.01 } }
+    },
+    {
+      "name": "f16_no_beta_output_only_hidden6_unaligned_row",
+      "provenance": { "notes": "Half precision with neither beta nor bias, at an unaligned hidden size." },
+      "requires": { "features": ["shader-f16"] },
+      "attrs": { "epsilon": 0.00001 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float16",
+          "shape": [2, 6],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "skipT": {
+          "dtype": "float16",
+          "shape": [2, 6],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "gammaT": {
+          "dtype": "float16",
+          "shape": [6],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "scale": 0.2 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float16", "shape": [2, 6], "tolerance": 0.002 } }
+    },
+    {
+      "name": "f16_no_beta_output_only",
+      "provenance": {
+        "notes": "Half precision with neither beta nor bias at an aligned hidden size, which is the vec4 arm of that pair."
+      },
+      "requires": { "features": ["shader-f16"] },
+      "attrs": { "epsilon": 0.00001 },
+      "inputs": {
+        "inputT": {
+          "dtype": "float16",
+          "shape": [3, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "skipT": {
+          "dtype": "float16",
+          "shape": [3, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "gammaT": {
+          "dtype": "float16",
+          "shape": [8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "scale": 0.2 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float16", "shape": [3, 8], "tolerance": 0.002 } }
     }
   ]
 }