sync 91d990483a17

Browse files

Files changed (8) hide show

README.md +14 -10
build/webgpu/bench.json +0 -1
build/webgpu/group-normalization-splitk-apply.wgsl.jinja +0 -3
build/webgpu/group-normalization-splitk-partials.wgsl.jinja +0 -3
build/webgpu/manifest.json +98 -174
build/webgpu/metadata.json +17 -9
build/webgpu/norm-row-stats.wgsl.jinja +126 -20
build/webgpu/test.json +2 -3

README.md CHANGED Viewed

@@ -18,17 +18,17 @@ See the [ONNX `GroupNormalization` spec](https://onnx.ai/onnx/operators/onnx__Gr
 ## Inputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `X` | `x` | `T` | — | — | Input data tensor of shape `(N x C x D1 x ... x Dn)` where `N` is batch size and `C` is the number of channels. | required |
-| `scale` | `scale` | `T` | `1` | — | Scale tensor of shape `(C)`, one value per channel. | required |
-| `bias` | `bias` | `T` | `1` | — | Bias tensor of shape `(C)`, one value per channel. | required |
 ## Outputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `Y` | `y` | `T` | same as `X` | same as `X` | Normalized output tensor of the same shape as `X`. | required |
 ## Attributes
@@ -37,8 +37,8 @@ Attributes and default values (overridable per request):
 | Attribute | Default | Description |
 | --- | --- | --- |
 | `epsilon` | `0.00001` | Small value added to the variance denominator to avoid division by zero. |
-| `stash_type` | `1` | TensorProto element type used for the normalization stage: `1` computes in float32, while `10` computes in float16. Normalized values are cast back to the input type before scale and bias are applied. |
 | `num_groups` | — | Required number of groups to divide the channels into; must be a divisor of `C`. |
 ## Type constraints
@@ -48,7 +48,7 @@ Attributes and default values (overridable per request):
 ## Files
-- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
 - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
@@ -59,10 +59,14 @@ Attributes and 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.

 ## Inputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `x` | `X` | `T` | — | — | Input data tensor of shape `(N x C x D1 x ... x Dn)` where `N` is batch size and `C` is the number of channels. | required |
+| `scale` | — | `T` | `1` | — | Scale tensor of shape `(C)`, one value per channel. | required |
+| `bias` | — | `T` | `1` | — | Bias tensor of shape `(C)`, one value per channel. | required |
 ## Outputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `y` | `Y` | `T` | same as `x` | same as `x` | Normalized output tensor of the same shape as `X`. | required |
 ## Attributes
 | Attribute | Default | Description |
 | --- | --- | --- |
 | `epsilon` | `0.00001` | Small value added to the variance denominator to avoid division by zero. |
 | `num_groups` | — | Required number of groups to divide the channels into; must be a divisor of `C`. |
+| `stash_type` | `1` | TensorProto element type used for the normalization stage: `1` computes in float32, while `10` computes in float16. Normalized values are cast back to the input type before scale and bias are applied. |
 ## Type constraints
 ## Files
+- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
 - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 ## Use with `@huggingface/kernels`
+```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.2
+```
+Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

build/webgpu/bench.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "ai.onnx.GroupNormalization",
   "cases": [
     {
       "name": "4x64x128x128_g32",

 {
   "cases": [
     {
       "name": "4x64x128x128_g32",

build/webgpu/group-normalization-splitk-apply.wgsl.jinja CHANGED Viewed

@@ -1,6 +1,3 @@
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 const HIDDEN: u32 = {{ hiddenSize }}u;





1	{{ env.wgsl.resourceDeclarations }}
2
3	const HIDDEN: u32 = {{ hiddenSize }}u;

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

@@ -1,6 +1,3 @@
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 const HIDDEN: u32 = {{ hiddenSize }}u;





1	{{ env.wgsl.resourceDeclarations }}
2
3	const HIDDEN: u32 = {{ hiddenSize }}u;

build/webgpu/manifest.json CHANGED Viewed

@@ -2,55 +2,31 @@
   "domain": "ai.onnx",
   "name": "GroupNormalization",
   "sinceVersion": 21,
-  "description": "Applies group normalization to the input: `y = scale * (x - mean) / sqrt(variance + epsilon) + bias`, where mean and variance are computed per instance per group of channels. The number of groups `num_groups` must divide the channel count `C` evenly; when `num_groups == C` this is equivalent to InstanceNormalization, and when `num_groups == 1` it is equivalent to LayerNormalization. The normalization stage supports TensorProto `stash_type` values `1` (float32) and `10` (float16).",
-  "inputs": [
-    {
-      "role": "X",
-      "dtype": "T",
-      "description": "Input data tensor of shape `(N x C x D1 x ... x Dn)` where `N` is batch size and `C` is the number of channels."
-    },
-    { "role": "scale", "dtype": "T", "rank": 1, "description": "Scale tensor of shape `(C)`, one value per channel." },
-    { "role": "bias", "dtype": "T", "rank": 1, "description": "Bias tensor of shape `(C)`, one value per channel." }
-  ],
-  "outputs": [
-    {
-      "role": "Y",
-      "dtype": "T",
-      "rank": "ranks.X",
-      "description": "Normalized output tensor of the same shape as `X`.",
-      "shape": "shapes.X"
-    }
-  ],
-  "attributes": { "epsilon": 0.00001, "stash_type": 1 },
-  "attributeDescriptions": {
-    "epsilon": "Small value added to the variance denominator to avoid division by zero.",
-    "num_groups": "Required number of groups to divide the channels into; must be a divisor of `C`.",
-    "stash_type": "TensorProto element type used for the normalization stage: `1` computes in float32, while `10` computes in float16. Normalized values are cast back to the input type before scale and bias are applied."
   },
   "attributeConstraints": { "num_groups": { "required": true }, "stash_type": { "values": [1, 10] } },
   "typeConstraints": { "T": ["float32", "float16"] },
-  "args": {
-    "x": { "kind": "tensor", "semantic": "X", "role": "input" },
-    "scale": { "kind": "tensor", "semantic": "scale", "role": "input" },
-    "bias": { "kind": "tensor", "semantic": "bias", "role": "input" },
-    "y": { "kind": "tensor", "semantic": "Y", "role": "output" }
-  },
   "tunables": {
-    "WORKGROUP_SIZE": 256,
-    "MAX_STATS_SPLITS": 64,
-    "STATS_VALUES_PER_SPLIT": 4096,
-    "SPLIT_STATS_MIN_HIDDEN": 65536,
-    "SPLIT_STATS_MAX_ROWS": 256
   },
   "derive": {
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
     "groupAttributesOk": "attrs.num_groups >= 1",
-    "groupShapeOk": "groupAttributesOk and f16Ok(dtypes.T) and ranks.X >= 3 and ranks.scale == 1 and ranks.bias == 1 and ranks.Y == ranks.X and sameShape(shapes.Y, shapes.X) and dim(shapes.scale, 0) == dim(shapes.X, 1) and dim(shapes.bias, 0) == dim(shapes.X, 1) and dim(shapes.X, 1) % attrs.num_groups == 0",
     "groupContractOk": "groupShapeOk and attrs.stash_type == onnxDtypeCode(\"float32\")",
     "groupStashF16Ok": "groupShapeOk and attrs.stash_type == onnxDtypeCode(\"float16\")",
-    "groupRows": "dim(shapes.X, 0) * attrs.num_groups if groupAttributesOk else 0",
-    "groupSpatial": "inner(shapes.X, 1)",
-    "groupChannelsPerGroup": "dim(shapes.X, 1) / attrs.num_groups if groupAttributesOk else 0",
     "groupHidden": "groupChannelsPerGroup * groupSpatial",
     "normDeviceWorkgroupCap": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
     "normWorkgroupCap": "max(1, pow2ceil(normDeviceWorkgroupCap + 1) / 2)",
@@ -59,98 +35,39 @@
     "hasSubgroupId": "device.features.has(\"subgroups\") and device.wgslLanguageFeatures.has(\"subgroup_id\")",
     "groupRowWorkgroupBytes": "normWorkgroupCap * 2 * 4",
     "groupRowCovered": "groupContractOk and groupRowWorkgroupBytes <= device.limits.maxComputeWorkgroupStorageSize",
-    "groupSplitCount": "min(tunables.MAX_STATS_SPLITS, device.limits.maxComputeWorkgroupsPerDimension, pow2ceil(ceilDiv(groupHidden, tunables.STATS_VALUES_PER_SPLIT)))",
     "groupPartialBytes": "groupRows * groupSplitCount * 2 * 4",
-    "groupSplitCovered": "groupRowCovered and groupRows <= tunables.SPLIT_STATS_MAX_ROWS and groupRows <= device.limits.maxComputeWorkgroupsPerDimension and groupHidden >= tunables.SPLIT_STATS_MIN_HIDDEN and groupPartialBytes <= device.limits.maxStorageBufferBindingSize and groupPartialBytes <= device.limits.maxBufferSize"
   },
-  "bindingSets": {
-    "norm": [
-      {
-        "name": "x",
-        "arg": "x",
-        "semantic": "X",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$ioElement"
-      },
-      {
-        "name": "scale",
-        "arg": "scale",
-        "semantic": "scale",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "bias",
-        "arg": "bias",
-        "semantic": "bias",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$ioElement" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "groupRows" },
-            {
-              "name": "rowStride",
-              "type": "u32",
-              "value": "max(1, min(groupRows, device.limits.maxComputeWorkgroupsPerDimension))"
-            }
-          ]
         }
-      }
-    ],
-    "splitPartials": [
-      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
-      { "name": "partials", "semantic": "partials", "buffer": { "type": "storage" }, "elementType": "vec2<f32>" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "rows", "type": "u32", "value": "groupRows" }] }
-      }
-    ],
-    "splitApply": [
-      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
-      {
-        "name": "scale",
-        "arg": "scale",
-        "semantic": "scale",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "bias",
-        "arg": "bias",
-        "semantic": "bias",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "partials",
-        "semantic": "partials",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "vec2<f32>"
-      },
-      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "rows", "type": "u32", "value": "groupRows" }] }
-      }
-    ]
   },
   "variants": [
     {
       "id": "group_stash_f16_serial",
       "priority": 1000,
-      "when": "groupStashF16Ok",
-      "constants": {
         "scalar": "dtypes.T",
         "ioElement": "dtypes.T",
         "usesF16": "dtypes.T == \"f16\"",
@@ -165,8 +82,8 @@
           "id": "main",
           "name": "GroupNormalization.StashF16Serial",
           "shader": "group-normalization-stash-f16-serial.wgsl.jinja",
-          "bindings": "norm",
-          "dispatch": { "workgroups": "groupRows" }
         }
       ]
     },
@@ -174,7 +91,7 @@
       "id": "group_splitk",
       "priority": 120,
       "when": ["groupSplitCovered"],
-      "constants": {
         "scalar": "dtypes.T",
         "usesF16": "dtypes.T == \"f16\"",
         "hiddenSize": "groupHidden",
@@ -191,15 +108,30 @@
           "id": "partials",
           "name": "GroupNormalization.SplitKPartials",
           "shader": "group-normalization-splitk-partials.wgsl.jinja",
-          "bindings": "splitPartials",
-          "dispatch": { "workgroups": "groupRows", "z": "groupSplitCount" }
         },
         {
           "id": "apply",
           "name": "GroupNormalization.SplitKApply",
           "shader": "group-normalization-splitk-apply.wgsl.jinja",
-          "bindings": "splitApply",
-          "dispatch": { "workgroups": "groupRows", "z": "groupSplitCount" }
         }
       ]
     },
@@ -207,36 +139,30 @@
       "id": "group_subgroup_vec4",
       "priority": 110,
       "when": ["groupRowCovered", "groupSpatial % 4 == 0"],
-      "constants": {
-        "scalar": "dtypes.T",
-        "ioElement": "\"vec4<\" ~ dtypes.T ~ \">\"",
-        "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\""
-      },
       "passes": [
         {
           "id": "main",
           "name": "GroupNormalization.group_subgroup_vec4",
-          "source": {
-            "shader": "norm-row-stats.wgsl.jinja",
-            "inputs": {
-              "mode": "\"group\"",
-              "vec4": true,
-              "scalar": "dtypes.T",
-              "usesF16": "dtypes.T == \"f16\"",
-              "hidden": "groupHidden",
-              "wg": "groupVec4Workgroup",
-              "epsilon": "attrs.epsilon",
-              "numGroups": "attrs.num_groups",
-              "cpg": "groupChannelsPerGroup",
-              "hiddenVec": "groupHidden / 4",
-              "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
-              "spatialVec": "groupSpatial / 4",
-              "combineSubgroups": "hasSubgroupId"
-            }
           },
-          "subgroupCollectivesWidth": "portable",
-          "bindings": "norm",
-          "dispatch": { "workgroups": "groupRows" }
         }
       ]
     },
@@ -244,30 +170,28 @@
       "id": "group_subgroup",
       "priority": 100,
       "when": ["groupRowCovered"],
-      "constants": { "scalar": "dtypes.T", "ioElement": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "GroupNormalization.group_subgroup",
-          "source": {
-            "shader": "norm-row-stats.wgsl.jinja",
-            "inputs": {
-              "mode": "\"group\"",
-              "vec4": false,
-              "scalar": "dtypes.T",
-              "usesF16": "dtypes.T == \"f16\"",
-              "hidden": "groupHidden",
-              "wg": "groupScalarWorkgroup",
-              "epsilon": "attrs.epsilon",
-              "numGroups": "attrs.num_groups",
-              "cpg": "groupChannelsPerGroup",
-              "spatial": "groupSpatial",
-              "combineSubgroups": "hasSubgroupId"
-            }
           },
-          "subgroupCollectivesWidth": "portable",
-          "bindings": "norm",
-          "dispatch": { "workgroups": "groupRows" }
         }
       ]
     }

   "domain": "ai.onnx",
   "name": "GroupNormalization",
   "sinceVersion": 21,
+  "inputs": {
+    "x": { "onnx": "X", "dtype": "T" },
+    "scale": { "dtype": "T", "rank": 1 },
+    "bias": { "dtype": "T", "rank": 1 }
   },
+  "outputs": { "y": { "onnx": "Y", "dtype": "T", "rank": "ranks.x", "shape": "shapes.x" } },
+  "attributes": { "epsilon": { "default": 0.00001 }, "stash_type": { "default": 1 }, "num_groups": {} },
   "attributeConstraints": { "num_groups": { "required": true }, "stash_type": { "values": [1, 10] } },
   "typeConstraints": { "T": ["float32", "float16"] },
   "tunables": {
+    "WORKGROUP_SIZE": { "default": 256 },
+    "MAX_STATS_SPLITS": { "default": 64 },
+    "STATS_VALUES_PER_SPLIT": { "default": 4096 },
+    "SPLIT_STATS_MIN_HIDDEN": { "default": 65536 },
+    "SPLIT_STATS_MAX_ROWS": { "default": 256 }
   },
   "derive": {
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
     "groupAttributesOk": "attrs.num_groups >= 1",
+    "groupShapeOk": "groupAttributesOk and f16Ok(dtypes.T) and ranks.x >= 3 and ranks.scale == 1 and ranks.bias == 1 and ranks.y == ranks.x and sameShape(shapes.y, shapes.x) and dim(shapes.scale, 0) == dim(shapes.x, 1) and dim(shapes.bias, 0) == dim(shapes.x, 1) and dim(shapes.x, 1) % attrs.num_groups == 0",
     "groupContractOk": "groupShapeOk and attrs.stash_type == onnxDtypeCode(\"float32\")",
     "groupStashF16Ok": "groupShapeOk and attrs.stash_type == onnxDtypeCode(\"float16\")",
+    "groupRows": "dim(shapes.x, 0) * attrs.num_groups if groupAttributesOk else 0",
+    "groupSpatial": "inner(shapes.x, 1)",
+    "groupChannelsPerGroup": "dim(shapes.x, 1) / attrs.num_groups if groupAttributesOk else 0",
     "groupHidden": "groupChannelsPerGroup * groupSpatial",
     "normDeviceWorkgroupCap": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
     "normWorkgroupCap": "max(1, pow2ceil(normDeviceWorkgroupCap + 1) / 2)",
     "hasSubgroupId": "device.features.has(\"subgroups\") and device.wgslLanguageFeatures.has(\"subgroup_id\")",
     "groupRowWorkgroupBytes": "normWorkgroupCap * 2 * 4",
     "groupRowCovered": "groupContractOk and groupRowWorkgroupBytes <= device.limits.maxComputeWorkgroupStorageSize",
+    "groupSplitCount": "min(tunables.MAX_STATS_SPLITS, min(device.limits.maxComputeWorkgroupsPerDimension, 65535), pow2ceil(ceilDiv(groupHidden, tunables.STATS_VALUES_PER_SPLIT)))",
     "groupPartialBytes": "groupRows * groupSplitCount * 2 * 4",
+    "groupSplitCovered": "groupRowCovered and groupRows <= tunables.SPLIT_STATS_MAX_ROWS and groupRows <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and groupHidden >= tunables.SPLIT_STATS_MIN_HIDDEN and groupPartialBytes <= device.limits.maxStorageBufferBindingSize and groupPartialBytes <= device.limits.maxBufferSize"
   },
+  "bindings": {
+    "x": { "buffer": "read-only-storage", "elementType": "$ioElement" },
+    "scale": { "buffer": "read-only-storage", "elementType": "$scalar" },
+    "bias": { "buffer": "read-only-storage", "elementType": "$scalar" },
+    "y": { "buffer": "storage", "elementType": "$ioElement" },
+    "params": {
+      "buffer": "uniform",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "groupRows" },
+        {
+          "name": "rowStride",
+          "type": "u32",
+          "value": "max(1, min(groupRows, min(device.limits.maxComputeWorkgroupsPerDimension, 65535)))"
         }
+      ]
+    },
+    "x_2": { "name": "x", "buffer": "read-only-storage", "elementType": "$scalar" },
+    "params_2": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [{ "name": "rows", "type": "u32", "value": "groupRows" }]
+    }
   },
   "variants": [
     {
       "id": "group_stash_f16_serial",
       "priority": 1000,
+      "when": ["groupStashF16Ok"],
+      "derive": {
         "scalar": "dtypes.T",
         "ioElement": "dtypes.T",
         "usesF16": "dtypes.T == \"f16\"",
           "id": "main",
           "name": "GroupNormalization.StashF16Serial",
           "shader": "group-normalization-stash-f16-serial.wgsl.jinja",
+          "bindings": ["x", "scale", "bias", "y", "params"],
+          "dispatch": { "x": "min(groupRows, 65535)", "y": "ceilDiv(groupRows, 65535)", "z": 1 }
         }
       ]
     },
       "id": "group_splitk",
       "priority": 120,
       "when": ["groupSplitCovered"],
+      "derive": {
         "scalar": "dtypes.T",
         "usesF16": "dtypes.T == \"f16\"",
         "hiddenSize": "groupHidden",
           "id": "partials",
           "name": "GroupNormalization.SplitKPartials",
           "shader": "group-normalization-splitk-partials.wgsl.jinja",
+          "bindings": ["x_2", { "name": "partials", "buffer": "storage", "elementType": "vec2<f32>" }, "params_2"],
+          "dispatch": {
+            "x": "min(groupRows, DISPATCH_FOLD_WIDTH)",
+            "y": "ceilDiv(groupRows, DISPATCH_FOLD_WIDTH)",
+            "z": "groupSplitCount"
+          }
         },
         {
           "id": "apply",
           "name": "GroupNormalization.SplitKApply",
           "shader": "group-normalization-splitk-apply.wgsl.jinja",
+          "bindings": [
+            "x_2",
+            "scale",
+            "bias",
+            { "name": "partials", "buffer": "read-only-storage", "elementType": "vec2<f32>" },
+            { "arg": "y", "elementType": "$scalar" },
+            "params_2"
+          ],
+          "dispatch": {
+            "x": "min(groupRows, DISPATCH_FOLD_WIDTH)",
+            "y": "ceilDiv(groupRows, DISPATCH_FOLD_WIDTH)",
+            "z": "groupSplitCount"
+          }
         }
       ]
     },
       "id": "group_subgroup_vec4",
       "priority": 110,
       "when": ["groupRowCovered", "groupSpatial % 4 == 0"],
+      "derive": { "scalar": "dtypes.T", "ioElement": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "GroupNormalization.group_subgroup_vec4",
+          "shader": "norm-row-stats.wgsl.jinja",
+          "derive": {
+            "modeSpec": "\"group\"",
+            "vec4": true,
+            "scalar": "dtypes.T",
+            "usesF16Spec": "dtypes.T == \"f16\"",
+            "hidden": "groupHidden",
+            "wg": "groupVec4Workgroup",
+            "epsilon": "attrs.epsilon",
+            "numGroupsSpec": "attrs.num_groups",
+            "cpg": "groupChannelsPerGroup",
+            "hiddenVec": "groupHidden / 4",
+            "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
+            "spatialVec": "groupSpatial / 4",
+            "combineSubgroups": "hasSubgroupId"
           },
+          "bindings": ["x", "scale", "bias", "y", "params"],
+          "dispatch": { "x": "min(groupRows, 65535)", "y": "ceilDiv(groupRows, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
       "id": "group_subgroup",
       "priority": 100,
       "when": ["groupRowCovered"],
+      "derive": { "scalar": "dtypes.T", "ioElement": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "GroupNormalization.group_subgroup",
+          "shader": "norm-row-stats.wgsl.jinja",
+          "derive": {
+            "modeSpec": "\"group\"",
+            "vec4": false,
+            "scalar": "dtypes.T",
+            "usesF16Spec": "dtypes.T == \"f16\"",
+            "hidden": "groupHidden",
+            "wg": "groupScalarWorkgroup",
+            "epsilon": "attrs.epsilon",
+            "numGroupsSpec": "attrs.num_groups",
+            "cpg": "groupChannelsPerGroup",
+            "spatial": "groupSpatial",
+            "combineSubgroups": "hasSubgroupId"
           },
+          "bindings": ["x", "scale", "bias", "y", "params"],
+          "dispatch": { "x": "min(groupRows, 65535)", "y": "ceilDiv(groupRows, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,21 +1,29 @@
 {
   "name": "ai.onnx.GroupNormalization",
-  "id": "_ai_onnx_groupnormalization_webgpu_da50174",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "R60LEBSM13qdMqFSIdwpekWs9JTpBlgud7FNGLzPDdg=",
-      "group-normalization-splitk-apply.wgsl.jinja": "7Q93kvTqmCixZN4Z7GpxYdlxwaxiAUMjH7AMKHq8pfc=",
-      "group-normalization-splitk-partials.wgsl.jinja": "6FNPr81N83OW3DWYMjJln/B5zUsKTdinV610BRzzRvs=",
       "group-normalization-stash-f16-serial.wgsl.jinja": "Wez9kqS+lzZASbm2BpWZSsHuNTT4rqsus0NiyhmGmvY=",
-      "manifest.json": "nWQpLNZpKQYFo8Pkd/yNT3xboHmv9HZJ5S28OKU86TE=",
-      "norm-row-stats.wgsl.jinja": "RCAuBcGKN1mg+npi5UTIqD4zDk0+B8DQj6u2TH+4IXU=",
-      "test.json": "KbF1vAUrTyrgdXUaq67pwqn41JjEAcHxnO2dImDHUiM="
     }
   },
-  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
-  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.GroupNormalization" }
 }

 {
   "name": "ai.onnx.GroupNormalization",
+  "id": "_ai_onnx_groupnormalization_webgpu_a85b638",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "gttctaT32ACO8eeDLTBubmz2KmY+JbRp40eLp4BLAiw=",
+      "group-normalization-splitk-apply.wgsl.jinja": "bJQ3aD6iCak7YjlH5yZZGPo8FIGY8CuCz3BwVnTnUPE=",
+      "group-normalization-splitk-partials.wgsl.jinja": "DVgljhomjpQ4XizLxIEh8NckAby7lIb6+plUaGh812A=",
       "group-normalization-stash-f16-serial.wgsl.jinja": "Wez9kqS+lzZASbm2BpWZSsHuNTT4rqsus0NiyhmGmvY=",
+      "manifest.json": "A56uOwydeaOGyMd1ulwyeOrEUd7G8uLRbxougyW6s8g=",
+      "norm-row-stats.wgsl.jinja": "CyRuHHc7bYmXEhtvfCxLRvjhidAMvicRA5nuJJESwxg=",
+      "test.json": "jfHQT3aDoszhxWXbi2YtXb9bqMfNRQyaM3HFdPdOWA0="
     }
   },
+  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.0",
+    "variants": {
+      "group_stash_f16_serial": ["group-normalization-stash-f16-serial.wgsl.jinja"],
+      "group_splitk": ["group-normalization-splitk-apply.wgsl.jinja", "group-normalization-splitk-partials.wgsl.jinja"],
+      "group_subgroup_vec4": ["norm-row-stats.wgsl.jinja"],
+      "group_subgroup": ["norm-row-stats.wgsl.jinja"]
+    }
+  }
 }

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

@@ -1,8 +1,18 @@
-{% if source.usesF16 %}
 enable f16;
 {% endif %}
-{% set combineSubgroups = source.combineSubgroups %}
-{% set scalarIo = source.scalarIo if source.scalarIo is defined else false %}
 {% set reduceThreadParameters = ", sg_lane: u32, sg_id: u32, num_sg: u32"
     if combineSubgroups else ", tid: u32" %}
 {% set reduceThreadArguments = ", sg_lane, sg_id, num_sg"
@@ -22,21 +32,64 @@ enable subgroups;
 //
 // Shifted moments avoid cancellation from a large common offset; normalize as
 // (x - mean) / sqrt(variance + EPSILON).
-const HIDDEN: u32 = {{ source.hidden }}u;
-{% if source.vec4 %}
-const HIDDEN_V: u32 = {{ source.hiddenVec }}u;
-{% endif %}
-const WG: u32 = {{ source.wg }}u;
-const EPSILON: f32 = {{ source.epsilon }};
-const NUM_GROUPS: u32 = {{ source.numGroups }}u;
-const CPG: u32 = {{ source.cpg }}u;
-{% if source.vec4 %}
-const SPATIAL_V: u32 = {{ source.spatialVec }}u;
 {% else %}
-const SPATIAL: u32 = {{ source.spatial }}u;
 {% endif %}
 {% if combineSubgroups %}
 var<workgroup> sg_partials: array<vec2<f32>, WG>;
@@ -95,29 +148,52 @@ fn main(
     return;
   }
   let tid = lid.x;
-{% if source.vec4 and not scalarIo %}
   let base = row * HIDDEN_V;
 {% else %}
   let base = row * HIDDEN;
 {% endif %}
-{% if source.vec4 %}
   let shift = f32(x[base].x);
 {% else %}
   let shift = f32(x[base]);
 {% endif %}
   var acc = vec2<f32>(0.0, 0.0);
-{% if source.vec4 %}
   for (var i = tid; i < HIDDEN_V; i = i + WG) {
     let v = vec4<f32>(x[base + i]);
     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);
   }
 {% else %}
   for (var i = tid; i < HIDDEN; i = i + WG) {
     let v = f32(x[base + i]);
     let d = v - shift;
     acc.x = acc.x + d;
     acc.y = acc.y + d * d;
@@ -132,21 +208,51 @@ fn main(
   let row_mean = shift + mean_d;
   let g_ch_base = (row % NUM_GROUPS) * CPG;
-{% if source.vec4 %}
   for (var i = tid; i < HIDDEN_V; i = i + WG) {
     let idx = base + i;
     let v = vec4<f32>(x[idx]);
     let ch = g_ch_base + i / SPATIAL_V;
     let normed = (v - vec4<f32>(row_mean)) / vec4<f32>(denom);
-    y[idx] = {{ source.vecType }}(normed * vec4<f32>(f32(scale[ch])) + vec4<f32>(f32(bias[ch])));
   }
 {% else %}
   for (var i = tid; i < HIDDEN; i = i + WG) {
     let idx = base + i;
     let v = f32(x[idx]);
     let ch = g_ch_base + i / SPATIAL;
     let normed = (v - row_mean) / denom;
-    y[idx] = {{ source.scalar }}(normed * f32(scale[ch]) + f32(bias[ch]));
   }
 {% endif %}
 }

+{% if usesF16Spec %}
 enable f16;
 {% endif %}
+{% set combineSubgroups = combineSubgroups %}
+{% set scalarIo = scalarIo if scalarIo is defined else false %}
+{% set packedBf16Embedding = packedBf16Embedding if packedBf16Embedding is defined else false %}
+{% set rmsChainNorm = rmsChainNorm if rmsChainNorm is defined else false %}
+{% set hiddenPairs = hiddenPairs | default(0) %}
+{% set numRows = numRows | default(0) %}
+{% set epsilon = epsilon | default("0.0") %}
+{% set epsilon2 = epsilon2 | default("0.0") %}
+{% set numGroupsSpec = numGroupsSpec | default(0) %}
+{% set cpg = cpg | default(0) %}
+{% set spatialVec = spatialVec | default(0) %}
+{% set spatial = spatial | default(0) %}
 {% set reduceThreadParameters = ", sg_lane: u32, sg_id: u32, num_sg: u32"
     if combineSubgroups else ", tid: u32" %}
 {% set reduceThreadArguments = ", sg_lane, sg_id, num_sg"
 //
 // Shifted moments avoid cancellation from a large common offset; normalize as
 // (x - mean) / sqrt(variance + EPSILON).
+const HIDDEN: u32 = {{ hidden }}u;
+{% if vec4 %}
+const HIDDEN_V: u32 = {{ hiddenVec }}u;
+{% endif %}
+{% if packedBf16Embedding %}
+const HIDDEN_PAIRS: u32 = {{ hiddenPairs }}u;
+const NUM_ROWS: u32 = {{ numRows }}u;
+{% endif %}
+const WG: u32 = {{ wg }}u;
+const EPSILON: f32 = {{ epsilon }};
+{% if rmsChainNorm %}
+const EPSILON2: f32 = {{ epsilon2 }};
+{% endif %}
+const NUM_GROUPS: u32 = {{ numGroupsSpec }}u;
+const CPG: u32 = {{ cpg }}u;
+{% if vec4 %}
+const SPATIAL_V: u32 = {{ spatialVec }}u;
 {% else %}
+const SPATIAL: u32 = {{ spatial }}u;
+{% endif %}
+{% if packedBf16Embedding %}
+{% if vec4 %}
+fn unpack_bf16_pair(word: u32) -> vec2<f32> {
+  let bits = vec2<u32>(word & 0xffffu, word >> 16u);
+  return bitcast<vec2<f32>>(bits << vec2<u32>(16u));
+}
+{% endif %}
+{% if not vec4 %}
+fn embedding_scalar(source_row: u32, hidden: u32) -> f32 {
+  if (source_row >= NUM_ROWS) {
+    return 0.0;
+  }
+  let word = x[source_row * HIDDEN_PAIRS + (hidden >> 1u)];
+  let bits = select(word & 0xffffu, word >> 16u, (hidden & 1u) != 0u);
+  return bitcast<f32>(bits << 16u);
+}
 {% endif %}
+{% if vec4 %}
+fn embedding_vec4(source_row: u32, hidden_vec: u32) -> vec4<f32> {
+  if (source_row >= NUM_ROWS) {
+    return vec4<f32>(0.0);
+  }
+  let base = source_row * HIDDEN_PAIRS + hidden_vec * 2u;
+  let low = unpack_bf16_pair(x[base]);
+  let high = unpack_bf16_pair(x[base + 1u]);
+  return vec4<f32>(low, high);
+}
+{% endif %}
+{% endif %}
+{% if vec4 and scalarIo %}
+fn load_vec4(index: u32) -> vec4<f32> {
+  return vec4<f32>(x[index], x[index + 1u], x[index + 2u], x[index + 3u]);
+}
+{% endif %}
 {% if combineSubgroups %}
 var<workgroup> sg_partials: array<vec2<f32>, WG>;
     return;
   }
   let tid = lid.x;
+{% if packedBf16Embedding %}
+  let source_row = indices[row];
+{% if vec4 %}
+  let base = row * HIDDEN_V;
+{% else %}
+  let base = row * HIDDEN;
+{% endif %}
+{% elif vec4 and not scalarIo %}
   let base = row * HIDDEN_V;
 {% else %}
   let base = row * HIDDEN;
 {% endif %}
+{% if vec4 %}
+{% if scalarIo %}
+  let shift = f32(x[base]);
+{% else %}
   let shift = f32(x[base].x);
+{% endif %}
 {% else %}
   let shift = f32(x[base]);
 {% endif %}
   var acc = vec2<f32>(0.0, 0.0);
+{% if vec4 %}
   for (var i = tid; i < HIDDEN_V; i = i + WG) {
+{% if packedBf16Embedding %}
+    let v = embedding_vec4(source_row, i);
+    embedding_out[base + i] = v;
+{% elif scalarIo %}
+    let v = load_vec4(base + i * 4u);
+{% else %}
     let v = vec4<f32>(x[base + i]);
+{% endif %}
     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);
   }
 {% else %}
   for (var i = tid; i < HIDDEN; i = i + WG) {
+{% if packedBf16Embedding %}
+    let v = embedding_scalar(source_row, i);
+    embedding_out[base + i] = v;
+{% else %}
     let v = f32(x[base + i]);
+{% endif %}
     let d = v - shift;
     acc.x = acc.x + d;
     acc.y = acc.y + d * d;
   let row_mean = shift + mean_d;
   let g_ch_base = (row % NUM_GROUPS) * CPG;
+{% if rmsChainNorm %}
+  var acc2 = 0.0;
+{% endif %}
+{% if vec4 %}
   for (var i = tid; i < HIDDEN_V; i = i + WG) {
+{% if packedBf16Embedding %}
+    let idx = base + i;
+    let v = embedding_vec4(source_row, i);
+{% elif scalarIo %}
+    let idx = base + i * 4u;
+    let v = load_vec4(idx);
+{% else %}
     let idx = base + i;
     let v = vec4<f32>(x[idx]);
+{% endif %}
     let ch = g_ch_base + i / SPATIAL_V;
     let normed = (v - vec4<f32>(row_mean)) / vec4<f32>(denom);
+    y[idx] = {{ vecType }}(normed * vec4<f32>(f32(scale[ch])) + vec4<f32>(f32(bias[ch])));
   }
+{% if rmsChainNorm %}
+  // The chained second norm reads the residual row this loop just stored. This
+  // barrier completes those stores and any preceding shared-scratch use before
+  // the next reduction reuses its scratch; each lane then re-reads only the
+  // elements it wrote itself.
+  workgroupBarrier();
+  let total2 = reduce_scalar(acc2{{ reduceThreadArguments }});
+  let inv2 = inverseSqrt(total2 / f32(HIDDEN) + EPSILON2);
+  for (var i = tid; i < HIDDEN_V; i = i + WG) {
+    let idx = base + i;
+    let hv = vec4<f32>(y[idx]);
+    normed2[idx] = {{ vecType }}(hv * inv2 * vec4<f32>(scale2[i]));
+  }
+{% endif %}
 {% else %}
   for (var i = tid; i < HIDDEN; i = i + WG) {
     let idx = base + i;
+{% if packedBf16Embedding %}
+    let v = embedding_scalar(source_row, i);
+{% else %}
     let v = f32(x[idx]);
+{% endif %}
     let ch = g_ch_base + i / SPATIAL;
     let normed = (v - row_mean) / denom;
+    y[idx] = {{ scalar }}(normed * f32(scale[ch]) + f32(bias[ch]));
   }
 {% endif %}
 }

build/webgpu/test.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "ai.onnx.GroupNormalization",
   "fixtureArrays": {
     "batch2_groups2_f32_input_x": [1, 2, 3, 4, 5, 6, 7, 8, -1, -2, -3, -4, 8, 7, 6, 5],
     "ort_rank3_equivalent_instance_norm_groups_equal_channels_input_x": [3.1513367, 9.283596, 1.4546119, 5.4617004, 8.519701, 1.2382338, 1.7930176, 5.1099434, 7.9195533, 7.638727, 8.065445, 3.8082376, 2.3667817, 2.8248506, 3.7754705, 5.861325, 5.058735, 3.2787242, 3.6843839, 9.755121, 2.7902672, 7.3974323, 8.283609, 8.488337],
@@ -107,7 +106,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/nn/group_norm_op_test.cc",
         "test": "GroupNormalizationOpTest.GroupSize_N",
-        "notes": "Larger grouped companion: subnormal scales should survive across channels and spatial positions."
       },
       "attrs": { "num_groups": 2, "epsilon": 0.00001 },
       "inputs": {
@@ -693,7 +692,7 @@
     {
       "name": "group_splitk_layernorm_equiv_65536_f16",
       "provenance": {
-        "notes": "float16 twin of the split-K group reduction. Only f32 cases reached the split-K partial and apply shaders, so neither had ever emitted its f16 declaration."
       },
       "attrs": { "num_groups": 1, "epsilon": 0.00001 },
       "inputs": {

 {
   "fixtureArrays": {
     "batch2_groups2_f32_input_x": [1, 2, 3, 4, 5, 6, 7, 8, -1, -2, -3, -4, 8, 7, 6, 5],
     "ort_rank3_equivalent_instance_norm_groups_equal_channels_input_x": [3.1513367, 9.283596, 1.4546119, 5.4617004, 8.519701, 1.2382338, 1.7930176, 5.1099434, 7.9195533, 7.638727, 8.065445, 3.8082376, 2.3667817, 2.8248506, 3.7754705, 5.861325, 5.058735, 3.2787242, 3.6843839, 9.755121, 2.7902672, 7.3974323, 8.283609, 8.488337],
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/nn/group_norm_op_test.cc",
         "test": "GroupNormalizationOpTest.GroupSize_N",
+        "notes": "A larger grouped input requires subnormal scales to survive across channels and spatial positions."
       },
       "attrs": { "num_groups": 2, "epsilon": 0.00001 },
       "inputs": {
     {
       "name": "group_splitk_layernorm_equiv_65536_f16",
       "provenance": {
+        "notes": "A float16 group of 65,536 elements exercises both split-K partial reduction and the float16 apply pass."
       },
       "attrs": { "num_groups": 1, "epsilon": 0.00001 },
       "inputs": {