sync 6fdf6301e2bb

Browse files

Files changed (9) hide show

README.md +2 -2
build/webgpu/manifest.json +17 -34
build/webgpu/metadata.json +10 -10
build/webgpu/norm-row-stats.wgsl.jinja +234 -139
build/webgpu/rms-normalization-splitk-normalize.wgsl.jinja +0 -1
build/webgpu/rms-normalization-splitk-partials.wgsl.jinja +5 -39
build/webgpu/rms-normalization-stash-f16-serial.wgsl.jinja +0 -4
build/webgpu/rms-normalization.wgsl.jinja +217 -43
build/webgpu/test.json +76 -2

README.md CHANGED Viewed

@@ -51,7 +51,7 @@ Default values (overridable per request):
 - [`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
 - [`norm-row-stats.wgsl.jinja`](build/webgpu/norm-row-stats.wgsl.jinja)
 - [`rms-normalization-splitk-normalize.wgsl.jinja`](build/webgpu/rms-normalization-splitk-normalize.wgsl.jinja)
 - [`rms-normalization-splitk-partials.wgsl.jinja`](build/webgpu/rms-normalization-splitk-partials.wgsl.jinja)
@@ -61,7 +61,7 @@ Default values (overridable per request):
 ## 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.

 - [`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 cases
 - [`norm-row-stats.wgsl.jinja`](build/webgpu/norm-row-stats.wgsl.jinja)
 - [`rms-normalization-splitk-normalize.wgsl.jinja`](build/webgpu/rms-normalization-splitk-normalize.wgsl.jinja)
 - [`rms-normalization-splitk-partials.wgsl.jinja`](build/webgpu/rms-normalization-splitk-partials.wgsl.jinja)
 ## Use with `@huggingface/kernels`
 ```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.3
 ```
 Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.

build/webgpu/manifest.json CHANGED Viewed

@@ -18,6 +18,7 @@
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
     "wave32Adapter": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 32 and device.adapterInfo.subgroupMaxSize == 32",
     "reportedNonWave32Adapter": "not wave32Adapter and (has(device.adapterInfo, \"subgroupMinSize\") or has(device.adapterInfo, \"subgroupMaxSize\"))",
     "normMaxWorkgroup": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
     "hasSubgroupId": "device.features.has(\"subgroups\") and device.wgslLanguageFeatures.has(\"subgroup_id\")",
     "axisNorm": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x",
@@ -32,11 +33,10 @@
     "suffixAxisOk": "baseOk and ranks.x >= 2 and not (attrs.axis == -1 or attrs.axis == ranks.x - 1)"
   },
   "bindings": {
-    "x": { "buffer": "read-only-storage", "elementType": "$xElement" },
-    "scale": { "buffer": "read-only-storage", "elementType": "$ioElement" },
-    "y": { "buffer": "storage", "elementType": "$ioElement" },
     "params": {
-      "buffer": "uniform",
       "struct": [
         { "name": "rows", "type": "u32", "value": "normalizedRows" },
         {
@@ -46,9 +46,8 @@
         }
       ]
     },
-    "params_2": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
         { "name": "rows", "type": "u32", "value": "splitRows" },
         {
@@ -68,7 +67,6 @@
         "scalar": "dtypes.V",
         "xElement": "dtypes.T",
         "ioElement": "dtypes.V",
-        "usesF16": "dtypes.T == \"f16\" or dtypes.V == \"f16\"",
         "hiddenSize": "normalizedHidden",
         "epsilon": "attrs.epsilon"
       },
@@ -95,8 +93,8 @@
     {
       "id": "suffix_axis_splitk",
       "priority": 15,
-      "when": ["baseOk", "ranks.x >= 2", "normalizedRows <= tunables.SPLIT_MAX_ROWS", "normalizedHidden >= tunables.SPLIT_MIN_HIDDEN", "min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(normalizedHidden, tunables.SPLIT_TARGET_ELEMENTS))) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "normalizedRows * min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(normalizedHidden, tunables.SPLIT_TARGET_ELEMENTS))) * 4 <= device.limits.maxStorageBufferBindingSize", "normalizedRows * min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(normalizedHidden, tunables.SPLIT_TARGET_ELEMENTS))) * 4 <= device.limits.maxBufferSize"],
-      "demoteWhen": ["reportedNonWave32Adapter"],
       "derive": {
         "splitRows": "normalizedRows",
         "splitHidden": "normalizedHidden",
@@ -104,7 +102,6 @@
         "scalar": "dtypes.V",
         "xElement": "dtypes.T",
         "ioElement": "dtypes.V",
-        "usesF16": "dtypes.T == \"f16\" or dtypes.V == \"f16\"",
         "hiddenSize": "splitHidden",
         "workgroupSize": "normMaxWorkgroup",
         "epsilon": "attrs.epsilon",
@@ -116,7 +113,7 @@
           "id": "partials",
           "name": "RMSNormalization.SplitKPartials",
           "shader": "rms-normalization-splitk-partials.wgsl.jinja",
-          "bindings": ["x", { "name": "partials", "buffer": "storage", "elementType": "f32" }, "params_2"],
           "dispatch": {
             "x": "min(splitRows, DISPATCH_FOLD_WIDTH)",
             "y": "ceilDiv(splitRows, DISPATCH_FOLD_WIDTH)",
@@ -132,8 +129,6 @@
             "scaleShape": "shapes.scale",
             "xRank": "ranks.x",
             "scaleRank": "ranks.scale",
-            "writeStats": false,
-            "rmsScaleAfterCast": true,
             "normalizeBlocks": "max(split, min(min(device.limits.maxComputeWorkgroupsPerDimension, 65535), ceilDiv(workgroupSize, max(1, normalizeRows)), ceilDiv(hiddenSize, workgroupSize * 4)))",
             "normalizeChunk": "ceilDiv(hiddenSize, normalizeBlocks)"
           },
@@ -142,7 +137,7 @@
             "scale",
             { "name": "partials", "buffer": "read-only-storage", "elementType": "f32" },
             "y",
-            "params_2"
           ],
           "dispatch": {
             "x": "min(splitRows, DISPATCH_FOLD_WIDTH)",
@@ -159,7 +154,6 @@
       "derive": {
         "scalar": "dtypes.V",
         "xElement": "dtypes.T",
-        "usesF16": "dtypes.T == \"f16\" or dtypes.V == \"f16\"",
         "ioElement": "dtypes.V",
         "hiddenSize": "normalizedHidden",
         "workgroupSize": "min(normMaxWorkgroup, pow2ceil(normalizedWorkgroupHidden))",
@@ -171,12 +165,11 @@
           "name": "RMSNormalization",
           "shader": "rms-normalization.wgsl.jinja",
           "derive": {
             "xShape": "shapes.x",
             "scaleShape": "shapes.scale",
             "xRank": "ranks.x",
-            "scaleRank": "ranks.scale",
-            "writeStats": false,
-            "rmsScaleAfterCast": true
           },
           "bindings": ["x", "scale", "y", "params"],
           "dispatch": {
@@ -194,7 +187,6 @@
       "derive": {
         "scalar": "dtypes.V",
         "xElement": "dtypes.T",
-        "usesF16": "dtypes.T == \"f16\" or dtypes.V == \"f16\"",
         "ioElement": "dtypes.V",
         "hiddenSize": "normalizedHidden",
         "workgroupSize": "min(normMaxWorkgroup, pow2ceil(normalizedWorkgroupHidden))",
@@ -206,12 +198,11 @@
           "name": "RMSNormalization.SuffixAxis",
           "shader": "rms-normalization.wgsl.jinja",
           "derive": {
             "xShape": "shapes.x",
             "scaleShape": "shapes.scale",
             "xRank": "ranks.x",
-            "scaleRank": "ranks.scale",
-            "writeStats": false,
-            "rmsScaleAfterCast": true
           },
           "bindings": ["x", "scale", "y", "params"],
           "dispatch": {
@@ -233,12 +224,9 @@
           "name": "RMSNormalization.LastAxisRow",
           "shader": "norm-row-stats.wgsl.jinja",
           "derive": {
-            "modeSpec": "\"rms\"",
             "vec4": true,
-            "writeStats": false,
-            "rmsScaleAfterCast": true,
             "scalar": "dtypes.T",
-            "usesF16Spec": "dtypes.T == \"f16\"",
             "hidden": "dim(shapes.x, -1)",
             "wg": "min(normMaxWorkgroup, pow2ceil(max(1, dim(shapes.x, -1) / 4)))",
             "epsilon": "attrs.epsilon",
@@ -251,8 +239,7 @@
             "x": "min(normalizedDispatchRows, 65535)",
             "y": "ceilDiv(normalizedDispatchRows, 65535)",
             "z": 1
-          },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
@@ -267,12 +254,9 @@
           "name": "RMSNormalization.LastAxisRow",
           "shader": "norm-row-stats.wgsl.jinja",
           "derive": {
-            "modeSpec": "\"rms\"",
             "vec4": false,
-            "writeStats": false,
-            "rmsScaleAfterCast": true,
             "scalar": "dtypes.T",
-            "usesF16Spec": "dtypes.T == \"f16\"",
             "hidden": "dim(shapes.x, -1)",
             "wg": "min(normMaxWorkgroup, pow2ceil(max(1, dim(shapes.x, -1))))",
             "epsilon": "attrs.epsilon",
@@ -285,8 +269,7 @@
             "x": "min(normalizedDispatchRows, 65535)",
             "y": "ceilDiv(normalizedDispatchRows, 65535)",
             "z": 1
-          },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     }

     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
     "wave32Adapter": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 32 and device.adapterInfo.subgroupMaxSize == 32",
     "reportedNonWave32Adapter": "not wave32Adapter and (has(device.adapterInfo, \"subgroupMinSize\") or has(device.adapterInfo, \"subgroupMaxSize\"))",
+    "variableSubgroup16To32": "device.features.has(\"subgroups\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 16 and device.adapterInfo.subgroupMaxSize == 32",
     "normMaxWorkgroup": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
     "hasSubgroupId": "device.features.has(\"subgroups\") and device.wgslLanguageFeatures.has(\"subgroup_id\")",
     "axisNorm": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x",
     "suffixAxisOk": "baseOk and ranks.x >= 2 and not (attrs.axis == -1 or attrs.axis == ranks.x - 1)"
   },
   "bindings": {
+    "x": { "elementType": "$xElement" },
+    "scale": { "elementType": "$ioElement" },
+    "y": { "elementType": "$ioElement" },
     "params": {
       "struct": [
         { "name": "rows", "type": "u32", "value": "normalizedRows" },
         {
         }
       ]
     },
+    "params__uniform": {
       "name": "params",
       "struct": [
         { "name": "rows", "type": "u32", "value": "splitRows" },
         {
         "scalar": "dtypes.V",
         "xElement": "dtypes.T",
         "ioElement": "dtypes.V",
         "hiddenSize": "normalizedHidden",
         "epsilon": "attrs.epsilon"
       },
     {
       "id": "suffix_axis_splitk",
       "priority": 15,
+      "when": ["baseOk", "ranks.x >= 2", "normalizedRows <= tunables.SPLIT_MAX_ROWS", "normalizedHidden >= tunables.SPLIT_MIN_HIDDEN", "min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(normalizedHidden, tunables.SPLIT_TARGET_ELEMENTS))) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "normalizedRows * min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(normalizedHidden, tunables.SPLIT_TARGET_ELEMENTS))) * 4 <= device.limits.maxStorageBufferBindingSize", "normalizedRows * min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(normalizedHidden, tunables.SPLIT_TARGET_ELEMENTS))) * 4 <= device.limits.maxBufferSize", "not (variableSubgroup16To32 and dtypes.T == \"f16\" and dtypes.V == \"f32\")"],
+      "demoteWhen": ["reportedNonWave32Adapter and not variableSubgroup16To32"],
       "derive": {
         "splitRows": "normalizedRows",
         "splitHidden": "normalizedHidden",
         "scalar": "dtypes.V",
         "xElement": "dtypes.T",
         "ioElement": "dtypes.V",
         "hiddenSize": "splitHidden",
         "workgroupSize": "normMaxWorkgroup",
         "epsilon": "attrs.epsilon",
           "id": "partials",
           "name": "RMSNormalization.SplitKPartials",
           "shader": "rms-normalization-splitk-partials.wgsl.jinja",
+          "bindings": ["x", { "name": "partials", "elementType": "f32" }, "params__uniform"],
           "dispatch": {
             "x": "min(splitRows, DISPATCH_FOLD_WIDTH)",
             "y": "ceilDiv(splitRows, DISPATCH_FOLD_WIDTH)",
             "scaleShape": "shapes.scale",
             "xRank": "ranks.x",
             "scaleRank": "ranks.scale",
             "normalizeBlocks": "max(split, min(min(device.limits.maxComputeWorkgroupsPerDimension, 65535), ceilDiv(workgroupSize, max(1, normalizeRows)), ceilDiv(hiddenSize, workgroupSize * 4)))",
             "normalizeChunk": "ceilDiv(hiddenSize, normalizeBlocks)"
           },
             "scale",
             { "name": "partials", "buffer": "read-only-storage", "elementType": "f32" },
             "y",
+            "params__uniform"
           ],
           "dispatch": {
             "x": "min(splitRows, DISPATCH_FOLD_WIDTH)",
       "derive": {
         "scalar": "dtypes.V",
         "xElement": "dtypes.T",
         "ioElement": "dtypes.V",
         "hiddenSize": "normalizedHidden",
         "workgroupSize": "min(normMaxWorkgroup, pow2ceil(normalizedWorkgroupHidden))",
           "name": "RMSNormalization",
           "shader": "rms-normalization.wgsl.jinja",
           "derive": {
+            "compensateHalfStats": true,
             "xShape": "shapes.x",
             "scaleShape": "shapes.scale",
             "xRank": "ranks.x",
+            "scaleRank": "ranks.scale"
           },
           "bindings": ["x", "scale", "y", "params"],
           "dispatch": {
       "derive": {
         "scalar": "dtypes.V",
         "xElement": "dtypes.T",
         "ioElement": "dtypes.V",
         "hiddenSize": "normalizedHidden",
         "workgroupSize": "min(normMaxWorkgroup, pow2ceil(normalizedWorkgroupHidden))",
           "name": "RMSNormalization.SuffixAxis",
           "shader": "rms-normalization.wgsl.jinja",
           "derive": {
+            "compensateHalfStats": true,
             "xShape": "shapes.x",
             "scaleShape": "shapes.scale",
             "xRank": "ranks.x",
+            "scaleRank": "ranks.scale"
           },
           "bindings": ["x", "scale", "y", "params"],
           "dispatch": {
           "name": "RMSNormalization.LastAxisRow",
           "shader": "norm-row-stats.wgsl.jinja",
           "derive": {
+            "compensateHalfStats": true,
             "vec4": true,
             "scalar": "dtypes.T",
             "hidden": "dim(shapes.x, -1)",
             "wg": "min(normMaxWorkgroup, pow2ceil(max(1, dim(shapes.x, -1) / 4)))",
             "epsilon": "attrs.epsilon",
             "x": "min(normalizedDispatchRows, 65535)",
             "y": "ceilDiv(normalizedDispatchRows, 65535)",
             "z": 1
+          }
         }
       ]
     },
           "name": "RMSNormalization.LastAxisRow",
           "shader": "norm-row-stats.wgsl.jinja",
           "derive": {
+            "compensateHalfStats": true,
             "vec4": false,
             "scalar": "dtypes.T",
             "hidden": "dim(shapes.x, -1)",
             "wg": "min(normMaxWorkgroup, pow2ceil(max(1, dim(shapes.x, -1))))",
             "epsilon": "attrs.epsilon",
             "x": "min(normalizedDispatchRows, 65535)",
             "y": "ceilDiv(normalizedDispatchRows, 65535)",
             "z": 1
+          }
         }
       ]
     }

build/webgpu/metadata.json CHANGED Viewed

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

 {
   "name": "ai.onnx.RMSNormalization",
+  "id": "_ai_onnx_rmsnormalization_webgpu_79e305f",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
     "algorithm": "sha256",
     "files": {
       "bench.json": "aQq6UMoOU5W6oOnS6tWV1CjBpS0e4r93laCsON1jvqk=",
+      "manifest.json": "yeH/EnaDGr2egsF8QVQs3KXw6ntjk/KVHH0h44Mdr58=",
+      "norm-row-stats.wgsl.jinja": "z9fPI5NfWmJ207kdnob9/md3Cc2yVtW3Uui89r8JKDs=",
+      "rms-normalization-splitk-normalize.wgsl.jinja": "LZOLQQd6JVX8rA5YOlcl3cw3cyaAvRYKh6h1lvVEiu4=",
+      "rms-normalization-splitk-partials.wgsl.jinja": "KnvI9/ZZe4zPBS+v4ieFHw7RTg/2c/MqVHH1Zj3T/F4=",
+      "rms-normalization-stash-f16-serial.wgsl.jinja": "7ZVEC2w1SegkxBz3V2T1DKhGKzmf7q32dqAfKXYr35Q=",
+      "rms-normalization.wgsl.jinja": "qpO/TCUizsL9TYtnpqNh1rwf3ZYy+Y5EboJ4j4jZqNc=",
+      "test.json": "cXXefTiwUMuJYGZSdIBjsMAyw0EO31uh2QXxrbZBNmU="
     }
   },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
   "webgpu": {
+    "manifestSpec": "2.1",
     "variants": {
       "stash_f16_serial": ["rms-normalization-stash-f16-serial.wgsl.jinja"],
       "suffix_axis_splitk": ["rms-normalization-splitk-normalize.wgsl.jinja", "rms-normalization-splitk-partials.wgsl.jinja"],

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

@@ -1,25 +1,7 @@
-{% if usesF16Spec %}
-enable f16;
-{% endif %}
-{% set combineSubgroups = combineSubgroups %}
-{% set scalarIo = scalarIo if scalarIo is defined else false %}
-{% set packedBf16Embedding = packedBf16Embedding if packedBf16Embedding is defined else false %}
-{% set writeStats = writeStats if writeStats is defined else false %}
-{% set rmsWeightOffset = rmsWeightOffset if rmsWeightOffset is defined else false %}
-{% set rmsScaleAfterCast = rmsScaleAfterCast if rmsScaleAfterCast is defined else false %}
-{% set rmsResidualAdd = rmsResidualAdd if rmsResidualAdd is defined else false %}
-{% set rmsChainNorm = rmsChainNorm if rmsChainNorm is defined else false %}
-{% set hiddenPairs = hiddenPairs | default(0) %}
-{% set numRows = numRows | default(0) %}
-{% set epsilon = epsilon | default("0.0") %}
-{% set epsilon2 = epsilon2 | default("0.0") %}
-{% if rmsWeightOffset %}
-{% set rmsScaleVec = "(vec4<f32>(1.0) + vec4<f32>(scale[i]))" %}
-{% set rmsScaleScalar = "(1.0 + f32(scale[i]))" %}
-{% else %}
 {% set rmsScaleVec = "vec4<f32>(scale[i])" %}
 {% set rmsScaleScalar = "f32(scale[i])" %}
-{% endif %}
 {% set reduceThreadParameters = ", sg_lane: u32, sg_id: u32, num_sg: u32"
     if combineSubgroups else ", tid: u32" %}
 {% set reduceThreadArguments = ", sg_lane, sg_id, num_sg"
@@ -28,67 +10,251 @@ enable f16;
 enable subgroups;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
-// Workgroup-parallel single-pass row statistics + fused normalize/affine.
-//
-// One workgroup owns one contiguous normalization span ("row": a last-axis
-// row, an instance plane, or a channel group). Threads stride the row once,
-// accumulating (sum, sum_sq) simultaneously. Partials are reduced either with
-// subgroupAdd plus a shared-memory combine or with a portable shared-memory
-// tree, then every thread applies the fused normalize + affine write.
-//
-// RMS mode uses sum_sq / HIDDEN without computing or subtracting a mean.
-const HIDDEN: u32 = {{ hidden }}u;
-{% if vec4 %}
-const HIDDEN_V: u32 = {{ hiddenVec }}u;
-{% endif %}
-{% if packedBf16Embedding %}
-const HIDDEN_PAIRS: u32 = {{ hiddenPairs }}u;
-const NUM_ROWS: u32 = {{ numRows }}u;
 {% endif %}
 const WG: u32 = {{ wg }}u;
 const EPSILON: f32 = {{ epsilon }};
-{% if rmsChainNorm %}
-const EPSILON2: f32 = {{ epsilon2 }};
-{% endif %}
-{% if packedBf16Embedding %}
-{% if vec4 %}
-fn unpack_bf16_pair(word: u32) -> vec2<f32> {
-  let bits = vec2<u32>(word & 0xffffu, word >> 16u);
-  return bitcast<vec2<f32>>(bits << vec2<u32>(16u));
 }
-{% endif %}
-{% if not vec4 %}
-fn embedding_scalar(source_row: u32, hidden: u32) -> f32 {
-  if (source_row >= NUM_ROWS) {
-    return 0.0;
   }
-  let word = x[source_row * HIDDEN_PAIRS + (hidden >> 1u)];
-  let bits = select(word & 0xffffu, word >> 16u, (hidden & 1u) != 0u);
-  return bitcast<f32>(bits << 16u);
 }
 {% endif %}
 {% if vec4 %}
-fn embedding_vec4(source_row: u32, hidden_vec: u32) -> vec4<f32> {
-  if (source_row >= NUM_ROWS) {
-    return vec4<f32>(0.0);
   }
-  let base = source_row * HIDDEN_PAIRS + hidden_vec * 2u;
-  let low = unpack_bf16_pair(x[base]);
-  let high = unpack_bf16_pair(x[base + 1u]);
-  return vec4<f32>(low, high);
-}
 {% endif %}
 {% endif %}
-{% if vec4 and scalarIo %}
-fn load_vec4(index: u32) -> vec4<f32> {
-  return vec4<f32>(x[index], x[index + 1u], x[index + 2u], x[index + 3u]);
 }
 {% endif %}
 {% if combineSubgroups %}
 var<workgroup> sg_partials: array<f32, WG>;
@@ -142,41 +308,21 @@ fn main(
     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 %}
   var acc = 0.0;
 {% if vec4 %}
   for (var i = tid; i < HIDDEN_V; i = i + WG) {
-{% if packedBf16Embedding %}
-    let v = embedding_vec4(source_row, i);
-    embedding_out[base + i] = v;
-{% elif scalarIo %}
-    let v = load_vec4(base + i * 4u);
-{% else %}
-    let v = vec4<f32>(x[base + i]);
-{% endif %}
     acc = acc + dot(v, v);
   }
 {% else %}
   for (var i = tid; i < HIDDEN; i = i + WG) {
-{% if packedBf16Embedding %}
-    let v = embedding_scalar(source_row, i);
-    embedding_out[base + i] = v;
-{% else %}
     let v = f32(x[base + i]);
-{% endif %}
     acc = acc + v * v;
   }
 {% endif %}
@@ -184,70 +330,19 @@ fn main(
   let total = reduce_scalar(acc{{ reduceThreadArguments }});
   let inv = inverseSqrt(total / f32(HIDDEN) + EPSILON);
-{% if writeStats %}
-  if (tid == 0u) {
-    inv_std_out[row] = inv;
-  }
-{% endif %}
-{% if rmsChainNorm %}
-  var acc2 = 0.0;
-{% endif %}
 {% if vec4 %}
   for (var i = tid; i < HIDDEN_V; i = i + WG) {
-{% if packedBf16Embedding %}
-    let idx = base + i;
-    let v = embedding_vec4(source_row, i);
-{% elif scalarIo %}
-    let idx = base + i * 4u;
-    let v = load_vec4(idx);
-{% else %}
     let idx = base + i;
-    let v = vec4<f32>(x[idx]);
-{% endif %}
-{% if rmsScaleAfterCast %}
     y[idx] = {{ vecType }}(v * inv) * {{ vecType }}({{ rmsScaleVec }});
-{% elif rmsChainNorm %}
-    // fma(a, b, 0.0) rounds the weighted product exactly as the decomposed pair's store does,
-    // and prevents the compiler from re-contracting it into the residual add.
-    let hv = y[idx] + fma(v * inv, {{ rmsScaleVec }}, vec4<f32>(0.0));
-    y[idx] = hv;
-    acc2 = acc2 + dot(hv, hv);
-{% elif rmsResidualAdd %}
-    // See the chained branch: fma(a, b, 0.0) pins the pre-add rounding of the decomposed pair.
-    y[idx] = y[idx] + fma(v * inv, {{ rmsScaleVec }}, vec4<f32>(0.0));
-{% else %}
-    y[idx] = {{ vecType }}(v * inv * {{ rmsScaleVec }});
-{% endif %}
   }
-{% if rmsChainNorm %}
-  // The chained second norm reads the residual row this loop just stored. This
-  // barrier completes those stores and any preceding shared-scratch use before
-  // the next reduction reuses its scratch; each lane then re-reads only the
-  // elements it wrote itself.
-  workgroupBarrier();
-  let total2 = reduce_scalar(acc2{{ reduceThreadArguments }});
-  let inv2 = inverseSqrt(total2 / f32(HIDDEN) + EPSILON2);
-  for (var i = tid; i < HIDDEN_V; i = i + WG) {
-    let idx = base + i;
-    let hv = vec4<f32>(y[idx]);
-    normed2[idx] = {{ vecType }}(hv * inv2 * vec4<f32>(scale2[i]));
-  }
-{% endif %}
 {% else %}
   for (var i = tid; i < HIDDEN; i = i + WG) {
     let idx = base + i;
-{% if packedBf16Embedding %}
-    let v = embedding_scalar(source_row, i);
-{% else %}
     let v = f32(x[idx]);
-{% endif %}
-{% if rmsScaleAfterCast %}
     y[idx] = {{ scalar }}(v * inv) * {{ scalar }}({{ rmsScaleScalar }});
-{% else %}
-    y[idx] = {{ scalar }}(v * inv * {{ rmsScaleScalar }});
-{% endif %}
   }
 {% endif %}
 }

+{% set scalarIo = false %}
+{% set packedF32 = "vec4<f32>" %}
 {% set rmsScaleVec = "vec4<f32>(scale[i])" %}
 {% set rmsScaleScalar = "f32(scale[i])" %}
 {% set reduceThreadParameters = ", sg_lane: u32, sg_id: u32, num_sg: u32"
     if combineSubgroups else ", tid: u32" %}
 {% set reduceThreadArguments = ", sg_lane, sg_id, num_sg"
 enable subgroups;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
+{% if scalar == "f16" and compensateHalfStats is defined and compensateHalfStats %}
+{% set halfOutputScalar = "f16" %}
+{% set halfStageVector = vec4 %}
+fn round_f16_bits_rte(value: f32) -> u32 {
+  let bits = bitcast<u32>(value);
+  let sign = (bits >> 16u) & 0x8000u;
+  let exponent_f32 = (bits >> 23u) & 0xffu;
+  let mantissa_f32 = bits & 0x7fffffu;
+  if (exponent_f32 == 0xffu) {
+    if (mantissa_f32 != 0u) {
+      return 0x7e00u;
+    }
+    return sign | 0x7c00u;
+  }
+  var exponent_f16 = i32(exponent_f32) - 127 + 15;
+  if (exponent_f16 >= 0x1f) {
+    return sign | 0x7c00u;
+  }
+  if (exponent_f16 <= 0) {
+    if (exponent_f16 < -10) {
+      return sign;
+    }
+    let significand = mantissa_f32 | 0x800000u;
+    let shift = u32(14 - exponent_f16);
+    let halfway = 1u << (shift - 1u);
+    let discarded = significand & ((1u << shift) - 1u);
+    var fraction = significand >> shift;
+    if (discarded > halfway || (discarded == halfway && (fraction & 1u) == 1u)) {
+      fraction = fraction + 1u;
+    }
+    return sign | fraction;
+  }
+  let halfway = 1u << 12u;
+  let discarded = mantissa_f32 & 0x1fffu;
+  var mantissa_f16 = mantissa_f32 >> 13u;
+  if (discarded > halfway || (discarded == halfway && (mantissa_f16 & 1u) == 1u)) {
+    mantissa_f16 = mantissa_f16 + 1u;
+    if (mantissa_f16 == 0x400u) {
+      mantissa_f16 = 0u;
+      exponent_f16 = exponent_f16 + 1;
+    }
+  }
+  if (exponent_f16 >= 0x1f) {
+    return sign | 0x7c00u;
+  }
+  return sign | (u32(exponent_f16) << 10u) | mantissa_f16;
+}
+fn widen_f16_bits(value: u32) -> f32 {
+  return unpack2x16float(value & 0xffffu).x;
+}
+// Typed ONNX edges must survive arithmetic fusion and narrow/wide casts.
+// Integer rounding also fixes the ties-to-even rule independently of the
+// implementation's floating-point conversion rounding mode.
+fn half_stage(value: f32) -> f32 {
+  return widen_f16_bits(round_f16_bits_rte(value));
+}
+{% if halfStageVector | default(false) %}
+fn half_stage4(value: vec4<f32>) -> vec4<f32> {
+  return vec4<f32>(half_stage(value.x), half_stage(value.y),
+                   half_stage(value.z), half_stage(value.w));
+}
 {% endif %}
+// Half output magnifies statistics errors at rounding midpoints. Keep a low
+// residual through the reduction and normalization, then round the typed
+// float32-normalized/half-scale/half-bias edges explicitly. This is the standard ONNX half path;
+// other normalization contracts retain their existing arithmetic.
+const HIDDEN: u32 = {{ hidden }}u;
 const WG: u32 = {{ wg }}u;
 const EPSILON: f32 = {{ epsilon }};
+fn pair_add(a: vec2<f32>, b: vec2<f32>) -> vec2<f32> {
+  let s = fma(a.x, 1.0, b.x);
+  // Materialize each rounded subtraction in the error-free transform. Plain
+  // cancellation expressions do not preserve the intended evaluation tree on
+  // every shader backend.
+  let bv = fma(-1.0, a.x, s);
+  let av = fma(-1.0, bv, s);
+  let a_error = fma(-1.0, av, a.x);
+  let b_error = fma(-1.0, bv, b.x);
+  let error = fma(a_error, 1.0, b_error);
+  let e = fma(fma(error, 1.0, a.y), 1.0, b.y);
+  let hi = fma(s, 1.0, e);
+  return vec2<f32>(hi, fma(-1.0, fma(-1.0, s, hi), e));
 }
+fn pair_mul(a: vec2<f32>, b: vec2<f32>) -> vec2<f32> {
+  let p = fma(a.x, b.x, 0.0);
+  let error = fma(a.x, b.x, -p);
+  return pair_add(vec2<f32>(p, 0.0), vec2<f32>(error + (a.x * b.y + b.x * a.y), 0.0));
+}
+fn pair_div(a: vec2<f32>, b: f32) -> vec2<f32> {
+  let q = a.x / b;
+  let residual = pair_add(a, -pair_mul(vec2<f32>(q, 0.0), vec2<f32>(b, 0.0)));
+  return pair_add(vec2<f32>(q, 0.0), vec2<f32>((residual.x + residual.y) / b, 0.0));
+}
+fn pair_inverse_sqrt(a: vec2<f32>) -> vec2<f32> {
+  let r = inverseSqrt(a.x);
+  let rr = pair_mul(vec2<f32>(r, 0.0), vec2<f32>(r, 0.0));
+  let residual = pair_add(vec2<f32>(1.0, 0.0), -pair_mul(a, rr));
+  return pair_add(vec2<f32>(r, 0.0), vec2<f32>((0.5 * r) * (residual.x + residual.y), 0.0));
+}
+fn half_normalized(value: vec2<f32>) -> f32 {
+  // stash_type=1 materializes Normalized as float32 before its cast to half.
+  // Collapse the compensated residual at that typed edge; rounding the pair
+  // directly to half can choose a different result at a float32 midpoint.
+  return widen_f16_bits(round_f16_bits_rte(fma(value.x, 1.0, value.y)));
+}
+{% if combineSubgroups %}
+var<workgroup> sg_partials: array<vec2<f32>, WG>;
+fn reduce_pair(value: vec2<f32>, sg_lane: u32, sg_id: u32, num_sg: u32, sg_size: u32) -> vec2<f32> {
+  var s = value;
+  let active_lanes = subgroupBallot(true);
+  let counts = countOneBits(active_lanes);
+  if (counts.x + counts.y + counts.z + counts.w == sg_size) {
+    for (var step = sg_size / 2u; step > 0u; step /= 2u) {
+      let other = subgroupShuffleDown(s, step);
+      if (sg_lane + step < sg_size) { s = pair_add(s, other); }
+    }
+    s = subgroupBroadcastFirst(s);
+  } else {
+    // Small workgroups need not fill a subgroup. Enumerate its actual active
+    // lanes instead of adding indeterminate shuffle results from inactive lanes.
+    s = vec2<f32>(0.0);
+    for (var word = 0u; word < 4u; word++) {
+      var mask = active_lanes[word];
+      while (mask != 0u) {
+        let lane = word * 32u + firstTrailingBit(mask);
+        s = pair_add(s, subgroupShuffle(value, lane));
+        mask &= mask - 1u;
+      }
+    }
   }
+  if (num_sg == 1u) { return s; }
+  if (subgroupElect()) { sg_partials[sg_id] = s; }
+  workgroupBarrier();
+  var total = vec2<f32>(0.0);
+  for (var i = 0u; i < num_sg; i++) { total = pair_add(total, sg_partials[i]); }
+  // Another reduction may immediately reuse the same storage.
+  workgroupBarrier();
+  return total;
+}
+{% else %}
+var<workgroup> partial: array<vec2<f32>, WG>;
+fn reduce_pair(value: vec2<f32>, tid: u32) -> vec2<f32> {
+  partial[tid] = value;
+  workgroupBarrier();
+  for (var step = WG / 2u; step > 0u; step /= 2u) {
+    if (tid < step) { partial[tid] = pair_add(partial[tid], partial[tid + step]); }
+    workgroupBarrier();
+  }
+  let total = partial[0];
+  workgroupBarrier();
+  return total;
 }
 {% endif %}
+{% set reduceArgs = "sg_lane, sg_id, num_sg, sg_size" if combineSubgroups else "tid" %}
+{% if not vec4 %}
+fn load_value(index: u32) -> f32 {
+  return f32(x[index]);
+}
+{% endif %}
+fn normalize_half_row(row: u32, tid: u32
+{% if combineSubgroups %}
+  , sg_lane: u32, sg_id: u32, num_sg: u32, sg_size: u32
+{% endif %}
+) {
+  if (row >= params.rows) { return; }
+  let base = row * HIDDEN;
+  let mean = vec2<f32>(0.0);
+  var local_square = vec2<f32>(0.0);
 {% if vec4 %}
+  for (var i = tid; i < HIDDEN / 4u; i += WG) {
+    let v = vec4<f32>(x[base / 4u + i]);
+{% for component in ["x", "y", "z", "w"] %}
+    let centered_{{ component }} = pair_add(vec2<f32>(v.{{ component }}, 0.0), -mean);
+    local_square = pair_add(local_square, pair_mul(centered_{{ component }}, centered_{{ component }}));
+{% endfor %}
+  }
+{% else %}
+  for (var i = tid; i < HIDDEN; i += WG) {
+    let centered = pair_add(vec2<f32>(load_value(base + i), 0.0), -mean);
+    local_square = pair_add(local_square, pair_mul(centered, centered));
   }
 {% endif %}
+  let variance = pair_div(reduce_pair(local_square, {{ reduceArgs }}), f32(HIDDEN));
+  let inv = pair_inverse_sqrt(pair_add(variance, vec2<f32>(EPSILON, 0.0)));
+{% if vec4 %}
+  for (var i = tid; i < HIDDEN / 4u; i += WG) {
+    let v = vec4<f32>(x[base / 4u + i]);
+    var normalized: vec4<f32>;
+    for (var component = 0u; component < 4u; component++) {
+      normalized[component] = half_normalized(pair_mul(pair_add(vec2<f32>(v[component], 0.0), -mean), inv));
+    }
+    var value = half_stage4(normalized * vec4<f32>(scale[i]));
+    y[base / 4u + i] = vec4<f16>(value);
+  }
+{% else %}
+  for (var i = tid; i < HIDDEN; i += WG) {
+    let normalized = half_normalized(pair_mul(pair_add(vec2<f32>(load_value(base + i), 0.0), -mean), inv));
+    var value = half_stage(normalized * f32(scale[{{ halfScaleOffset | default("i") }}]));
+    y[base + i] = {{ halfOutputScalar }}(value);
+  }
 {% endif %}
 }
+@compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(workgroup_id) wg_id: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>
+{% if combineSubgroups %}
+  , @builtin(subgroup_invocation_id) sg_lane: u32, @builtin(subgroup_id) sg_id: u32,
+  @builtin(num_subgroups) num_sg: u32, @builtin(subgroup_size) sg_size: u32
 {% endif %}
+) {
+  let row = wg_id.x + wg_id.y * params.rowStride;
+  normalize_half_row(row, lid.x{% if combineSubgroups %}, sg_lane, sg_id, num_sg, sg_size{% endif %});
+}
+{% else %}
+// Workgroup-parallel single-pass row statistics + fused normalize/affine.
+//
+// One workgroup owns one contiguous normalization span ("row": a last-axis
+// row, an instance plane, or a channel group). Threads stride the row once,
+// accumulating (sum, sum_sq) simultaneously. Partials are reduced either with
+// subgroupAdd plus a shared-memory combine or with a portable shared-memory
+// tree, then every thread applies the fused normalize + affine write.
+//
+// RMS mode uses sum_sq / HIDDEN without computing or subtracting a mean.
+const HIDDEN: u32 = {{ hidden }}u;
+{% if vec4 %}
+const HIDDEN_V: u32 = {{ hiddenVec }}u;
+{% endif %}
+const WG: u32 = {{ wg }}u;
+const EPSILON: f32 = {{ epsilon }};
 {% if combineSubgroups %}
 var<workgroup> sg_partials: array<f32, WG>;
     return;
   }
   let tid = lid.x;
+{% if vec4 and not scalarIo %}
   let base = row * HIDDEN_V;
 {% else %}
   let base = row * HIDDEN;
 {% endif %}
   var acc = 0.0;
 {% if vec4 %}
   for (var i = tid; i < HIDDEN_V; i = i + WG) {
+    let v = {{ packedF32 }}(x[base + i]);
     acc = acc + dot(v, v);
   }
 {% else %}
   for (var i = tid; i < HIDDEN; i = i + WG) {
     let v = f32(x[base + i]);
     acc = acc + v * v;
   }
 {% endif %}
   let total = reduce_scalar(acc{{ reduceThreadArguments }});
   let inv = inverseSqrt(total / f32(HIDDEN) + EPSILON);
 {% if vec4 %}
   for (var i = tid; i < HIDDEN_V; i = i + WG) {
     let idx = base + i;
+    let v = {{ packedF32 }}(x[idx]);
     y[idx] = {{ vecType }}(v * inv) * {{ vecType }}({{ rmsScaleVec }});
   }
 {% else %}
   for (var i = tid; i < HIDDEN; i = i + WG) {
     let idx = base + i;
     let v = f32(x[idx]);
     y[idx] = {{ scalar }}(v * inv) * {{ scalar }}({{ rmsScaleScalar }});
   }
 {% endif %}
 }
+{% endif %}

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

@@ -54,7 +54,6 @@ fn scale_offset({% if scaleRank > 0 %}out_index: u32{% endif %}) -> u32 {
 {% endif %}
 }
 var<workgroup> shared_inv: f32;
 @compute @workgroup_size(WG, 1, 1)

 {% endif %}
 }
 var<workgroup> shared_inv: f32;
 @compute @workgroup_size(WG, 1, 1)

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

@@ -1,49 +1,15 @@
 {% macro wgsl_tree_fold_stmt(a, op, idx, svar) %}
-{% if op == "max" %}
-{{ a }}[{{ idx }}] = max({{ a }}[{{ idx }}], {{ a }}[{{ idx }} + {{ svar }}]);
-{%- else %}
-{{ a }}[{{ idx }}] = {{ a }}[{{ idx }}] + {{ a }}[{{ idx }} + {{ svar }}];
-{%- endif %}
-{% endmacro %}
-{% macro wgsl_tree_fold(arrays, op="add", idx="lid", wg="WORKGROUP_SIZE", svar="stride", typed=false, form="tail", breakInline=false, bodyInline=false, barrierFirst=false) %}
   var {{ svar }}{{ ": u32 " if typed else " " }}= {{ wg }} / 2u;
   loop {
-{% if form == "head" %}
-{% if breakInline %}
     if ({{ svar }} == 0u) { break; }
-{% else %}
-    if ({{ svar }} == 0u) {
-      break;
-    }
-{% endif %}
-{% endif %}
-{% if bodyInline %}
     if ({{ idx }} < {{ svar }}) { {{ wgsl_tree_fold_stmt(arrays[0], op, idx, svar) }} }
-{% else %}
-    if ({{ idx }} < {{ svar }}) {
-{% for a in arrays %}
-      {{ wgsl_tree_fold_stmt(a, op, idx, svar) }}
-{% endfor %}
-    }
-{% endif %}
-{% if form == "head" %}
-{% if barrierFirst %}
-    workgroupBarrier();
-    {{ svar }} = {{ svar }} / 2u;
-{% else %}
     {{ svar }} = {{ svar }} / 2u;
     workgroupBarrier();
-{% endif %}
-{% else %}
-    workgroupBarrier();
-    if ({{ svar }} == 1u) {
-      break;
-    }
-    {{ svar }} = {{ svar }} / 2u;
-{% endif %}
-  }
-{%- endmacro %}
 /* Split-K partial sum-of-squares for tensors with few rows and a large hidden
    dimension. A workgroup-per-row kernel exposes too little parallelism in this
    regime, so this pass splits each row across SPLIT workgroups

 {% macro wgsl_tree_fold_stmt(a, op, idx, svar) %}
+{% if op == "max" or op == "min" %}
+{{ a }}[{{ idx }}] = {{ op }}({{ a }}[{{ idx }}], {{ a }}[{{ idx }} + {{ svar }}]);{% else %}
+{{ a }}[{{ idx }}] = {{ a }}[{{ idx }}] {{ "*" if op == "prod" else "+" }} {{ a }}[{{ idx }} + {{ svar }}];{% endif %}{% endmacro %}
+{% macro wgsl_tree_fold(arrays, op="add", idx="lid", wg="WORKGROUP_SIZE", svar="stride", typed=false, form="tail", breakInline=false, bodyInline=false, barrierFirst=false, reuse=false) %}
   var {{ svar }}{{ ": u32 " if typed else " " }}= {{ wg }} / 2u;
   loop {
     if ({{ svar }} == 0u) { break; }
     if ({{ idx }} < {{ svar }}) { {{ wgsl_tree_fold_stmt(arrays[0], op, idx, svar) }} }
     {{ svar }} = {{ svar }} / 2u;
     workgroupBarrier();
+  }{% endmacro %}
 /* Split-K partial sum-of-squares for tensors with few rows and a large hidden
    dimension. A workgroup-per-row kernel exposes too little parallelism in this
    regime, so this pass splits each row across SPLIT workgroups

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

@@ -1,5 +1,3 @@
-{% if usesF16 %}enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 const HIDDEN: u32 = {{ hiddenSize }}u;
@@ -58,7 +56,6 @@ fn widen_f16_bits(value: u32) -> f32 {
   return unpack2x16float(value & 0xffffu).x;
 }
 {% if scaleRank > 0 %}
 const X_RANK: u32 = {{ xRank }}u;
 const SCALE_RANK: u32 = {{ scaleRank }}u;
@@ -104,7 +101,6 @@ fn scale_offset({% if scaleRank > 0 %}out_index: u32{% endif %}) -> u32 {
 {% endif %}
 }
 // ONNX RMSNormalization-23 expresses the stash_type=FLOAT16 stage as a graph
 // of f16 tensor operators. Keeping one normalization row on one invocation
 // preserves f16 accumulation and every stage boundary.

 {{ env.wgsl.resourceDeclarations }}
 const HIDDEN: u32 = {{ hiddenSize }}u;
   return unpack2x16float(value & 0xffffu).x;
 }
 {% if scaleRank > 0 %}
 const X_RANK: u32 = {{ xRank }}u;
 const SCALE_RANK: u32 = {{ scaleRank }}u;
 {% endif %}
 }
 // ONNX RMSNormalization-23 expresses the stash_type=FLOAT16 stage as a graph
 // of f16 tensor operators. Keeping one normalization row on one invocation
 // preserves f16 accumulation and every stage boundary.

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

@@ -4,8 +4,6 @@ const HIDDEN: u32 = {{ hiddenSize }}u;
 const EPSILON: f32 = {{ epsilon }};
 const WG: u32 = {{ workgroupSize }}u;
-var<workgroup> partial: array<f32, WG>;
 {% if scaleRank > 0 %}
 const X_RANK: u32 = {{ xRank }}u;
 const SCALE_RANK: u32 = {{ scaleRank }}u;
@@ -51,70 +49,245 @@ fn scale_offset({% if scaleRank > 0 %}out_index: u32{% endif %}) -> u32 {
 {% endif %}
 }
 {% macro wgsl_tree_fold_stmt(a, op, idx, svar) %}
-{% if op == "max" %}
-{{ a }}[{{ idx }}] = max({{ a }}[{{ idx }}], {{ a }}[{{ idx }} + {{ svar }}]);
-{%- else %}
-{{ a }}[{{ idx }}] = {{ a }}[{{ idx }}] + {{ a }}[{{ idx }} + {{ svar }}];
-{%- endif %}
-{% endmacro %}
-{% macro wgsl_tree_fold(arrays, op="add", idx="lid", wg="WORKGROUP_SIZE", svar="stride", typed=false, form="tail", breakInline=false, bodyInline=false, barrierFirst=false) %}
   var {{ svar }}{{ ": u32 " if typed else " " }}= {{ wg }} / 2u;
   loop {
-{% if form == "head" %}
-{% if breakInline %}
-    if ({{ svar }} == 0u) { break; }
-{% else %}
     if ({{ svar }} == 0u) {
       break;
     }
-{% endif %}
-{% endif %}
-{% if bodyInline %}
-    if ({{ idx }} < {{ svar }}) { {{ wgsl_tree_fold_stmt(arrays[0], op, idx, svar) }} }
-{% else %}
     if ({{ idx }} < {{ svar }}) {
 {% for a in arrays %}
       {{ wgsl_tree_fold_stmt(a, op, idx, svar) }}
 {% endfor %}
     }
-{% endif %}
-{% if form == "head" %}
-{% if barrierFirst %}
-    workgroupBarrier();
     {{ svar }} = {{ svar }} / 2u;
-{% else %}
-    {{ svar }} = {{ svar }} / 2u;
-    workgroupBarrier();
-{% endif %}
-{% else %}
     workgroupBarrier();
-    if ({{ svar }} == 1u) {
-      break;
-    }
-    {{ svar }} = {{ svar }} / 2u;
-{% endif %}
-  }
-{%- endmacro %}
 // Reusing partial after this reduction requires a barrier between the read of
 // partial[0] and the next write, or the next round can race the prior readers.
-{% set trailingBarrier = trailingBarrier is defined and trailingBarrier %}
 fn reduce_sum(value: f32, tid: u32) -> f32 {
   partial[tid] = value;
   workgroupBarrier();
 {{ wgsl_tree_fold(["partial"], idx="tid", wg="WG", form="head") }}
-{% if trailingBarrier %}
-  let total = partial[0];
-  workgroupBarrier();
-  return total;
-{% else %}
   return partial[0];
-{% endif %}
 }
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
   let row = wg.x + wg.y * params.rowStride;
@@ -139,3 +312,4 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
     y[base + d] = {{ scalar }}(value);
   }
 }

 const EPSILON: f32 = {{ epsilon }};
 const WG: u32 = {{ workgroupSize }}u;
 {% if scaleRank > 0 %}
 const X_RANK: u32 = {{ xRank }}u;
 const SCALE_RANK: u32 = {{ scaleRank }}u;
 {% endif %}
 }
+{% if compensateHalfStats is defined and compensateHalfStats and xElement == "f16" %}
+{% set halfOutputScalar = scalar %}
+{% set halfScaleOffset = "scale_offset(base + i)" if scaleRank > 0 else "scale_offset()" %}
+{% set halfOutputScalar = halfOutputScalar | default("f16") %}
+{% if halfOutputScalar == "f16" %}
+fn round_f16_bits_rte(value: f32) -> u32 {
+  let bits = bitcast<u32>(value);
+  let sign = (bits >> 16u) & 0x8000u;
+  let exponent_f32 = (bits >> 23u) & 0xffu;
+  let mantissa_f32 = bits & 0x7fffffu;
+  if (exponent_f32 == 0xffu) {
+    if (mantissa_f32 != 0u) {
+      return 0x7e00u;
+    }
+    return sign | 0x7c00u;
+  }
+  var exponent_f16 = i32(exponent_f32) - 127 + 15;
+  if (exponent_f16 >= 0x1f) {
+    return sign | 0x7c00u;
+  }
+  if (exponent_f16 <= 0) {
+    if (exponent_f16 < -10) {
+      return sign;
+    }
+    let significand = mantissa_f32 | 0x800000u;
+    let shift = u32(14 - exponent_f16);
+    let halfway = 1u << (shift - 1u);
+    let discarded = significand & ((1u << shift) - 1u);
+    var fraction = significand >> shift;
+    if (discarded > halfway || (discarded == halfway && (fraction & 1u) == 1u)) {
+      fraction = fraction + 1u;
+    }
+    return sign | fraction;
+  }
+  let halfway = 1u << 12u;
+  let discarded = mantissa_f32 & 0x1fffu;
+  var mantissa_f16 = mantissa_f32 >> 13u;
+  if (discarded > halfway || (discarded == halfway && (mantissa_f16 & 1u) == 1u)) {
+    mantissa_f16 = mantissa_f16 + 1u;
+    if (mantissa_f16 == 0x400u) {
+      mantissa_f16 = 0u;
+      exponent_f16 = exponent_f16 + 1;
+    }
+  }
+  if (exponent_f16 >= 0x1f) {
+    return sign | 0x7c00u;
+  }
+  return sign | (u32(exponent_f16) << 10u) | mantissa_f16;
+}
+fn widen_f16_bits(value: u32) -> f32 {
+  return unpack2x16float(value & 0xffffu).x;
+}
+// Typed ONNX edges must survive arithmetic fusion and narrow/wide casts.
+// Integer rounding also fixes the ties-to-even rule independently of the
+// implementation's floating-point conversion rounding mode.
+fn half_stage(value: f32) -> f32 {
+  return widen_f16_bits(round_f16_bits_rte(value));
+}
+{% else %}
+fn round_f16_bits_rte(value: f32) -> u32 {
+  let bits = bitcast<u32>(value);
+  let sign = (bits >> 16u) & 0x8000u;
+  let exponent_f32 = (bits >> 23u) & 0xffu;
+  let mantissa_f32 = bits & 0x7fffffu;
+  if (exponent_f32 == 0xffu) {
+    if (mantissa_f32 != 0u) {
+      return 0x7e00u;
+    }
+    return sign | 0x7c00u;
+  }
+  var exponent_f16 = i32(exponent_f32) - 127 + 15;
+  if (exponent_f16 >= 0x1f) {
+    return sign | 0x7c00u;
+  }
+  if (exponent_f16 <= 0) {
+    if (exponent_f16 < -10) {
+      return sign;
+    }
+    let significand = mantissa_f32 | 0x800000u;
+    let shift = u32(14 - exponent_f16);
+    let halfway = 1u << (shift - 1u);
+    let discarded = significand & ((1u << shift) - 1u);
+    var fraction = significand >> shift;
+    if (discarded > halfway || (discarded == halfway && (fraction & 1u) == 1u)) {
+      fraction = fraction + 1u;
+    }
+    return sign | fraction;
+  }
+  let halfway = 1u << 12u;
+  let discarded = mantissa_f32 & 0x1fffu;
+  var mantissa_f16 = mantissa_f32 >> 13u;
+  if (discarded > halfway || (discarded == halfway && (mantissa_f16 & 1u) == 1u)) {
+    mantissa_f16 = mantissa_f16 + 1u;
+    if (mantissa_f16 == 0x400u) {
+      mantissa_f16 = 0u;
+      exponent_f16 = exponent_f16 + 1;
+    }
+  }
+  if (exponent_f16 >= 0x1f) {
+    return sign | 0x7c00u;
+  }
+  return sign | (u32(exponent_f16) << 10u) | mantissa_f16;
+}
+fn widen_f16_bits(value: u32) -> f32 {
+  return unpack2x16float(value & 0xffffu).x;
+}
+{% endif %}
+// Half output magnifies statistics errors at rounding midpoints. Keep a low
+// residual through the reduction and normalization, then round the typed
+// float32-normalized/half-scale/half-bias edges explicitly. This is the standard ONNX half path;
+// other normalization contracts retain their existing arithmetic.
+fn pair_add(a: vec2<f32>, b: vec2<f32>) -> vec2<f32> {
+  let s = fma(a.x, 1.0, b.x);
+  // Materialize each rounded subtraction in the error-free transform. Plain
+  // cancellation expressions do not preserve the intended evaluation tree on
+  // every shader backend.
+  let bv = fma(-1.0, a.x, s);
+  let av = fma(-1.0, bv, s);
+  let a_error = fma(-1.0, av, a.x);
+  let b_error = fma(-1.0, bv, b.x);
+  let error = fma(a_error, 1.0, b_error);
+  let e = fma(fma(error, 1.0, a.y), 1.0, b.y);
+  let hi = fma(s, 1.0, e);
+  return vec2<f32>(hi, fma(-1.0, fma(-1.0, s, hi), e));
+}
+fn pair_mul(a: vec2<f32>, b: vec2<f32>) -> vec2<f32> {
+  let p = fma(a.x, b.x, 0.0);
+  let error = fma(a.x, b.x, -p);
+  return pair_add(vec2<f32>(p, 0.0), vec2<f32>(error + (a.x * b.y + b.x * a.y), 0.0));
+}
+fn pair_div(a: vec2<f32>, b: f32) -> vec2<f32> {
+  let q = a.x / b;
+  let residual = pair_add(a, -pair_mul(vec2<f32>(q, 0.0), vec2<f32>(b, 0.0)));
+  return pair_add(vec2<f32>(q, 0.0), vec2<f32>((residual.x + residual.y) / b, 0.0));
+}
+fn pair_inverse_sqrt(a: vec2<f32>) -> vec2<f32> {
+  let r = inverseSqrt(a.x);
+  let rr = pair_mul(vec2<f32>(r, 0.0), vec2<f32>(r, 0.0));
+  let residual = pair_add(vec2<f32>(1.0, 0.0), -pair_mul(a, rr));
+  return pair_add(vec2<f32>(r, 0.0), vec2<f32>((0.5 * r) * (residual.x + residual.y), 0.0));
+}
+fn half_normalized(value: vec2<f32>) -> f32 {
+  // stash_type=1 materializes Normalized as float32 before its cast to half.
+  // Collapse the compensated residual at that typed edge; rounding the pair
+  // directly to half can choose a different result at a float32 midpoint.
+  return widen_f16_bits(round_f16_bits_rte(fma(value.x, 1.0, value.y)));
+}
+var<workgroup> partial: array<vec2<f32>, WG>;
+fn reduce_pair(value: vec2<f32>, tid: u32) -> vec2<f32> {
+  partial[tid] = value;
+  workgroupBarrier();
+  for (var step = WG / 2u; step > 0u; step /= 2u) {
+    if (tid < step) { partial[tid] = pair_add(partial[tid], partial[tid + step]); }
+    workgroupBarrier();
+  }
+  let total = partial[0];
+  workgroupBarrier();
+  return total;
+}
+{% set reduceArgs = "tid" %}
+fn load_value(index: u32) -> f32 {
+  return f32(x[index]);
+}
+fn normalize_half_row(row: u32, tid: u32
+) {
+  if (row >= params.rows) { return; }
+  let base = row * HIDDEN;
+  let mean = vec2<f32>(0.0);
+  var local_square = vec2<f32>(0.0);
+  for (var i = tid; i < HIDDEN; i += WG) {
+    let centered = pair_add(vec2<f32>(load_value(base + i), 0.0), -mean);
+    local_square = pair_add(local_square, pair_mul(centered, centered));
+  }
+  let variance = pair_div(reduce_pair(local_square, {{ reduceArgs }}), f32(HIDDEN));
+  let inv = pair_inverse_sqrt(pair_add(variance, vec2<f32>(EPSILON, 0.0)));
+  for (var i = tid; i < HIDDEN; i += WG) {
+    let normalized = half_normalized(pair_mul(pair_add(vec2<f32>(load_value(base + i), 0.0), -mean), inv));
+{% if halfOutputScalar == "f16" %}
+    var value = half_stage(normalized * f32(scale[{{ halfScaleOffset | default("i") }}]));
+{% else %}
+    var value = normalized * f32(scale[{{ halfScaleOffset | default("i") }}]);
+{% endif %}
+    y[base + i] = {{ halfOutputScalar }}(value);
+  }
+}
+@compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
+  let row = wg.x + wg.y * params.rowStride;
+  normalize_half_row(row, lid.x);
+}
+{% else %}
+var<workgroup> partial: array<f32, WG>;
 {% macro wgsl_tree_fold_stmt(a, op, idx, svar) %}
+{% if op == "max" or op == "min" %}
+{{ a }}[{{ idx }}] = {{ op }}({{ a }}[{{ idx }}], {{ a }}[{{ idx }} + {{ svar }}]);{% else %}
+{{ a }}[{{ idx }}] = {{ a }}[{{ idx }}] {{ "*" if op == "prod" else "+" }} {{ a }}[{{ idx }} + {{ svar }}];{% endif %}{% endmacro %}
+{% macro wgsl_tree_fold(arrays, op="add", idx="lid", wg="WORKGROUP_SIZE", svar="stride", typed=false, form="tail", breakInline=false, bodyInline=false, barrierFirst=false, reuse=false) %}
   var {{ svar }}{{ ": u32 " if typed else " " }}= {{ wg }} / 2u;
   loop {
     if ({{ svar }} == 0u) {
       break;
     }
     if ({{ idx }} < {{ svar }}) {
 {% for a in arrays %}
       {{ wgsl_tree_fold_stmt(a, op, idx, svar) }}
 {% endfor %}
     }
     {{ svar }} = {{ svar }} / 2u;
     workgroupBarrier();
+  }{% endmacro %}
 // Reusing partial after this reduction requires a barrier between the read of
 // partial[0] and the next write, or the next round can race the prior readers.
 fn reduce_sum(value: f32, tid: u32) -> f32 {
   partial[tid] = value;
   workgroupBarrier();
 {{ wgsl_tree_fold(["partial"], idx="tid", wg="WG", form="head") }}
   return partial[0];
 }
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
   let row = wg.x + wg.y * params.rowStride;
     y[base + d] = {{ scalar }}(value);
   }
 }
+{% endif %}

build/webgpu/test.json CHANGED Viewed

@@ -9,6 +9,80 @@
     "ort_f16_axis1_outer_inner_broadcast_scale_input_x": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23]
   },
   "cases": [
     {
       "name": "subgroup_vec4_2x512",
       "attrs": { "epsilon": 0.000001, "axis": -1 },
@@ -1391,7 +1465,7 @@
     {
       "name": "onnx23_f16_cast_before_scale_row_vec4_exact",
       "provenance": {
-        "notes": "ONNX RMSNormalization-23 stage one casts Normalized back to X's dtype before Scale. This vector-scale case selects the optimized vec4 row kernel and differs by one f16 ULP if the cast is deferred until after Scale."
       },
       "requires": { "features": ["shader-f16"] },
       "attrs": { "epsilon": 0.00001, "axis": -1 },
@@ -1449,7 +1523,7 @@
     {
       "name": "onnx23_f16_cast_before_scale_splitk_exact",
       "provenance": {
-        "notes": "Forces the split-K shared kernel on the same exact f16 boundary case, pinning the standardized cast of Normalized back to X before Scale."
       },
       "requires": { "features": ["shader-f16"] },
       "tunables": { "SPLIT_MIN_HIDDEN": 1, "SPLIT_TARGET_ELEMENTS": 1 },

     "ort_f16_axis1_outer_inner_broadcast_scale_input_x": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23]
   },
   "cases": [
+    {
+      "name": "half_normalized_midpoint_scalar",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [1, 7],
+          "data": {
+            "kind": "values",
+            "values": [0.333984375, -0.82177734375, 0.35009765625, -1.189453125, -0.8701171875, 0.043914794921875, -0.35595703125]
+          }
+        },
+        "scale": {
+          "dtype": "float16",
+          "shape": [7],
+          "data": {
+            "kind": "values",
+            "values": [1.1455078125, 1.0283203125, 0.92626953125, 1.0771484375, 1.232421875, 0.89990234375, 1.060546875]
+          }
+        }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "float16",
+          "shape": [1, 7],
+          "data": {
+            "kind": "values",
+            "values": [0.56494140625, -1.248046875, 0.47900390625, -1.892578125, -1.583984375, 0.058319091796875, -0.5576171875]
+          },
+          "tolerance": 0.000002
+        }
+      },
+      "provenance": {
+        "source": "ONNX RMSNormalization typed normalization and scale stages",
+        "notes": "Float32 Normalized midpoint followed by cast to half and affine half rounding. The CPU reference and independent PyTorch float32 typed stages agree; rounding float64 Normalized directly to half gives a different result."
+      }
+    },
+    {
+      "name": "half_normalized_midpoint_vec4",
+      "attrs": { "axis": -1, "epsilon": 0.000001 },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [1, 28],
+          "data": {
+            "kind": "values",
+            "values": [0.333984375, -0.82177734375, 0.35009765625, -1.189453125, -0.8701171875, 0.043914794921875, -0.35595703125, 0.333984375, -0.82177734375, 0.35009765625, -1.189453125, -0.8701171875, 0.043914794921875, -0.35595703125, 0.333984375, -0.82177734375, 0.35009765625, -1.189453125, -0.8701171875, 0.043914794921875, -0.35595703125, 0.333984375, -0.82177734375, 0.35009765625, -1.189453125, -0.8701171875, 0.043914794921875, -0.35595703125]
+          }
+        },
+        "scale": {
+          "dtype": "float16",
+          "shape": [28],
+          "data": {
+            "kind": "values",
+            "values": [1.1455078125, 1.0283203125, 0.92626953125, 1.0771484375, 1.232421875, 0.89990234375, 1.060546875, 1.1455078125, 1.0283203125, 0.92626953125, 1.0771484375, 1.232421875, 0.89990234375, 1.060546875, 1.1455078125, 1.0283203125, 0.92626953125, 1.0771484375, 1.232421875, 0.89990234375, 1.060546875, 1.1455078125, 1.0283203125, 0.92626953125, 1.0771484375, 1.232421875, 0.89990234375, 1.060546875]
+          }
+        }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "float16",
+          "shape": [1, 28],
+          "data": {
+            "kind": "values",
+            "values": [0.56494140625, -1.248046875, 0.47900390625, -1.892578125, -1.583984375, 0.058319091796875, -0.5576171875, 0.56494140625, -1.248046875, 0.47900390625, -1.892578125, -1.583984375, 0.058319091796875, -0.5576171875, 0.56494140625, -1.248046875, 0.47900390625, -1.892578125, -1.583984375, 0.058319091796875, -0.5576171875, 0.56494140625, -1.248046875, 0.47900390625, -1.892578125, -1.583984375, 0.058319091796875, -0.5576171875]
+          },
+          "tolerance": 0.000002
+        }
+      },
+      "provenance": {
+        "source": "ONNX RMSNormalization typed normalization and scale stages",
+        "notes": "Float32 Normalized midpoint followed by cast to half and affine half rounding. The CPU reference and independent PyTorch float32 typed stages agree; rounding float64 Normalized directly to half gives a different result."
+      }
+    },
     {
       "name": "subgroup_vec4_2x512",
       "attrs": { "epsilon": 0.000001, "axis": -1 },
     {
       "name": "onnx23_f16_cast_before_scale_row_vec4_exact",
       "provenance": {
+        "notes": "ONNX RMSNormalization-23 casts Normalized back to X's dtype before Scale. This vector-scale case differs by one float16 ULP if the cast is deferred."
       },
       "requires": { "features": ["shader-f16"] },
       "attrs": { "epsilon": 0.00001, "axis": -1 },
     {
       "name": "onnx23_f16_cast_before_scale_splitk_exact",
       "provenance": {
+        "notes": "A four-element float16 row with a scalar Scale checks the standardized cast of Normalized back to X's dtype before Scale."
       },
       "requires": { "features": ["shader-f16"] },
       "tunables": { "SPLIT_MIN_HIDDEN": 1, "SPLIT_TARGET_ELEMENTS": 1 },