sync 91d990483a17

Browse files

Files changed (8) hide show

README.md +13 -9
build/webgpu/bench.json +0 -1
build/webgpu/lp-norm-divide.wgsl.jinja +0 -3
build/webgpu/lp-norm-reduce.wgsl.jinja +5 -8
build/webgpu/manifest.json +150 -212
build/webgpu/metadata.json +17 -9
build/webgpu/norm-row-stats.wgsl.jinja +112 -14
build/webgpu/test.json +17 -42

README.md CHANGED Viewed

@@ -18,15 +18,15 @@ See the [ONNX `LpNormalization` spec](https://onnx.ai/onnx/operators/onnx__LpNor
 ## Inputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
-| --- | --- | --- | --- | --- | --- | --- |
-| `input` | `input` | `T` | — | — | Input tensor to normalize. | required |
 ## Outputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
-| --- | --- | --- | --- | --- | --- | --- |
-| `output` | `output` | `T` | same as `input` | same as `input` | Tensor after Lp-normalization; same shape as the input. | required |
 ## Attributes
@@ -45,7 +45,7 @@ Default values (overridable per request):
 ## Files
-- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
 - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
@@ -56,10 +56,14 @@ Default values (overridable per request):
 ## Use with `@huggingface/kernels`
-The loader derives every required output's shape and logical dtype from the manifest contract and this call.
-It then allocates the result tensors automatically.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

 ## Inputs
+| Name | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- |
+| `input` | `T` | — | — | Input tensor to normalize. | required |
 ## Outputs
+| Name | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- |
+| `output` | `T` | same as `input` | same as `input` | Tensor after Lp-normalization; same shape as the input. | required |
 ## Attributes
 ## Files
+- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
 - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 ## Use with `@huggingface/kernels`
+```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.2
+```
+Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

build/webgpu/bench.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "ai.onnx.LpNormalization",
   "tunableSpace": {
     "WORKGROUP_SIZE": [64, 128, 256],
     "SPLIT_TARGET_DIM": [128, 256, 512, 1024],

 {
   "tunableSpace": {
     "WORKGROUP_SIZE": [64, 128, 256],
     "SPLIT_TARGET_DIM": [128, 256, 512, 1024],

build/webgpu/lp-norm-divide.wgsl.jinja CHANGED Viewed

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





1	{{ env.wgsl.resourceDeclarations }}
2
3	const WG: u32 = {{ workgroupSize }}u;

build/webgpu/lp-norm-reduce.wgsl.jinja CHANGED Viewed

@@ -1,14 +1,11 @@
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ workgroupSize }}u;
-{% if source.segmented %}
 const SPLIT: u32 = {{ split }}u;
 {% endif %}
-{% if source.segmented %}
 // Reduce a strided normalization axis in segments. Workgroup rows divide a long
 // axis and write partial sums for the final reduction.
 {% else %}
@@ -18,14 +15,14 @@ const SPLIT: u32 = {{ split }}u;
 // tensor.
 @compute @workgroup_size(WG)
 fn main(
-  @builtin(global_invocation_id) gid: vec3<u32>{% if source.segmented %},
   @builtin(workgroup_id) wid: vec3<u32>,
 {% else %},
 {% endif %}
   @builtin(num_workgroups) nwg: vec3<u32>
 ) {
   let stride = nwg.x * WG;
-{% if source.segmented %}
   let segment = wid.y;
   let chunk = (params.dim + SPLIT - 1u) / SPLIT;
   let begin = segment * chunk;
@@ -43,7 +40,7 @@ fn main(
       let v = f32(input[base + j * params.inner]);
       norm += select(v * v, abs(v), params.p == 1u);
     }
-{% if source.segmented %}
     normOutput[segment * params.rows + r] = norm;
 {% else %}
     if (params.p == 2u) { norm = sqrt(norm); }

 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ workgroupSize }}u;
+{% if segmented %}
 const SPLIT: u32 = {{ split }}u;
 {% endif %}
+{% if segmented %}
 // Reduce a strided normalization axis in segments. Workgroup rows divide a long
 // axis and write partial sums for the final reduction.
 {% else %}
 // tensor.
 @compute @workgroup_size(WG)
 fn main(
+  @builtin(global_invocation_id) gid: vec3<u32>{% if segmented %},
   @builtin(workgroup_id) wid: vec3<u32>,
 {% else %},
 {% endif %}
   @builtin(num_workgroups) nwg: vec3<u32>
 ) {
   let stride = nwg.x * WG;
+{% if segmented %}
   let segment = wid.y;
   let chunk = (params.dim + SPLIT - 1u) / SPLIT;
   let begin = segment * chunk;
       let v = f32(input[base + j * params.inner]);
       norm += select(v * v, abs(v), params.p == 1u);
     }
+{% if segmented %}
     normOutput[segment * params.rows + r] = norm;
 {% else %}
     if (params.p == 2u) { norm = sqrt(norm); }

build/webgpu/manifest.json CHANGED Viewed

@@ -2,185 +2,61 @@
   "domain": "ai.onnx",
   "name": "LpNormalization",
   "sinceVersion": 1,
-  "description": "Applies Lp-normalization to the input tensor along the specified axis: `output = input / Lp_norm(input, axis)`. Supports L1 (`p=1`) and L2 (`p=2`) norms. Where the Lp norm is zero, the output is defined as zero to avoid division by zero.",
-  "inputs": [{ "role": "input", "dtype": "T", "description": "Input tensor to normalize." }],
-  "outputs": [
-    {
-      "role": "output",
-      "dtype": "T",
-      "rank": "ranks.input",
-      "description": "Tensor after Lp-normalization; same shape as the input.",
-      "shape": "shapes.input"
-    }
-  ],
-  "attributes": { "axis": -1, "p": 2 },
-  "attributeDescriptions": {
-    "axis": "The axis along which normalization is applied; `-1` means the last axis.",
-    "p": "The order of the Lp norm to use; only `1` (L1) or `2` (L2) are supported."
-  },
   "attributeConstraints": { "p": { "values": [1, 2] } },
   "typeConstraints": { "T": ["float32", "float16"] },
-  "args": {
-    "input": { "kind": "tensor", "semantic": "input", "role": "input" },
-    "output": { "kind": "tensor", "semantic": "output", "role": "output" }
-  },
   "tunables": {
-    "WORKGROUP_SIZE": 256,
-    "SPLIT_MIN_DIM": 512,
-    "SPLIT_MIN_ROWS": 32,
-    "SPLIT_TARGET_DIM": 256,
-    "MAX_SPLITS": 128
   },
   "derive": {
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
     "workgroupSize": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
-    "hasSubgroupId": "device.features.has(\"subgroups\") and device.wgslLanguageFeatures.has(\"subgroup_id\")"
   },
-  "constants": { "scalar": "dtypes.T" },
-  "bindingSets": {
-    "rowStats": [
-      {
-        "name": "x",
-        "arg": "input",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$ioElement"
-      },
-      {
-        "name": "y",
-        "arg": "output",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$ioElement"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "numel(shapes.input) / dim(shapes.input, -1)" },
-            {
-              "name": "rowStride",
-              "type": "u32",
-              "value": "max(1, min(numel(shapes.input) / dim(shapes.input, -1), device.limits.maxComputeWorkgroupsPerDimension))"
-            }
-          ]
-        }
-      }
-    ],
-    "axisReduce": [
-      {
-        "name": "input",
-        "arg": "input",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      { "name": "normOutput", "semantic": "rowNorms", "buffer": { "type": "storage" }, "elementType": "f32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "axisRows" },
-            { "name": "dim", "type": "u32", "value": "axisDim" },
-            { "name": "inner", "type": "u32", "value": "axisInner" },
-            { "name": "p", "type": "u32", "value": "attrs.p" }
-          ]
-        }
-      }
-    ],
-    "splitReduce": [
-      {
-        "name": "input",
-        "arg": "input",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      { "name": "normOutput", "semantic": "partials", "buffer": { "type": "storage" }, "elementType": "f32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "axisRows" },
-            { "name": "dim", "type": "u32", "value": "axisDim" },
-            { "name": "inner", "type": "u32", "value": "axisInner" },
-            { "name": "p", "type": "u32", "value": "attrs.p" }
-          ]
-        }
-      }
-    ],
-    "splitCombine": [
-      { "name": "partials", "semantic": "partials", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
-      { "name": "rowNorms", "semantic": "rowNorms", "buffer": { "type": "storage" }, "elementType": "f32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "axisRows" },
-            { "name": "p", "type": "u32", "value": "attrs.p" }
-          ]
-        }
-      }
-    ],
-    "axisDivide": [
-      {
-        "name": "input",
-        "arg": "input",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$divElem"
-      },
-      { "name": "rowNorms", "semantic": "rowNorms", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
-      {
-        "name": "output",
-        "arg": "output",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$divElem"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "count", "type": "u32", "value": "divideCount" },
-            { "name": "dim", "type": "u32", "value": "axisDim" },
-            { "name": "inner", "type": "u32", "value": "axisInner" }
-          ]
-        }
-      }
-    ]
   },
   "variants": [
     {
       "id": "axis_splitk",
       "priority": 20,
       "derive": {
         "axisDim": "dim(shapes.input, attrs.axis)",
         "axisInner": "inner(shapes.input, attrs.axis)",
         "axisRows": "numel(shapes.input) / axisDim",
         "split": "min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(axisDim, tunables.SPLIT_TARGET_DIM)))",
         "divideVec4": "numel(shapes.output) % 4 == 0 and axisInner % 4 == 0",
-        "divideCount": "numel(shapes.output) / 4 if divideVec4 else numel(shapes.output)"
-      },
-      "when": ["ranks.input >= 2", "ranks.output == ranks.input", "numel(shapes.input) == numel(shapes.output)", "inner(shapes.input, attrs.axis) > 1", "dim(shapes.input, attrs.axis) >= tunables.SPLIT_MIN_DIM", "numel(shapes.input) / dim(shapes.input, attrs.axis) >= tunables.SPLIT_MIN_ROWS", "min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(dim(shapes.input, attrs.axis), tunables.SPLIT_TARGET_DIM))) <= device.limits.maxComputeWorkgroupsPerDimension", "numel(shapes.input) / dim(shapes.input, attrs.axis) * 4 <= device.limits.maxStorageBufferBindingSize", "numel(shapes.input) / dim(shapes.input, attrs.axis) * 4 <= device.limits.maxBufferSize", "numel(shapes.input) / dim(shapes.input, attrs.axis) * min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(dim(shapes.input, attrs.axis), tunables.SPLIT_TARGET_DIM))) * 4 <= device.limits.maxStorageBufferBindingSize", "numel(shapes.input) / dim(shapes.input, attrs.axis) * min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(dim(shapes.input, attrs.axis), tunables.SPLIT_TARGET_DIM))) * 4 <= device.limits.maxBufferSize", "f16Ok(dtypes.T)"],
-      "constants": {
-        "usesF16": "dtypes.T == \"f16\"",
-        "split": "split",
         "divVec4": "divideVec4",
         "divElem": "(\"vec4<\" ~ dtypes.T ~ \">\") if divideVec4 else dtypes.T"
       },
@@ -192,114 +68,176 @@
         {
           "id": "split_reduce",
           "name": "LpNormalization.SplitReduce",
-          "source": { "shader": "lp-norm-reduce.wgsl.jinja", "inputs": { "segmented": true } },
-          "bindings": "splitReduce",
           "dispatch": {
-            "x": "min(ceilDiv(axisRows, workgroupSize), device.limits.maxComputeWorkgroupsPerDimension)",
             "y": "split"
           }
         },
         {
           "id": "combine",
           "name": "LpNormalization.SplitCombine",
-          "source": { "shader": "lp-norm-split-combine.wgsl.jinja", "inputs": { "split": "split" } },
-          "bindings": "splitCombine",
-          "dispatch": { "gridStride": "axisRows", "workgroupSize": "workgroupSize" }
         },
         {
           "id": "divide",
           "name": "LpNormalization.Divide",
           "shader": "lp-norm-divide.wgsl.jinja",
-          "bindings": "axisDivide",
-          "dispatch": { "gridStride": "divideCount", "workgroupSize": "workgroupSize" }
         }
       ]
     },
     {
       "id": "axis",
       "derive": {
         "axisDim": "dim(shapes.input, attrs.axis)",
         "axisInner": "inner(shapes.input, attrs.axis)",
         "axisRows": "numel(shapes.input) / axisDim",
-        "divideCount": "numel(shapes.output)"
       },
-      "when": ["ranks.input >= 1", "ranks.output == ranks.input", "numel(shapes.input) == numel(shapes.output)", "numel(shapes.input) / dim(shapes.input, attrs.axis) * 4 <= device.limits.maxStorageBufferBindingSize", "numel(shapes.input) / dim(shapes.input, attrs.axis) * 4 <= device.limits.maxBufferSize", "f16Ok(dtypes.T)"],
-      "constants": { "usesF16": "dtypes.T == \"f16\"", "divVec4": false, "divElem": "dtypes.T" },
       "intermediates": [{ "id": "rowNorms", "dtype": "float32", "shape": "[axisRows]" }],
       "passes": [
         {
           "id": "reduce",
           "name": "LpNormalization.RowReduce",
-          "source": { "shader": "lp-norm-reduce.wgsl.jinja", "inputs": { "segmented": false } },
-          "bindings": "axisReduce",
-          "dispatch": { "gridStride": "axisRows", "workgroupSize": "workgroupSize" }
         },
         {
           "id": "divide",
           "name": "LpNormalization.Divide",
           "shader": "lp-norm-divide.wgsl.jinja",
-          "bindings": "axisDivide",
-          "dispatch": { "gridStride": "divideCount", "workgroupSize": "workgroupSize" }
         }
       ]
     },
     {
       "id": "last_axis_row_vec4",
       "priority": 110,
-      "when": ["ranks.input >= 1", "ranks.output == ranks.input", "numel(shapes.input) == numel(shapes.output)", "(attrs.axis == -1 or attrs.axis == ranks.input - 1)", "dim(shapes.input, -1) % 4 == 0", "f16Ok(dtypes.T)"],
-      "constants": { "ioElement": "\"vec4<\" ~ dtypes.T ~ \">\"", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "LpNormalization.LastAxisRow",
-          "source": {
-            "shader": "norm-row-stats.wgsl.jinja",
-            "inputs": {
-              "mode": "\"lp\"",
-              "vec4": true,
-              "scalar": "dtypes.T",
-              "usesF16": "dtypes.T == \"f16\"",
-              "hidden": "dim(shapes.input, -1)",
-              "wg": "min(workgroupSize, pow2ceil(dim(shapes.input, -1) / 4))",
-              "p": "attrs.p",
-              "hiddenVec": "dim(shapes.input, -1) / 4",
-              "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
-              "combineSubgroups": "hasSubgroupId"
             }
           },
-          "subgroupCollectivesWidth": "portable",
-          "bindings": "rowStats",
-          "dispatch": { "workgroups": "numel(shapes.input) / dim(shapes.input, -1)" }
         }
       ]
     },
     {
       "id": "last_axis_row",
       "priority": 100,
-      "when": ["ranks.input >= 1", "ranks.output == ranks.input", "numel(shapes.input) == numel(shapes.output)", "(attrs.axis == -1 or attrs.axis == ranks.input - 1)", "true", "f16Ok(dtypes.T)"],
-      "constants": { "ioElement": "dtypes.T", "vectorScalar": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "LpNormalization.LastAxisRow",
-          "source": {
-            "shader": "norm-row-stats.wgsl.jinja",
-            "inputs": {
-              "mode": "\"lp\"",
-              "vec4": false,
-              "scalar": "dtypes.T",
-              "usesF16": "dtypes.T == \"f16\"",
-              "hidden": "dim(shapes.input, -1)",
-              "wg": "min(workgroupSize, pow2ceil(dim(shapes.input, -1)))",
-              "p": "attrs.p",
-              "hiddenVec": 1,
-              "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
-              "combineSubgroups": "hasSubgroupId"
             }
           },
-          "subgroupCollectivesWidth": "portable",
-          "bindings": "rowStats",
-          "dispatch": { "workgroups": "numel(shapes.input) / dim(shapes.input, -1)" }
         }
       ]
     }

   "domain": "ai.onnx",
   "name": "LpNormalization",
   "sinceVersion": 1,
+  "inputs": { "input": { "dtype": "T" } },
+  "outputs": { "output": { "dtype": "T", "rank": "ranks.input", "shape": "shapes.input" } },
+  "attributes": { "axis": { "default": -1 }, "p": { "default": 2 } },
   "attributeConstraints": { "p": { "values": [1, 2] } },
   "typeConstraints": { "T": ["float32", "float16"] },
   "tunables": {
+    "WORKGROUP_SIZE": { "default": 256 },
+    "SPLIT_MIN_DIM": { "default": 512 },
+    "SPLIT_MIN_ROWS": { "default": 32 },
+    "SPLIT_TARGET_DIM": { "default": 256 },
+    "MAX_SPLITS": { "default": 128 }
   },
   "derive": {
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
     "workgroupSize": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
+    "hasSubgroupId": "device.features.has(\"subgroups\") and device.wgslLanguageFeatures.has(\"subgroup_id\")",
+    "scalar": "dtypes.T"
   },
+  "when": ["ranks.output == ranks.input", "numel(shapes.input) == numel(shapes.output)", "f16Ok(dtypes.T)"],
+  "bindings": {
+    "input": { "buffer": "read-only-storage", "elementType": "$scalar" },
+    "params": {
+      "buffer": "uniform",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "axisRows" },
+        { "name": "dim", "type": "u32", "value": "axisDim" },
+        { "name": "inner", "type": "u32", "value": "axisInner" },
+        { "name": "p", "type": "u32", "value": "attrs.p" }
+      ]
+    },
+    "input_2": { "name": "input", "buffer": "read-only-storage", "elementType": "$divElem" },
+    "rowNorms": { "buffer": "read-only-storage", "elementType": "f32" },
+    "output": { "buffer": "storage", "elementType": "$divElem" },
+    "params_3": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "count", "type": "u32", "value": "divideCount" },
+        { "name": "dim", "type": "u32", "value": "axisDim" },
+        { "name": "inner", "type": "u32", "value": "axisInner" }
+      ]
+    }
   },
   "variants": [
     {
       "id": "axis_splitk",
       "priority": 20,
+      "when": ["ranks.input >= 2", "inner(shapes.input, attrs.axis) > 1", "dim(shapes.input, attrs.axis) >= tunables.SPLIT_MIN_DIM", "numel(shapes.input) / dim(shapes.input, attrs.axis) >= tunables.SPLIT_MIN_ROWS", "min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(dim(shapes.input, attrs.axis), tunables.SPLIT_TARGET_DIM))) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "numel(shapes.input) / dim(shapes.input, attrs.axis) * 4 <= device.limits.maxStorageBufferBindingSize", "numel(shapes.input) / dim(shapes.input, attrs.axis) * 4 <= device.limits.maxBufferSize", "numel(shapes.input) / dim(shapes.input, attrs.axis) * min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(dim(shapes.input, attrs.axis), tunables.SPLIT_TARGET_DIM))) * 4 <= device.limits.maxStorageBufferBindingSize", "numel(shapes.input) / dim(shapes.input, attrs.axis) * min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(dim(shapes.input, attrs.axis), tunables.SPLIT_TARGET_DIM))) * 4 <= device.limits.maxBufferSize"],
       "derive": {
         "axisDim": "dim(shapes.input, attrs.axis)",
         "axisInner": "inner(shapes.input, attrs.axis)",
         "axisRows": "numel(shapes.input) / axisDim",
         "split": "min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(axisDim, tunables.SPLIT_TARGET_DIM)))",
         "divideVec4": "numel(shapes.output) % 4 == 0 and axisInner % 4 == 0",
+        "divideCount": "numel(shapes.output) / 4 if divideVec4 else numel(shapes.output)",
         "divVec4": "divideVec4",
         "divElem": "(\"vec4<\" ~ dtypes.T ~ \">\") if divideVec4 else dtypes.T"
       },
         {
           "id": "split_reduce",
           "name": "LpNormalization.SplitReduce",
+          "shader": "lp-norm-reduce.wgsl.jinja",
+          "derive": { "segmented": true },
+          "bindings": ["input", { "scratch": "partials", "name": "normOutput", "elementType": "f32" }, "params"],
           "dispatch": {
+            "x": "min(ceilDiv(axisRows, workgroupSize), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
             "y": "split"
           }
         },
         {
           "id": "combine",
           "name": "LpNormalization.SplitCombine",
+          "shader": "lp-norm-split-combine.wgsl.jinja",
+          "derive": {},
+          "bindings": [
+            { "name": "partials", "buffer": "read-only-storage", "elementType": "f32" },
+            { "name": "rowNorms", "buffer": "storage", "elementType": "f32" },
+            {
+              "name": "params",
+              "struct": [
+                { "name": "rows", "type": "u32", "value": "axisRows" },
+                { "name": "p", "type": "u32", "value": "attrs.p" }
+              ]
+            }
+          ],
+          "dispatch": {
+            "x": "min(ceilDiv((axisRows), (workgroupSize)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
+            "y": 1,
+            "z": 1
+          }
         },
         {
           "id": "divide",
           "name": "LpNormalization.Divide",
           "shader": "lp-norm-divide.wgsl.jinja",
+          "bindings": ["input_2", "rowNorms", "output", "params_3"],
+          "dispatch": {
+            "x": "min(ceilDiv((divideCount), (workgroupSize)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
+            "y": 1,
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "axis",
+      "when": ["ranks.input >= 1", "numel(shapes.input) / dim(shapes.input, attrs.axis) * 4 <= device.limits.maxStorageBufferBindingSize", "numel(shapes.input) / dim(shapes.input, attrs.axis) * 4 <= device.limits.maxBufferSize"],
       "derive": {
         "axisDim": "dim(shapes.input, attrs.axis)",
         "axisInner": "inner(shapes.input, attrs.axis)",
         "axisRows": "numel(shapes.input) / axisDim",
+        "divideCount": "numel(shapes.output)",
+        "divVec4": false,
+        "divElem": "dtypes.T"
       },
       "intermediates": [{ "id": "rowNorms", "dtype": "float32", "shape": "[axisRows]" }],
       "passes": [
         {
           "id": "reduce",
           "name": "LpNormalization.RowReduce",
+          "shader": "lp-norm-reduce.wgsl.jinja",
+          "derive": { "segmented": false },
+          "bindings": ["input", { "scratch": "rowNorms", "name": "normOutput", "elementType": "f32" }, "params"],
+          "dispatch": {
+            "x": "min(ceilDiv((axisRows), (workgroupSize)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
+            "y": 1,
+            "z": 1
+          }
         },
         {
           "id": "divide",
           "name": "LpNormalization.Divide",
           "shader": "lp-norm-divide.wgsl.jinja",
+          "bindings": ["input_2", "rowNorms", "output", "params_3"],
+          "dispatch": {
+            "x": "min(ceilDiv((divideCount), (workgroupSize)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
+            "y": 1,
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "last_axis_row_vec4",
       "priority": 110,
+      "when": ["ranks.input >= 1", "(attrs.axis == -1 or attrs.axis == ranks.input - 1)", "dim(shapes.input, -1) % 4 == 0"],
+      "derive": { "ioElement": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "LpNormalization.LastAxisRow",
+          "shader": "norm-row-stats.wgsl.jinja",
+          "derive": {
+            "modeSpec": "\"lp\"",
+            "vec4": true,
+            "scalar": "dtypes.T",
+            "usesF16Spec": "dtypes.T == \"f16\"",
+            "hidden": "dim(shapes.input, -1)",
+            "wg": "min(workgroupSize, pow2ceil(dim(shapes.input, -1) / 4))",
+            "p": "attrs.p",
+            "hiddenVec": "dim(shapes.input, -1) / 4",
+            "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
+            "combineSubgroups": "hasSubgroupId"
+          },
+          "bindings": [
+            { "arg": "input", "name": "x", "elementType": "$ioElement" },
+            { "arg": "output", "name": "y", "elementType": "$ioElement" },
+            {
+              "name": "params",
+              "struct": [
+                { "name": "rows", "type": "u32", "value": "numel(shapes.input) / dim(shapes.input, -1)" },
+                {
+                  "name": "rowStride",
+                  "type": "u32",
+                  "value": "max(1, min(numel(shapes.input) / dim(shapes.input, -1), min(device.limits.maxComputeWorkgroupsPerDimension, 65535)))"
+                }
+              ]
             }
+          ],
+          "dispatch": {
+            "x": "min(numel(shapes.input) / dim(shapes.input, -1), 65535)",
+            "y": "ceilDiv(numel(shapes.input) / dim(shapes.input, -1), 65535)",
+            "z": 1
           },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "last_axis_row",
       "priority": 100,
+      "when": ["ranks.input >= 1", "(attrs.axis == -1 or attrs.axis == ranks.input - 1)", "true"],
+      "derive": { "ioElement": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "LpNormalization.LastAxisRow",
+          "shader": "norm-row-stats.wgsl.jinja",
+          "derive": {
+            "modeSpec": "\"lp\"",
+            "vec4": false,
+            "scalar": "dtypes.T",
+            "usesF16Spec": "dtypes.T == \"f16\"",
+            "hidden": "dim(shapes.input, -1)",
+            "wg": "min(workgroupSize, pow2ceil(dim(shapes.input, -1)))",
+            "p": "attrs.p",
+            "hiddenVec": 1,
+            "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
+            "combineSubgroups": "hasSubgroupId"
+          },
+          "bindings": [
+            { "arg": "input", "name": "x", "elementType": "$ioElement" },
+            { "arg": "output", "name": "y", "elementType": "$ioElement" },
+            {
+              "name": "params",
+              "struct": [
+                { "name": "rows", "type": "u32", "value": "numel(shapes.input) / dim(shapes.input, -1)" },
+                {
+                  "name": "rowStride",
+                  "type": "u32",
+                  "value": "max(1, min(numel(shapes.input) / dim(shapes.input, -1), min(device.limits.maxComputeWorkgroupsPerDimension, 65535)))"
+                }
+              ]
             }
+          ],
+          "dispatch": {
+            "x": "min(numel(shapes.input) / dim(shapes.input, -1), 65535)",
+            "y": "ceilDiv(numel(shapes.input) / dim(shapes.input, -1), 65535)",
+            "z": 1
           },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,21 +1,29 @@
 {
   "name": "ai.onnx.LpNormalization",
-  "id": "_ai_onnx_lpnormalization_webgpu_4171c72",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "4Qg+bNUooWdz8Aj8xcupfNIzsDtFpPO/oVyDuhMdCKQ=",
-      "lp-norm-divide.wgsl.jinja": "27WmHX7/sfxmHRGh068Q06O+5Up+UA1yNukXGXUmCLE=",
-      "lp-norm-reduce.wgsl.jinja": "XTvPbUz+NedeLvmFx2JZ0VnbBdx0t1TJ8QdaQlvhw+w=",
       "lp-norm-split-combine.wgsl.jinja": "PzQVkT5bUQ3+l2lqn3MgA2P0MbokIMa8Q6fVT4lFq9Q=",
-      "manifest.json": "rdr9JjuddPq9Xhf4rXnzFfMglX1ORDsP1HvkriIryy8=",
-      "norm-row-stats.wgsl.jinja": "Ru5DVZqQDj7YfQjw02snVOuPdbiVeFOxafk1dQmKQoU=",
-      "test.json": "7CtVpdOA+xLNbzNOVKMe12xYpagIX0a1f5YABmQa1Hc="
     }
   },
-  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
-  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.LpNormalization" }
 }

 {
   "name": "ai.onnx.LpNormalization",
+  "id": "_ai_onnx_lpnormalization_webgpu_7afe9b5",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "Zqjg8X2eQ7EWdSXd2uSW+U227BKDhM9aUvyzA2Fd75o=",
+      "lp-norm-divide.wgsl.jinja": "bIMXO9qn/IyThJQr5tapg04aNYkQqYt6LK/yz9D9ZBo=",
+      "lp-norm-reduce.wgsl.jinja": "Yu7Q9BM+uAAsEfA4VC1PQR7VOxdTyn+A+mjIryfLJwc=",
       "lp-norm-split-combine.wgsl.jinja": "PzQVkT5bUQ3+l2lqn3MgA2P0MbokIMa8Q6fVT4lFq9Q=",
+      "manifest.json": "EVQo8vZnPKz7vHj4+ZebBiS7UI2rAlpagBmiNdNEp6c=",
+      "norm-row-stats.wgsl.jinja": "iwey3jfLc6FXtfY4bBqo5+YG3ALWvo0RHLn12HHeJ0Q=",
+      "test.json": "7Vll6iQ9EAVx7E3pLedW2baT6dCaQLOfgy3qgY3xzyQ="
     }
   },
+  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.0",
+    "variants": {
+      "axis_splitk": ["lp-norm-divide.wgsl.jinja", "lp-norm-reduce.wgsl.jinja", "lp-norm-split-combine.wgsl.jinja"],
+      "axis": ["lp-norm-divide.wgsl.jinja", "lp-norm-reduce.wgsl.jinja"],
+      "last_axis_row_vec4": ["norm-row-stats.wgsl.jinja"],
+      "last_axis_row": ["norm-row-stats.wgsl.jinja"]
+    }
+  }
 }

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

@@ -1,9 +1,15 @@
-{% if source.usesF16 %}
 enable f16;
 {% endif %}
-{% set combineSubgroups = source.combineSubgroups %}
-{% set scalarIo = source.scalarIo if source.scalarIo is defined else false %}
-{% set lpOrder = source.p if source.p is defined else 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"
@@ -22,15 +28,58 @@ enable subgroups;
 // tree, then every thread applies the fused normalize + affine write.
 //
 // Lp mode divides by the norm and maps a zero norm to zero without epsilon.
-{% if source.mode != "lp" or not source.vec4 %}
-const HIDDEN: u32 = {{ source.hidden }}u;
 {% endif %}
-{% if source.vec4 %}
-const HIDDEN_V: u32 = {{ source.hiddenVec }}u;
 {% endif %}
-const WG: u32 = {{ source.wg }}u;
 {% if combineSubgroups %}
 var<workgroup> sg_partials: array<f32, WG>;
@@ -84,7 +133,14 @@ fn main(
     return;
   }
   let tid = lid.x;
-{% if source.vec4 and not scalarIo %}
   let base = row * HIDDEN_V;
 {% else %}
   let base = row * HIDDEN;
@@ -92,9 +148,16 @@ fn main(
   var acc = 0.0;
-{% if source.vec4 %}
   for (var i = tid; i < HIDDEN_V; i = i + WG) {
     let v = vec4<f32>(x[base + i]);
 {% if lpOrder == 1 %}
     let a = abs(v);
     acc = acc + a.x + a.y + a.z + a.w;
@@ -104,7 +167,12 @@ fn main(
   }
 {% else %}
   for (var i = tid; i < HIDDEN; i = i + WG) {
     let v = f32(x[base + i]);
 {% if lpOrder == 1 %}
     acc = acc + abs(v);
 {% else %}
@@ -121,19 +189,49 @@ fn main(
   let norm = total;
 {% endif %}
-{% if source.vec4 %}
   for (var i = tid; i < HIDDEN_V; i = i + WG) {
     let idx = base + i;
     let v = vec4<f32>(x[idx]);
     let normalized = select(v / vec4<f32>(norm), vec4<f32>(0.0), vec4<bool>(norm == 0.0));
-    y[idx] = {{ source.vecType }}(normalized);
   }
 {% else %}
   for (var i = tid; i < HIDDEN; i = i + WG) {
     let idx = base + i;
     let v = f32(x[idx]);
     let normalized = select(v / norm, 0.0, norm == 0.0);
-    y[idx] = {{ source.scalar }}(normalized);
   }
 {% 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 lpOrder = p if p is defined else 0 %}
+{% 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 reduceThreadParameters = ", sg_lane: u32, sg_id: u32, num_sg: u32"
     if combineSubgroups else ", tid: u32" %}
 {% set reduceThreadArguments = ", sg_lane, sg_id, num_sg"
 // tree, then every thread applies the fused normalize + affine write.
 //
 // Lp mode divides by the norm and maps a zero norm to zero without epsilon.
+{% if modeSpec != "lp" or not vec4 %}
+const HIDDEN: u32 = {{ hidden }}u;
 {% endif %}
+{% 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;
+{% if rmsChainNorm %}
+const EPSILON2: f32 = {{ epsilon2 }};
+{% endif %}
+{% if packedBf16Embedding %}
+{% if vec4 %}
+fn unpack_bf16_pair(word: u32) -> vec2<f32> {
+  let bits = vec2<u32>(word & 0xffffu, word >> 16u);
+  return bitcast<vec2<f32>>(bits << vec2<u32>(16u));
+}
+{% endif %}
+{% if not vec4 %}
+fn embedding_scalar(source_row: u32, hidden: u32) -> f32 {
+  if (source_row >= NUM_ROWS) {
+    return 0.0;
+  }
+  let word = x[source_row * HIDDEN_PAIRS + (hidden >> 1u)];
+  let bits = select(word & 0xffffu, word >> 16u, (hidden & 1u) != 0u);
+  return bitcast<f32>(bits << 16u);
+}
 {% endif %}
+{% if vec4 %}
+fn embedding_vec4(source_row: u32, hidden_vec: u32) -> vec4<f32> {
+  if (source_row >= NUM_ROWS) {
+    return vec4<f32>(0.0);
+  }
+  let base = source_row * HIDDEN_PAIRS + hidden_vec * 2u;
+  let low = unpack_bf16_pair(x[base]);
+  let high = unpack_bf16_pair(x[base + 1u]);
+  return vec4<f32>(low, high);
+}
+{% endif %}
+{% endif %}
+{% if vec4 and scalarIo %}
+fn load_vec4(index: u32) -> vec4<f32> {
+  return vec4<f32>(x[index], x[index + 1u], x[index + 2u], x[index + 3u]);
+}
+{% endif %}
 {% if combineSubgroups %}
 var<workgroup> sg_partials: array<f32, WG>;
     return;
   }
   let tid = lid.x;
+{% if packedBf16Embedding %}
+  let source_row = indices[row];
+{% if vec4 %}
+  let base = row * HIDDEN_V;
+{% else %}
+  let base = row * HIDDEN;
+{% endif %}
+{% elif vec4 and not scalarIo %}
   let base = row * HIDDEN_V;
 {% else %}
   let base = row * HIDDEN;
   var acc = 0.0;
+{% if vec4 %}
   for (var i = tid; i < HIDDEN_V; i = i + WG) {
+{% if packedBf16Embedding %}
+    let v = embedding_vec4(source_row, i);
+    embedding_out[base + i] = v;
+{% elif scalarIo %}
+    let v = load_vec4(base + i * 4u);
+{% else %}
     let v = vec4<f32>(x[base + i]);
+{% endif %}
 {% if lpOrder == 1 %}
     let a = abs(v);
     acc = acc + a.x + a.y + a.z + a.w;
   }
 {% 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 %}
 {% if lpOrder == 1 %}
     acc = acc + abs(v);
 {% else %}
   let norm = total;
 {% 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 %}
     let normalized = select(v / vec4<f32>(norm), vec4<f32>(0.0), vec4<bool>(norm == 0.0));
+    y[idx] = {{ vecType }}(normalized);
   }
+{% 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 normalized = select(v / norm, 0.0, norm == 0.0);
+    y[idx] = {{ scalar }}(normalized);
   }
 {% endif %}
 }

build/webgpu/test.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "ai.onnx.LpNormalization",
   "fixtureArrays": {
     "ort_l1_axis1_rank3_input_input": [5.93932154, 7.4367043, 6.42487038, 5.90394865, 4.81289319, 6.81304702, 4.9382849, 9.02595701, 9.67296484, 4.45097367, 8.12552534, 5.76005428, 6.11240105, 9.33036974, 1.63932452, 1.7841637, 1.18196558, 8.49357861, 8.00341076, 8.83010933, 9.80756508, 8.19242708, 5.15331426, 8.02476259]
   },
@@ -72,7 +71,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/nn/lp_norm_op_test.cc",
         "test": "LpNormalizationTest.L2NormalizationWithZeroNorm",
-        "notes": "Axis-0 companion: normal inputs whose squared L2 sums are finite subnormals must still produce nonzero normalized values."
       },
       "attrs": { "axis": 0, "p": 2 },
       "inputs": {
@@ -93,7 +92,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/nn/lp_norm_op_test.cc",
         "test": "LpNormalizationTest.L2NormalizationWithZeroNorm",
-        "notes": "Vec4 last-axis companion: the L2 norm is tiny but nonzero, so the direction of the vector must survive normalization."
       },
       "attrs": { "axis": -1, "p": 2 },
       "inputs": {
@@ -142,7 +141,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/nn/lp_norm_op_test.cc",
         "test": "LpNormalizationTest.L1NormalizationWithZeroNorm",
-        "notes": "Axis-0 companion: the strided two-pass p1 path must not collapse finite subnormal L1 norms to zero."
       },
       "attrs": { "axis": 0, "p": 1 },
       "inputs": {
@@ -170,7 +169,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/nn/lp_norm_op_test.cc",
         "test": "LpNormalizationTest.L1NormalizationWithZeroNorm",
-        "notes": "Vec4 last-axis companion: finite subnormal L1 totals should still produce stable signed ratios."
       },
       "attrs": { "axis": -1, "p": 1 },
       "inputs": {
@@ -207,7 +206,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/nn/lp_norm_op_test.cc",
         "test": "LpNormalizationTest.L2Normalization_FP16",
-        "notes": "Uses large f16 values to preserve ORT's float-accumulation regression coverage."
       },
       "attrs": { "axis": -1, "p": 2 },
       "inputs": { "input": { "dtype": "float16", "shape": [2, 128], "data": { "kind": "constant", "value": 100.0 } } },
@@ -219,17 +218,16 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/nn/lp_norm_op_test.cc",
         "test": "LpNormalizationTest.L1Normalization_FP16",
-        "notes": "Uses large f16 values to preserve ORT's float-accumulation regression coverage."
       },
       "attrs": { "axis": -1, "p": 1 },
       "inputs": { "input": { "dtype": "float16", "shape": [2, 128], "data": { "kind": "constant", "value": 200.0 } } },
-      "outputs": { "output": { "dtype": "float16", "shape": [2, 128] } },
-      "tolerance": 0.002
     },
     {
       "name": "f16_p1_last_axis_constant200_norm_scale_lock",
       "provenance": {
-        "notes": "Mirrors the ORT-derived ort_f16_p1_last_axis_large_values_float_accumulation case, whose 0.002 tolerance against a 0.0078125 output is a min detectable uniform scale error of 0.256 - a 25% error in the L1 norm passes. The ORT case is left as imported; this one keeps the identical input (128 columns of 200.0, so the answer 200/25600 = 2^-7 is exact in f16) and asserts it to 1e-6, which pins the divisor to the column count."
       },
       "attrs": { "axis": -1, "p": 1 },
       "inputs": { "input": { "dtype": "float16", "shape": [2, 128], "data": { "kind": "constant", "value": 200.0 } } },
@@ -459,16 +457,8 @@
     },
     {
       "name": "f16_p2_axis0_splitk_decode",
-      "attrs": { "axis": 0, "p": 2 },
-      "inputs": {
-        "input": { "dtype": "float16", "shape": [512, 256], "data": { "kind": "linspace", "start": -1.0, "end": 1.0 } }
-      },
-      "outputs": { "output": { "dtype": "float16", "shape": [512, 256], "tolerance": 0.01 } }
-    },
-    {
-      "name": "f16_p2_axis0_splitk_decode_norm_scale_lock",
       "provenance": {
-        "notes": "The p=2 half of the axis_splitk route: f16_p2_axis0_splitk_decode's 0.01 absolute tolerance against a 0.0765 maximum output is a min detectable uniform scale error of 0.131, so a partial sum-of-squares combined without the square root, or a rsqrt applied per split instead of once, is within tolerance. Same shape/attrs/data with a tolerance at f16 output resolution."
       },
       "attrs": { "axis": 0, "p": 2 },
       "inputs": {
@@ -478,20 +468,8 @@
     },
     {
       "name": "f16_p1_rank3_axis0_splitk_scalardiv",
-      "attrs": { "axis": 0, "p": 1 },
-      "inputs": {
-        "input": {
-          "dtype": "float16",
-          "shape": [512, 86, 3],
-          "data": { "kind": "linspace", "start": -2.0, "end": 3.0 }
-        }
-      },
-      "outputs": { "output": { "dtype": "float16", "shape": [512, 86, 3], "tolerance": 0.01 } }
-    },
-    {
-      "name": "f16_p1_rank3_axis0_splitk_scalardiv_norm_scale_lock",
       "provenance": {
-        "notes": "LpNormalization output is scale-invariant in the input, so no data choice can lift it above ~1/n along the normalized axis; the only way a fixture sees a denominator error is a tolerance commensurate with that magnitude. f16_p1_rank3_axis0_splitk_scalardiv pairs a 0.01 absolute tolerance with a 4.5e-3 maximum output (min detectable uniform scale error 2.22), so the split-k L1 partial sum, its combine, and the scalar divide could each be off by any factor and pass. Same shape, axis, p and data on the axis_splitk route, with the tolerance retightened to a few f16 ulp at the true output scale."
       },
       "attrs": { "axis": 0, "p": 1 },
       "inputs": {
@@ -507,16 +485,8 @@
     },
     {
       "name": "f16_p2_last_axis_hidden2048",
-      "attrs": { "axis": -1, "p": 2 },
-      "inputs": {
-        "input": { "dtype": "float16", "shape": [256, 2048], "data": { "kind": "linspace", "start": -1.0, "end": 1.0 } }
-      },
-      "outputs": { "output": { "dtype": "float16", "shape": [256, 2048], "tolerance": 0.01 } }
-    },
-    {
-      "name": "f16_p2_last_axis_hidden2048_norm_scale_lock",
       "provenance": {
-        "notes": "last_axis_row_vec4 over a 2048-wide row: the sibling's 0.01 absolute tolerance against a 0.0383 maximum output is a min detectable uniform scale error of 0.261, so the row norm could be a quarter or four times too large and pass. Same shape/attrs/data with the tolerance scaled to the output."
       },
       "attrs": { "axis": -1, "p": 2 },
       "inputs": {
@@ -546,7 +516,12 @@
           "data": { "kind": "linspace", "start": -0.5, "end": 0.5 }
         }
       },
-      "outputs": { "output": { "dtype": "float16", "shape": [8, 1024, 768], "tolerance": 0.01 } }
     },
     {
       "name": "f16_vec4_last_axis_zero_norm_guard",

 {
   "fixtureArrays": {
     "ort_l1_axis1_rank3_input_input": [5.93932154, 7.4367043, 6.42487038, 5.90394865, 4.81289319, 6.81304702, 4.9382849, 9.02595701, 9.67296484, 4.45097367, 8.12552534, 5.76005428, 6.11240105, 9.33036974, 1.63932452, 1.7841637, 1.18196558, 8.49357861, 8.00341076, 8.83010933, 9.80756508, 8.19242708, 5.15331426, 8.02476259]
   },
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/nn/lp_norm_op_test.cc",
         "test": "LpNormalizationTest.L2NormalizationWithZeroNorm",
+        "notes": "Along axis 0, normal inputs whose squared L2 sums are finite subnormals must still produce nonzero normalized values."
       },
       "attrs": { "axis": 0, "p": 2 },
       "inputs": {
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/nn/lp_norm_op_test.cc",
         "test": "LpNormalizationTest.L2NormalizationWithZeroNorm",
+        "notes": "On the vec4 last-axis path, a tiny nonzero L2 norm must preserve the vector's direction."
       },
       "attrs": { "axis": -1, "p": 2 },
       "inputs": {
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/nn/lp_norm_op_test.cc",
         "test": "LpNormalizationTest.L1NormalizationWithZeroNorm",
+        "notes": "Along axis 0, the strided two-pass p=1 path must not collapse finite subnormal L1 norms to zero."
       },
       "attrs": { "axis": 0, "p": 1 },
       "inputs": {
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/nn/lp_norm_op_test.cc",
         "test": "LpNormalizationTest.L1NormalizationWithZeroNorm",
+        "notes": "On the vec4 last-axis path, finite subnormal L1 totals must produce stable signed ratios."
       },
       "attrs": { "axis": -1, "p": 1 },
       "inputs": {
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/nn/lp_norm_op_test.cc",
         "test": "LpNormalizationTest.L2Normalization_FP16",
+        "notes": "Large float16 values require float32 accumulation for the p=2 norm."
       },
       "attrs": { "axis": -1, "p": 2 },
       "inputs": { "input": { "dtype": "float16", "shape": [2, 128], "data": { "kind": "constant", "value": 100.0 } } },
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/nn/lp_norm_op_test.cc",
         "test": "LpNormalizationTest.L1Normalization_FP16",
+        "notes": "Large float16 values require float32 accumulation for the p=1 norm."
       },
       "attrs": { "axis": -1, "p": 1 },
       "inputs": { "input": { "dtype": "float16", "shape": [2, 128], "data": { "kind": "constant", "value": 200.0 } } },
+      "outputs": { "output": { "dtype": "float16", "shape": [2, 128], "tolerance": 0.00002, "relTolerance": 0.002 } }
     },
     {
       "name": "f16_p1_last_axis_constant200_norm_scale_lock",
       "provenance": {
+        "notes": "Each 128-element row contains 200, so its L1 norm is 25,600 and every normalized value is exactly 2^-7 in float16. A 1e-6 tolerance makes an incorrect divisor observable."
       },
       "attrs": { "axis": -1, "p": 1 },
       "inputs": { "input": { "dtype": "float16", "shape": [2, 128], "data": { "kind": "constant", "value": 200.0 } } },
     },
     {
       "name": "f16_p2_axis0_splitk_decode",
       "provenance": {
+        "notes": "Shape [512,256] uses a split reduction along axis 0 for the p=2 norm. Float16-resolution tolerances expose an omitted final square root or a reciprocal square root applied separately to partial sums."
       },
       "attrs": { "axis": 0, "p": 2 },
       "inputs": {
     },
     {
       "name": "f16_p1_rank3_axis0_splitk_scalardiv",
       "provenance": {
+        "notes": "Shape [512,86,3] uses a split L1 reduction along axis 0 followed by scalar division. The normalized outputs peak near 4.5e-3, so float16-resolution tolerances make partial-sum, combine, and divisor errors observable."
       },
       "attrs": { "axis": 0, "p": 1 },
       "inputs": {
     },
     {
       "name": "f16_p2_last_axis_hidden2048",
       "provenance": {
+        "notes": "A 2,048-wide row exercises vec4 last-axis normalization with tolerance scaled to the approximately 0.0383 output magnitude, making uniform norm-scale errors observable."
       },
       "attrs": { "axis": -1, "p": 2 },
       "inputs": {
           "data": { "kind": "linspace", "start": -0.5, "end": 0.5 }
         }
       },
+      "outputs": {
+        "output": { "dtype": "float16", "shape": [8, 1024, 768], "tolerance": 0.0001, "relTolerance": 0.002 }
+      },
+      "provenance": {
+        "notes": "A rank-3 p=2 norm over 1,024 elements of axis 1 runs split partials and a combine. Float16-resolution tolerances scaled to the roughly 0.05 outputs expose an omitted square root or a rescaled divisor."
+      }
     },
     {
       "name": "f16_vec4_last_axis_zero_norm_guard",