sync 6fdf6301e2bb

Browse files

Files changed (6) hide show

README.md +13 -2
build/webgpu/layer-normalization.wgsl.jinja +193 -54
build/webgpu/manifest.json +263 -80
build/webgpu/metadata.json +14 -8
build/webgpu/norm-row-stats.wgsl.jinja +370 -119
build/webgpu/test.json +518 -0

README.md CHANGED Viewed

@@ -48,19 +48,30 @@ Default values (overridable per request):
 | --- | --- |
 | `T` | `float32`, `float16` |
 ## 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
 - [`layer-normalization.wgsl.jinja`](build/webgpu/layer-normalization.wgsl.jinja)
 - [`norm-row-stats.wgsl.jinja`](build/webgpu/norm-row-stats.wgsl.jinja)
 ## 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.

 | --- | --- |
 | `T` | `float32`, `float16` |
+## Implementation variants
+One implementation is selected per call from the device capabilities, the request shapes and the dtypes; these notes say what each one covers.
+- `last_axis_row_vec2` — Uses paired float32 loads for even last-axis rows that cannot use four-wide storage.
+- `last_axis_bias_row_vec2` — Uses paired float32 loads for even last-axis rows that cannot use four-wide storage.
+- `last_axis_broadcast_row_vec4` — Normalizes packed rows with affine tensors broadcast across leading dimensions; short rows share a workgroup within device limits.
+- `last_axis_broadcast_bias_row_vec4` — Normalizes packed rows with affine tensors broadcast across leading dimensions; short rows share a workgroup within device limits.
+- `last_axis_broadcast_rows_vec4` — Normalizes packed rows with affine tensors broadcast across leading dimensions; short rows share a workgroup within device limits.
+- `last_axis_broadcast_bias_rows_vec4` — Normalizes packed rows with affine tensors broadcast across leading dimensions; short rows share a workgroup within device limits.
 ## 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 cases
 - [`layer-normalization.wgsl.jinja`](build/webgpu/layer-normalization.wgsl.jinja)
 - [`norm-row-stats.wgsl.jinja`](build/webgpu/norm-row-stats.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/layer-normalization.wgsl.jinja CHANGED Viewed

@@ -33,25 +33,22 @@ fn {{ fn_name }}({% if out_numel != 0 and op_numel != 1 %}out_index: u32{% endif
 {% endfor %}
   return offset;
 {% endif %}
-}
-{%- endmacro %}{% macro broadcast_offset_call(fn_name, opShape, outShape, out_index) %}
 {% set op_numel = namespace(value=1) %}
-{% for d in opShape %}{% set op_numel.value = op_numel.value * d %}{% endfor %}
 {% set out_numel = namespace(value=1) %}
-{% for d in outShape %}{% set out_numel.value = out_numel.value * d %}{% endfor %}
-{{ fn_name }}({% if out_numel.value != 0 and op_numel.value != 1 %}{{ out_index }}{% endif %})
-{%- endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 const HIDDEN: u32 = {{ hiddenSize }}u;
 const EPSILON: f32 = {{ epsilon }};
 const WG: u32 = {{ workgroupSize }}u;
-var<workgroup> partial: array<f32, WG>;
-var<workgroup> row_mean: f32;
-var<workgroup> row_inv: f32;
 {% set xNumel = namespace(value=1) %}
 {% for dim in xShape %}
 {% set xNumel.value = xNumel.value * dim %}
@@ -61,82 +58,223 @@ var<workgroup> row_inv: f32;
 {% set scaleNumel.value = scaleNumel.value * dim %}
 {% endfor %}
 {% if scaleNumel.value != 1 %}
 {{ offset_fn("scale_offset", scaleShape, scaleShape | length, scaleShape == xShape, scaleNumel.value, xShape, xShape | length, xNumel.value) }}
 {% endif %}
 {% if hasBias %}
 {% set biasNumel = namespace(value=1) %}
 {% for dim in biasShape %}
 {% set biasNumel.value = biasNumel.value * dim %}
 {% endfor %}
 {% if biasNumel.value != 1 %}
 {{ offset_fn("bias_offset", biasShape, biasShape | length, biasShape == xShape, biasNumel.value, xShape, xShape | length, xNumel.value) }}
 {% endif %}
 {% 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;
@@ -187,3 +325,4 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
     y[index] = {{ scalar }}(value);
   }
 }

 {% endfor %}
   return offset;
 {% endif %}
+}{% endmacro %}
+{% macro broadcast_offset_call(fn_name, opShape, outShape, out_index) %}
 {% set op_numel = namespace(value=1) %}
+{% for d in opShape %}
+{% set op_numel.value = op_numel.value * d %}
+{% endfor %}
 {% set out_numel = namespace(value=1) %}
+{% for d in outShape %}
+{% set out_numel.value = out_numel.value * d %}
+{% endfor %}
+{{ fn_name }}({% if out_numel.value != 0 and op_numel.value != 1 %}{{ out_index }}{% endif %}){% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 const HIDDEN: u32 = {{ hiddenSize }}u;
 const EPSILON: f32 = {{ epsilon }};
 const WG: u32 = {{ workgroupSize }}u;
 {% set xNumel = namespace(value=1) %}
 {% for dim in xShape %}
 {% set xNumel.value = xNumel.value * dim %}
 {% set scaleNumel.value = scaleNumel.value * dim %}
 {% endfor %}
 {% if scaleNumel.value != 1 %}
 {{ offset_fn("scale_offset", scaleShape, scaleShape | length, scaleShape == xShape, scaleNumel.value, xShape, xShape | length, xNumel.value) }}
 {% endif %}
 {% if hasBias %}
 {% set biasNumel = namespace(value=1) %}
 {% for dim in biasShape %}
 {% set biasNumel.value = biasNumel.value * dim %}
 {% endfor %}
 {% if biasNumel.value != 1 %}
 {{ offset_fn("bias_offset", biasShape, biasShape | length, biasShape == xShape, biasNumel.value, xShape, xShape | length, xNumel.value) }}
 {% endif %}
+{% endif %}
+{% if scalar == "f16" %}
+{% set modeSpec = "layer" %}
+{% set halfWriteMean = writeMean %}
+{% set halfWriteInv = writeInvStdDev %}
+{% set halfScaleOffset = "0u" if scaleNumel.value == 1 else broadcast_offset_call("scale_offset", scaleShape, xShape, "base + i") %}
+{% if hasBias %}
+{% set halfBiasOffset = "0u" if biasNumel.value == 1 else broadcast_offset_call("bias_offset", biasShape, xShape, "base + i") %}
+{% endif %}
+{% set 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));
+}
+// 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.
+{% set halfWriteMean = halfWriteMean if halfWriteMean is defined else (writeStats and modeSpec == "layer") %}
+{% set halfWriteInv = halfWriteInv if halfWriteInv is defined else writeStats %}
+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;
+  var local_sum = vec2<f32>(0.0);
+  for (var i = tid; i < HIDDEN; i += WG) {
+    local_sum = pair_add(local_sum, vec2<f32>(load_value(base + i), 0.0));
+  }
+  let mean = pair_div(reduce_pair(local_sum, {{ reduceArgs }}), f32(HIDDEN));
+  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)));
+{% if halfWriteMean %}
+  if (tid == 0u) { mean_out[row] = mean.x; }
+{% endif %}
+{% if halfWriteInv %}
+  if (tid == 0u) { inv_std_out[row] = inv.x; }
 {% endif %}
+  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") }}]));
+{% if modeSpec == "layer" and hasBias %}
+    value = half_stage(value + f32(bias[{{ halfBiasOffset | 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>;
+var<workgroup> row_mean: f32;
+var<workgroup> row_inv: f32;
 {% 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[index] = {{ scalar }}(value);
   }
 }
+{% endif %}

build/webgpu/manifest.json CHANGED Viewed

@@ -15,7 +15,7 @@
   "attributes": { "axis": { "default": -1 }, "epsilon": { "default": 0.00001 }, "stash_type": { "default": 1 } },
   "attributeConstraints": { "stash_type": { "values": [1] } },
   "typeConstraints": { "T": ["float32", "float16"] },
-  "tunables": { "MAX_WORKGROUP_SIZE": { "default": 256 }, "SCALAR_FAST_MAX_HIDDEN": { "default": 1024 } },
   "derive": {
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
     "normWorkgroupCap": "min(tunables.MAX_WORKGROUP_SIZE, deviceWorkgroupCap)",
@@ -54,25 +54,82 @@
   },
   "when": ["f16Ok(dtypes.T)"],
   "bindings": {
-    "x": { "buffer": "read-only-storage", "elementType": "$vectorScalar" },
-    "scale": { "buffer": "read-only-storage", "elementType": "$vectorScalar" },
-    "y": { "buffer": "storage", "elementType": "$vectorScalar" },
     "params": {
-      "buffer": "uniform",
       "struct": [
         { "name": "rows", "type": "u32", "value": "normRows" },
         { "name": "rowStride", "type": "u32", "value": "normRowStride" }
       ]
     },
-    "bias": { "arg": "b", "buffer": "read-only-storage", "elementType": "$vectorScalar" },
-    "mean_out": { "arg": "mean", "buffer": "storage", "elementType": "f32" },
-    "inv_std_out": { "arg": "invStdDev", "buffer": "storage", "elementType": "f32" },
-    "x_2": { "name": "x", "buffer": "read-only-storage", "elementType": "$scalar" },
-    "scale_2": { "name": "scale", "buffer": "read-only-storage", "elementType": "$scalar" },
-    "y_2": { "name": "y", "buffer": "storage", "elementType": "$scalar" },
-    "bias_2": { "arg": "b", "name": "bias", "buffer": "read-only-storage", "elementType": "$scalar" }
   },
   "variants": [
     {
       "id": "last_axis_row_vec4",
       "priority": 110,
@@ -85,11 +142,10 @@
           "shader": "norm-row-stats.wgsl.jinja",
           "derive": {
             "modeSpec": "\"layer\"",
             "vec4": true,
-            "hasBias": false,
-            "writeStats": false,
-            "scalar": "dtypes.T",
-            "usesF16Spec": "dtypes.T == \"f16\"",
             "hidden": "dim(shapes.x, -1)",
             "wg": "lastAxisWgVec4",
             "epsilon": "attrs.epsilon",
@@ -98,8 +154,7 @@
             "combineSubgroups": "hasSubgroupId"
           },
           "bindings": ["x", "scale", "y", "params"],
-          "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
@@ -116,19 +171,17 @@
           "shader": "norm-row-stats.wgsl.jinja",
           "derive": {
             "modeSpec": "\"layer\"",
             "vec4": false,
-            "hasBias": false,
-            "writeStats": false,
-            "scalar": "dtypes.T",
-            "usesF16Spec": "dtypes.T == \"f16\"",
             "hidden": "dim(shapes.x, -1)",
             "wg": "lastAxisWg",
             "epsilon": "attrs.epsilon",
             "combineSubgroups": "hasSubgroupId"
           },
           "bindings": ["x", "scale", "y", "params"],
-          "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
@@ -144,11 +197,10 @@
           "shader": "norm-row-stats.wgsl.jinja",
           "derive": {
             "modeSpec": "\"layer\"",
             "vec4": true,
-            "hasBias": true,
-            "writeStats": false,
-            "scalar": "dtypes.T",
-            "usesF16Spec": "dtypes.T == \"f16\"",
             "hidden": "dim(shapes.x, -1)",
             "wg": "lastAxisWgVec4",
             "epsilon": "attrs.epsilon",
@@ -157,8 +209,7 @@
             "combineSubgroups": "hasSubgroupId"
           },
           "bindings": ["x", "scale", "bias", "y", "params"],
-          "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
@@ -175,19 +226,17 @@
           "shader": "norm-row-stats.wgsl.jinja",
           "derive": {
             "modeSpec": "\"layer\"",
             "vec4": false,
-            "hasBias": true,
-            "writeStats": false,
-            "scalar": "dtypes.T",
-            "usesF16Spec": "dtypes.T == \"f16\"",
             "hidden": "dim(shapes.x, -1)",
             "wg": "lastAxisWg",
             "epsilon": "attrs.epsilon",
             "combineSubgroups": "hasSubgroupId"
           },
           "bindings": ["x", "scale", "bias", "y", "params"],
-          "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
@@ -203,11 +252,10 @@
           "shader": "norm-row-stats.wgsl.jinja",
           "derive": {
             "modeSpec": "\"layer\"",
             "vec4": true,
-            "hasBias": false,
-            "writeStats": true,
-            "scalar": "dtypes.T",
-            "usesF16Spec": "dtypes.T == \"f16\"",
             "hidden": "dim(shapes.x, -1)",
             "wg": "lastAxisWgVec4",
             "epsilon": "attrs.epsilon",
@@ -216,8 +264,7 @@
             "combineSubgroups": "hasSubgroupId"
           },
           "bindings": ["x", "scale", "y", "mean_out", "inv_std_out", "params"],
-          "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
@@ -234,19 +281,17 @@
           "shader": "norm-row-stats.wgsl.jinja",
           "derive": {
             "modeSpec": "\"layer\"",
             "vec4": false,
-            "hasBias": false,
-            "writeStats": true,
-            "scalar": "dtypes.T",
-            "usesF16Spec": "dtypes.T == \"f16\"",
             "hidden": "dim(shapes.x, -1)",
             "wg": "lastAxisWg",
             "epsilon": "attrs.epsilon",
             "combineSubgroups": "hasSubgroupId"
           },
           "bindings": ["x", "scale", "y", "mean_out", "inv_std_out", "params"],
-          "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
@@ -262,11 +307,10 @@
           "shader": "norm-row-stats.wgsl.jinja",
           "derive": {
             "modeSpec": "\"layer\"",
             "vec4": true,
-            "hasBias": true,
-            "writeStats": true,
-            "scalar": "dtypes.T",
-            "usesF16Spec": "dtypes.T == \"f16\"",
             "hidden": "dim(shapes.x, -1)",
             "wg": "lastAxisWgVec4",
             "epsilon": "attrs.epsilon",
@@ -275,8 +319,7 @@
             "combineSubgroups": "hasSubgroupId"
           },
           "bindings": ["x", "scale", "bias", "y", "mean_out", "inv_std_out", "params"],
-          "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
@@ -293,19 +336,17 @@
           "shader": "norm-row-stats.wgsl.jinja",
           "derive": {
             "modeSpec": "\"layer\"",
             "vec4": false,
-            "hasBias": true,
-            "writeStats": true,
-            "scalar": "dtypes.T",
-            "usesF16Spec": "dtypes.T == \"f16\"",
             "hidden": "dim(shapes.x, -1)",
             "wg": "lastAxisWg",
             "epsilon": "attrs.epsilon",
             "combineSubgroups": "hasSubgroupId"
           },
           "bindings": ["x", "scale", "bias", "y", "mean_out", "inv_std_out", "params"],
-          "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
@@ -321,11 +362,10 @@
           "shader": "norm-row-stats.wgsl.jinja",
           "derive": {
             "modeSpec": "\"layer\"",
             "vec4": true,
             "hasBias": true,
             "writeStats": false,
-            "scalar": "dtypes.T",
-            "usesF16Spec": "dtypes.T == \"f16\"",
             "hidden": "suffixAxisSize",
             "wg": "suffixAxisWgVec4",
             "epsilon": "attrs.epsilon",
@@ -334,8 +374,7 @@
             "combineSubgroups": "hasSubgroupId"
           },
           "bindings": ["x", "scale", "bias", "y", "params"],
-          "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
@@ -544,9 +583,9 @@
       "priority": 31,
       "when": ["not present.b and meanOnlyOutputs and meanRowsOk", "lastAxisBroadcastScaleOk or suffixAxisBroadcastScaleOk"],
       "derive": {
-        "hasBias": false,
-        "writeMean": true,
-        "writeInvStdDev": false,
         "scalar": "dtypes.T",
         "hiddenSize": "genericHiddenSize",
         "workgroupSize": "genericWorkgroupSize",
@@ -558,7 +597,7 @@
           "name": "LayerNormalization.MeanOnly",
           "shader": "layer-normalization.wgsl.jinja",
           "derive": { "xShape": "shapes.x", "scaleShape": "shapes.scale" },
-          "bindings": ["x_2", "scale_2", "y_2", "mean_out", "params"],
           "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
         }
       ]
@@ -568,9 +607,9 @@
       "priority": 32,
       "when": ["present.b and meanOnlyOutputs and biasBroadcastOk and meanRowsOk", "lastAxisBroadcastScaleOk or suffixAxisBroadcastScaleOk"],
       "derive": {
-        "hasBias": true,
-        "writeMean": true,
-        "writeInvStdDev": false,
         "scalar": "dtypes.T",
         "hiddenSize": "genericHiddenSize",
         "workgroupSize": "genericWorkgroupSize",
@@ -582,7 +621,7 @@
           "name": "LayerNormalization.BiasMeanOnly",
           "shader": "layer-normalization.wgsl.jinja",
           "derive": { "xShape": "shapes.x", "scaleShape": "shapes.scale", "biasShape": "shapes.b" },
-          "bindings": ["x_2", "scale_2", "bias_2", "y_2", "mean_out", "params"],
           "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
         }
       ]
@@ -592,9 +631,9 @@
       "priority": 33,
       "when": ["not present.b and invStdOnlyOutputs and invStdRowsOk", "lastAxisBroadcastScaleOk or suffixAxisBroadcastScaleOk"],
       "derive": {
-        "hasBias": false,
-        "writeMean": false,
-        "writeInvStdDev": true,
         "scalar": "dtypes.T",
         "hiddenSize": "genericHiddenSize",
         "workgroupSize": "genericWorkgroupSize",
@@ -606,7 +645,7 @@
           "name": "LayerNormalization.InvStdDevOnly",
           "shader": "layer-normalization.wgsl.jinja",
           "derive": { "xShape": "shapes.x", "scaleShape": "shapes.scale" },
-          "bindings": ["x_2", "scale_2", "y_2", "inv_std_out", "params"],
           "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
         }
       ]
@@ -616,9 +655,9 @@
       "priority": 34,
       "when": ["present.b and invStdOnlyOutputs and biasBroadcastOk and invStdRowsOk", "lastAxisBroadcastScaleOk or suffixAxisBroadcastScaleOk"],
       "derive": {
-        "hasBias": true,
-        "writeMean": false,
-        "writeInvStdDev": true,
         "scalar": "dtypes.T",
         "hiddenSize": "genericHiddenSize",
         "workgroupSize": "genericWorkgroupSize",
@@ -630,10 +669,154 @@
           "name": "LayerNormalization.BiasInvStdDevOnly",
           "shader": "layer-normalization.wgsl.jinja",
           "derive": { "xShape": "shapes.x", "scaleShape": "shapes.scale", "biasShape": "shapes.b" },
-          "bindings": ["x_2", "scale_2", "bias_2", "y_2", "inv_std_out", "params"],
           "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
         }
       ]
     }
   ]
 }

   "attributes": { "axis": { "default": -1 }, "epsilon": { "default": 0.00001 }, "stash_type": { "default": 1 } },
   "attributeConstraints": { "stash_type": { "values": [1] } },
   "typeConstraints": { "T": ["float32", "float16"] },
+  "tunables": { "MAX_WORKGROUP_SIZE": { "default": 256 }, "SCALAR_FAST_MAX_HIDDEN": { "default": 4096 } },
   "derive": {
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
     "normWorkgroupCap": "min(tunables.MAX_WORKGROUP_SIZE, deviceWorkgroupCap)",
   },
   "when": ["f16Ok(dtypes.T)"],
   "bindings": {
+    "x": { "elementType": "$vectorScalar" },
+    "scale": { "elementType": "$vectorScalar" },
+    "y": { "elementType": "$vectorScalar" },
     "params": {
       "struct": [
         { "name": "rows", "type": "u32", "value": "normRows" },
         { "name": "rowStride", "type": "u32", "value": "normRowStride" }
       ]
     },
+    "bias": { "arg": "b", "elementType": "$vectorScalar" },
+    "mean_out": { "arg": "mean", "elementType": "f32" },
+    "inv_std_out": { "arg": "invStdDev", "elementType": "f32" },
+    "x_main": { "name": "x", "elementType": "$scalar" },
+    "scale_main": { "name": "scale", "elementType": "$scalar" },
+    "y_main": { "name": "y", "elementType": "$scalar" },
+    "bias_b": { "arg": "b", "name": "bias", "elementType": "$scalar" }
   },
   "variants": [
+    {
+      "id": "last_axis_row_vec2",
+      "priority": 105,
+      "when": ["dtypes.T == \"f32\"", "lastAxisExactScaleOk", "dim(shapes.x, -1) % 4 == 2", "noStatsOutputs", "not present.b"],
+      "derive": { "scalar": "dtypes.T", "vectorScalar": "\"vec2<f32>\"" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "LayerNormalization.LastAxisRowVec2",
+          "shader": "norm-row-stats.wgsl.jinja",
+          "derive": {
+            "modeSpec": "\"layer\"",
+            "vec4": true,
+            "hasBias": "present.b",
+            "writeStats": false,
+            "hidden": "dim(shapes.x, -1)",
+            "wg": "min(normWorkgroupCap, pow2ceil(dim(shapes.x, -1) / 2))",
+            "epsilon": "attrs.epsilon",
+            "hiddenVec": "dim(shapes.x, -1) / 2",
+            "vecType": "\"vec2<f32>\"",
+            "combineSubgroups": "hasSubgroupId",
+            "packedWidth": 2
+          },
+          "bindings": ["x", "scale", "y", "params"],
+          "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
+        }
+      ],
+      "demoteWhen": ["not lastAxisScalarFastOk"]
+    },
+    {
+      "id": "last_axis_bias_row_vec2",
+      "priority": 106,
+      "when": ["dtypes.T == \"f32\"", "lastAxisExactScaleOk", "dim(shapes.x, -1) % 4 == 2", "noStatsOutputs", "present.b and biasExactOk"],
+      "derive": { "scalar": "dtypes.T", "vectorScalar": "\"vec2<f32>\"" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "LayerNormalization.LastAxisRowVec2",
+          "shader": "norm-row-stats.wgsl.jinja",
+          "derive": {
+            "modeSpec": "\"layer\"",
+            "vec4": true,
+            "hasBias": "present.b",
+            "writeStats": false,
+            "hidden": "dim(shapes.x, -1)",
+            "wg": "min(normWorkgroupCap, pow2ceil(dim(shapes.x, -1) / 2))",
+            "epsilon": "attrs.epsilon",
+            "hiddenVec": "dim(shapes.x, -1) / 2",
+            "vecType": "\"vec2<f32>\"",
+            "combineSubgroups": "hasSubgroupId",
+            "packedWidth": 2
+          },
+          "bindings": ["x", "scale", "bias", "y", "params"],
+          "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
+        }
+      ],
+      "demoteWhen": ["not lastAxisScalarFastOk"]
+    },
     {
       "id": "last_axis_row_vec4",
       "priority": 110,
           "shader": "norm-row-stats.wgsl.jinja",
           "derive": {
             "modeSpec": "\"layer\"",
+            "compensateHalfStats": true,
             "vec4": true,
+            "hasBias": "present.b",
+            "writeStats": "fullStatsOutputs",
             "hidden": "dim(shapes.x, -1)",
             "wg": "lastAxisWgVec4",
             "epsilon": "attrs.epsilon",
             "combineSubgroups": "hasSubgroupId"
           },
           "bindings": ["x", "scale", "y", "params"],
+          "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
         }
       ]
     },
           "shader": "norm-row-stats.wgsl.jinja",
           "derive": {
             "modeSpec": "\"layer\"",
+            "compensateHalfStats": true,
             "vec4": false,
+            "hasBias": "present.b",
+            "writeStats": "fullStatsOutputs",
             "hidden": "dim(shapes.x, -1)",
             "wg": "lastAxisWg",
             "epsilon": "attrs.epsilon",
             "combineSubgroups": "hasSubgroupId"
           },
           "bindings": ["x", "scale", "y", "params"],
+          "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
         }
       ]
     },
           "shader": "norm-row-stats.wgsl.jinja",
           "derive": {
             "modeSpec": "\"layer\"",
+            "compensateHalfStats": true,
             "vec4": true,
+            "hasBias": "present.b",
+            "writeStats": "fullStatsOutputs",
             "hidden": "dim(shapes.x, -1)",
             "wg": "lastAxisWgVec4",
             "epsilon": "attrs.epsilon",
             "combineSubgroups": "hasSubgroupId"
           },
           "bindings": ["x", "scale", "bias", "y", "params"],
+          "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
         }
       ]
     },
           "shader": "norm-row-stats.wgsl.jinja",
           "derive": {
             "modeSpec": "\"layer\"",
+            "compensateHalfStats": true,
             "vec4": false,
+            "hasBias": "present.b",
+            "writeStats": "fullStatsOutputs",
             "hidden": "dim(shapes.x, -1)",
             "wg": "lastAxisWg",
             "epsilon": "attrs.epsilon",
             "combineSubgroups": "hasSubgroupId"
           },
           "bindings": ["x", "scale", "bias", "y", "params"],
+          "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
         }
       ]
     },
           "shader": "norm-row-stats.wgsl.jinja",
           "derive": {
             "modeSpec": "\"layer\"",
+            "compensateHalfStats": true,
             "vec4": true,
+            "hasBias": "present.b",
+            "writeStats": "fullStatsOutputs",
             "hidden": "dim(shapes.x, -1)",
             "wg": "lastAxisWgVec4",
             "epsilon": "attrs.epsilon",
             "combineSubgroups": "hasSubgroupId"
           },
           "bindings": ["x", "scale", "y", "mean_out", "inv_std_out", "params"],
+          "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
         }
       ]
     },
           "shader": "norm-row-stats.wgsl.jinja",
           "derive": {
             "modeSpec": "\"layer\"",
+            "compensateHalfStats": true,
             "vec4": false,
+            "hasBias": "present.b",
+            "writeStats": "fullStatsOutputs",
             "hidden": "dim(shapes.x, -1)",
             "wg": "lastAxisWg",
             "epsilon": "attrs.epsilon",
             "combineSubgroups": "hasSubgroupId"
           },
           "bindings": ["x", "scale", "y", "mean_out", "inv_std_out", "params"],
+          "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
         }
       ]
     },
           "shader": "norm-row-stats.wgsl.jinja",
           "derive": {
             "modeSpec": "\"layer\"",
+            "compensateHalfStats": true,
             "vec4": true,
+            "hasBias": "present.b",
+            "writeStats": "fullStatsOutputs",
             "hidden": "dim(shapes.x, -1)",
             "wg": "lastAxisWgVec4",
             "epsilon": "attrs.epsilon",
             "combineSubgroups": "hasSubgroupId"
           },
           "bindings": ["x", "scale", "bias", "y", "mean_out", "inv_std_out", "params"],
+          "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
         }
       ]
     },
           "shader": "norm-row-stats.wgsl.jinja",
           "derive": {
             "modeSpec": "\"layer\"",
+            "compensateHalfStats": true,
             "vec4": false,
+            "hasBias": "present.b",
+            "writeStats": "fullStatsOutputs",
             "hidden": "dim(shapes.x, -1)",
             "wg": "lastAxisWg",
             "epsilon": "attrs.epsilon",
             "combineSubgroups": "hasSubgroupId"
           },
           "bindings": ["x", "scale", "bias", "y", "mean_out", "inv_std_out", "params"],
+          "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
         }
       ]
     },
           "shader": "norm-row-stats.wgsl.jinja",
           "derive": {
             "modeSpec": "\"layer\"",
+            "compensateHalfStats": true,
             "vec4": true,
             "hasBias": true,
             "writeStats": false,
             "hidden": "suffixAxisSize",
             "wg": "suffixAxisWgVec4",
             "epsilon": "attrs.epsilon",
             "combineSubgroups": "hasSubgroupId"
           },
           "bindings": ["x", "scale", "bias", "y", "params"],
+          "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
         }
       ]
     },
       "priority": 31,
       "when": ["not present.b and meanOnlyOutputs and meanRowsOk", "lastAxisBroadcastScaleOk or suffixAxisBroadcastScaleOk"],
       "derive": {
+        "hasBias": "present.b",
+        "writeMean": "present.mean",
+        "writeInvStdDev": "present.invStdDev",
         "scalar": "dtypes.T",
         "hiddenSize": "genericHiddenSize",
         "workgroupSize": "genericWorkgroupSize",
           "name": "LayerNormalization.MeanOnly",
           "shader": "layer-normalization.wgsl.jinja",
           "derive": { "xShape": "shapes.x", "scaleShape": "shapes.scale" },
+          "bindings": ["x_main", "scale_main", "y_main", "mean_out", "params"],
           "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
         }
       ]
       "priority": 32,
       "when": ["present.b and meanOnlyOutputs and biasBroadcastOk and meanRowsOk", "lastAxisBroadcastScaleOk or suffixAxisBroadcastScaleOk"],
       "derive": {
+        "hasBias": "present.b",
+        "writeMean": "present.mean",
+        "writeInvStdDev": "present.invStdDev",
         "scalar": "dtypes.T",
         "hiddenSize": "genericHiddenSize",
         "workgroupSize": "genericWorkgroupSize",
           "name": "LayerNormalization.BiasMeanOnly",
           "shader": "layer-normalization.wgsl.jinja",
           "derive": { "xShape": "shapes.x", "scaleShape": "shapes.scale", "biasShape": "shapes.b" },
+          "bindings": ["x_main", "scale_main", "bias_b", "y_main", "mean_out", "params"],
           "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
         }
       ]
       "priority": 33,
       "when": ["not present.b and invStdOnlyOutputs and invStdRowsOk", "lastAxisBroadcastScaleOk or suffixAxisBroadcastScaleOk"],
       "derive": {
+        "hasBias": "present.b",
+        "writeMean": "present.mean",
+        "writeInvStdDev": "present.invStdDev",
         "scalar": "dtypes.T",
         "hiddenSize": "genericHiddenSize",
         "workgroupSize": "genericWorkgroupSize",
           "name": "LayerNormalization.InvStdDevOnly",
           "shader": "layer-normalization.wgsl.jinja",
           "derive": { "xShape": "shapes.x", "scaleShape": "shapes.scale" },
+          "bindings": ["x_main", "scale_main", "y_main", "inv_std_out", "params"],
           "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
         }
       ]
       "priority": 34,
       "when": ["present.b and invStdOnlyOutputs and biasBroadcastOk and invStdRowsOk", "lastAxisBroadcastScaleOk or suffixAxisBroadcastScaleOk"],
       "derive": {
+        "hasBias": "present.b",
+        "writeMean": "present.mean",
+        "writeInvStdDev": "present.invStdDev",
         "scalar": "dtypes.T",
         "hiddenSize": "genericHiddenSize",
         "workgroupSize": "genericWorkgroupSize",
           "name": "LayerNormalization.BiasInvStdDevOnly",
           "shader": "layer-normalization.wgsl.jinja",
           "derive": { "xShape": "shapes.x", "scaleShape": "shapes.scale", "biasShape": "shapes.b" },
+          "bindings": ["x_main", "scale_main", "bias_b", "y_main", "inv_std_out", "params"],
           "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
         }
       ]
+    },
+    {
+      "id": "last_axis_broadcast_row_vec4",
+      "priority": 90,
+      "when": ["dtypes.T == \"f32\"", "lastAxisBroadcastScaleOk", "dim(shapes.x, -1) >= 4", "dim(shapes.x, -1) % 4 == 0", "ranks.scale >= 1", "dim(shapes.scale, -1) == dim(shapes.x, -1)", "noStatsOutputs", "not present.b", "true"],
+      "demoteWhen": ["false"],
+      "derive": { "scalar": "dtypes.T", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "LayerNormalization.BroadcastRowsVec4",
+          "shader": "norm-row-stats.wgsl.jinja",
+          "derive": {
+            "modeSpec": "\"layer\"",
+            "vec4": true,
+            "hasBias": "present.b",
+            "writeStats": false,
+            "hidden": "dim(shapes.x, -1)",
+            "wg": "lastAxisWgVec4",
+            "epsilon": "attrs.epsilon",
+            "hiddenVec": "dim(shapes.x, -1) / 4",
+            "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
+            "combineSubgroups": "hasSubgroupId",
+            "affineRowBroadcast": true,
+            "xRowShape": "prefix(shapes.x, ranks.x - 1)",
+            "scaleRowShape": "prefix(shapes.scale, ranks.scale - 1)",
+            "biasRowShape": "prefix(shapes.b, ranks.b - 1) if present.b else []",
+            "batchRows": "1",
+            "batchLanes": "lastAxisWgVec4"
+          },
+          "bindings": ["x", "scale", "y", "params"],
+          "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
+        }
+      ]
+    },
+    {
+      "id": "last_axis_broadcast_bias_row_vec4",
+      "priority": 91,
+      "when": ["dtypes.T == \"f32\"", "lastAxisBroadcastScaleOk", "dim(shapes.x, -1) >= 4", "dim(shapes.x, -1) % 4 == 0", "ranks.scale >= 1", "dim(shapes.scale, -1) == dim(shapes.x, -1)", "noStatsOutputs", "present.b and biasBroadcastOk and ranks.b >= 1 and dim(shapes.b, -1) == dim(shapes.x, -1)", "true"],
+      "demoteWhen": ["false"],
+      "derive": { "scalar": "dtypes.T", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "LayerNormalization.BroadcastRowsVec4",
+          "shader": "norm-row-stats.wgsl.jinja",
+          "derive": {
+            "modeSpec": "\"layer\"",
+            "vec4": true,
+            "hasBias": "present.b",
+            "writeStats": false,
+            "hidden": "dim(shapes.x, -1)",
+            "wg": "lastAxisWgVec4",
+            "epsilon": "attrs.epsilon",
+            "hiddenVec": "dim(shapes.x, -1) / 4",
+            "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
+            "combineSubgroups": "hasSubgroupId",
+            "affineRowBroadcast": true,
+            "xRowShape": "prefix(shapes.x, ranks.x - 1)",
+            "scaleRowShape": "prefix(shapes.scale, ranks.scale - 1)",
+            "biasRowShape": "prefix(shapes.b, ranks.b - 1) if present.b else []",
+            "batchRows": "1",
+            "batchLanes": "lastAxisWgVec4"
+          },
+          "bindings": ["x", "scale", "bias", "y", "params"],
+          "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }
+        }
+      ]
+    },
+    {
+      "id": "last_axis_broadcast_rows_vec4",
+      "priority": 95,
+      "when": ["dtypes.T == \"f32\"", "lastAxisBroadcastScaleOk", "dim(shapes.x, -1) >= 4", "dim(shapes.x, -1) % 4 == 0", "ranks.scale >= 1", "dim(shapes.scale, -1) == dim(shapes.x, -1)", "noStatsOutputs", "not present.b", "floor(normWorkgroupCap / lastAxisWgVec4) > 1 and normWorkgroupCap * 8 <= device.limits.maxComputeWorkgroupStorageSize"],
+      "demoteWhen": ["normRows < floor(normWorkgroupCap / lastAxisWgVec4)"],
+      "derive": { "scalar": "dtypes.T", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "LayerNormalization.BroadcastRowsVec4",
+          "shader": "norm-row-stats.wgsl.jinja",
+          "derive": {
+            "modeSpec": "\"layer\"",
+            "vec4": true,
+            "hasBias": "present.b",
+            "writeStats": false,
+            "hidden": "dim(shapes.x, -1)",
+            "wg": "lastAxisWgVec4 * floor(normWorkgroupCap / lastAxisWgVec4)",
+            "epsilon": "attrs.epsilon",
+            "hiddenVec": "dim(shapes.x, -1) / 4",
+            "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
+            "combineSubgroups": "false",
+            "affineRowBroadcast": true,
+            "xRowShape": "prefix(shapes.x, ranks.x - 1)",
+            "scaleRowShape": "prefix(shapes.scale, ranks.scale - 1)",
+            "biasRowShape": "prefix(shapes.b, ranks.b - 1) if present.b else []",
+            "batchRows": "floor(normWorkgroupCap / lastAxisWgVec4)",
+            "batchLanes": "lastAxisWgVec4"
+          },
+          "bindings": ["x", "scale", "y", "params"],
+          "dispatch": {
+            "x": "min(ceilDiv(normRows, floor(normWorkgroupCap / lastAxisWgVec4)), 65535)",
+            "y": "ceilDiv(ceilDiv(normRows, floor(normWorkgroupCap / lastAxisWgVec4)), 65535)",
+            "z": 1
+          }
+        }
+      ]
+    },
+    {
+      "id": "last_axis_broadcast_bias_rows_vec4",
+      "priority": 96,
+      "when": ["dtypes.T == \"f32\"", "lastAxisBroadcastScaleOk", "dim(shapes.x, -1) >= 4", "dim(shapes.x, -1) % 4 == 0", "ranks.scale >= 1", "dim(shapes.scale, -1) == dim(shapes.x, -1)", "noStatsOutputs", "present.b and biasBroadcastOk and ranks.b >= 1 and dim(shapes.b, -1) == dim(shapes.x, -1)", "floor(normWorkgroupCap / lastAxisWgVec4) > 1 and normWorkgroupCap * 8 <= device.limits.maxComputeWorkgroupStorageSize"],
+      "demoteWhen": ["normRows < floor(normWorkgroupCap / lastAxisWgVec4)"],
+      "derive": { "scalar": "dtypes.T", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "LayerNormalization.BroadcastRowsVec4",
+          "shader": "norm-row-stats.wgsl.jinja",
+          "derive": {
+            "modeSpec": "\"layer\"",
+            "vec4": true,
+            "hasBias": "present.b",
+            "writeStats": false,
+            "hidden": "dim(shapes.x, -1)",
+            "wg": "lastAxisWgVec4 * floor(normWorkgroupCap / lastAxisWgVec4)",
+            "epsilon": "attrs.epsilon",
+            "hiddenVec": "dim(shapes.x, -1) / 4",
+            "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
+            "combineSubgroups": "false",
+            "affineRowBroadcast": true,
+            "xRowShape": "prefix(shapes.x, ranks.x - 1)",
+            "scaleRowShape": "prefix(shapes.scale, ranks.scale - 1)",
+            "biasRowShape": "prefix(shapes.b, ranks.b - 1) if present.b else []",
+            "batchRows": "floor(normWorkgroupCap / lastAxisWgVec4)",
+            "batchLanes": "lastAxisWgVec4"
+          },
+          "bindings": ["x", "scale", "bias", "y", "params"],
+          "dispatch": {
+            "x": "min(ceilDiv(normRows, floor(normWorkgroupCap / lastAxisWgVec4)), 65535)",
+            "y": "ceilDiv(ceilDiv(normRows, floor(normWorkgroupCap / lastAxisWgVec4)), 65535)",
+            "z": 1
+          }
+        }
+      ]
     }
   ]
 }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "ai.onnx.LayerNormalization",
-  "id": "_ai_onnx_layernormalization_webgpu_7b13eb1",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
@@ -8,16 +8,18 @@
     "algorithm": "sha256",
     "files": {
       "bench.json": "IPZzLq64+ycfl0fAzym0hDLorrSrZ5YDwLnpVEGgGHc=",
-      "layer-normalization.wgsl.jinja": "NJ1/CeeYHnToR+Ki5VHv4U4gxuHKa9zYg9PG2QOndME=",
-      "manifest.json": "aZkl41hMG0XVB5WthNHRcxLOw39oCNULp/kF2RczI8Y=",
-      "norm-row-stats.wgsl.jinja": "HUUqntKH7vqbffRpSu33tnmSU7PudcOCtl5QhV1xlGk=",
-      "test.json": "qJxDvb9POxT3rI+vLcahsq2mqIwAfen9b53MgOSG5Go="
     }
   },
-  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
   "webgpu": {
-    "manifestSpec": "2.0",
     "variants": {
       "last_axis_row_vec4": ["norm-row-stats.wgsl.jinja"],
       "last_axis_row": ["norm-row-stats.wgsl.jinja"],
       "last_axis_bias_row_vec4": ["norm-row-stats.wgsl.jinja"],
@@ -38,7 +40,11 @@
       "mean_only": ["layer-normalization.wgsl.jinja"],
       "bias_mean_only": ["layer-normalization.wgsl.jinja"],
       "inv_std_dev_only": ["layer-normalization.wgsl.jinja"],
-      "bias_inv_std_dev_only": ["layer-normalization.wgsl.jinja"]
     }
   }
 }

 {
   "name": "ai.onnx.LayerNormalization",
+  "id": "_ai_onnx_layernormalization_webgpu_cb0cade",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
     "algorithm": "sha256",
     "files": {
       "bench.json": "IPZzLq64+ycfl0fAzym0hDLorrSrZ5YDwLnpVEGgGHc=",
+      "layer-normalization.wgsl.jinja": "7h2F0b9zKAOqO4f5hhaJiX/TKGURziAlfQWUpO4IrWk=",
+      "manifest.json": "GWp3ACAaJzLWehyonJd7wD1r+5FDN4j8MJMiFV55SV8=",
+      "norm-row-stats.wgsl.jinja": "66BZfZ7q6X6xnqSuDyWvUov8wSdBYvxkmmfuqUl9uSk=",
+      "test.json": "qdUKVCiVNSQXTPzMr4JSIXCceUNzOoV4a5bLM2LNSo0="
     }
   },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
   "webgpu": {
+    "manifestSpec": "2.1",
     "variants": {
+      "last_axis_row_vec2": ["norm-row-stats.wgsl.jinja"],
+      "last_axis_bias_row_vec2": ["norm-row-stats.wgsl.jinja"],
       "last_axis_row_vec4": ["norm-row-stats.wgsl.jinja"],
       "last_axis_row": ["norm-row-stats.wgsl.jinja"],
       "last_axis_bias_row_vec4": ["norm-row-stats.wgsl.jinja"],
       "mean_only": ["layer-normalization.wgsl.jinja"],
       "bias_mean_only": ["layer-normalization.wgsl.jinja"],
       "inv_std_dev_only": ["layer-normalization.wgsl.jinja"],
+      "bias_inv_std_dev_only": ["layer-normalization.wgsl.jinja"],
+      "last_axis_broadcast_row_vec4": ["norm-row-stats.wgsl.jinja"],
+      "last_axis_broadcast_bias_row_vec4": ["norm-row-stats.wgsl.jinja"],
+      "last_axis_broadcast_rows_vec4": ["norm-row-stats.wgsl.jinja"],
+      "last_axis_broadcast_bias_rows_vec4": ["norm-row-stats.wgsl.jinja"]
     }
   }
 }

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

@@ -1,16 +1,6 @@
-{% 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 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 hasBias = hasBias is defined and hasBias %}
 {% set reduceThreadParameters = ", sg_lane: u32, sg_id: u32, num_sg: u32"
     if combineSubgroups else ", tid: u32" %}
 {% set reduceThreadArguments = ", sg_lane, sg_id, num_sg"
@@ -19,68 +9,359 @@ 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.
-//
-// Shifted moments avoid cancellation from a large common offset; scaling uses
-// inverseSqrt(variance + EPSILON).
-const HIDDEN: u32 = {{ hidden }}u;
-{% if vec4 %}
-const HIDDEN_V: u32 = {{ hiddenVec }}u;
 {% endif %}
-{% if packedBf16Embedding %}
-const HIDDEN_PAIRS: u32 = {{ hiddenPairs }}u;
-const NUM_ROWS: u32 = {{ numRows }}u;
 {% endif %}
-const WG: u32 = {{ wg }}u;
-const EPSILON: f32 = {{ epsilon }};
-{% if rmsChainNorm %}
-const EPSILON2: f32 = {{ epsilon2 }};
 {% endif %}
-{% if packedBf16Embedding %}
-{% if vec4 %}
-fn unpack_bf16_pair(word: u32) -> vec2<f32> {
-  let bits = vec2<u32>(word & 0xffffu, word >> 16u);
-  return bitcast<vec2<f32>>(bits << vec2<u32>(16u));
 }
 {% endif %}
-{% if not vec4 %}
-fn embedding_scalar(source_row: u32, hidden: u32) -> f32 {
-  if (source_row >= NUM_ROWS) {
-    return 0.0;
   }
-  let word = x[source_row * HIDDEN_PAIRS + (hidden >> 1u)];
-  let bits = select(word & 0xffffu, word >> 16u, (hidden & 1u) != 0u);
-  return bitcast<f32>(bits << 16u);
 }
 {% endif %}
 {% if vec4 %}
-fn embedding_vec4(source_row: u32, hidden_vec: u32) -> vec4<f32> {
-  if (source_row >= NUM_ROWS) {
-    return vec4<f32>(0.0);
   }
-  let base = source_row * HIDDEN_PAIRS + hidden_vec * 2u;
-  let low = unpack_bf16_pair(x[base]);
-  let high = unpack_bf16_pair(x[base + 1u]);
-  return vec4<f32>(low, high);
-}
 {% endif %}
 {% endif %}
-{% if vec4 and scalarIo %}
-fn load_vec4(index: u32) -> vec4<f32> {
-  return vec4<f32>(x[index], x[index + 1u], x[index + 2u], x[index + 3u]);
 }
 {% endif %}
 {% if combineSubgroups %}
 var<workgroup> sg_partials: array<vec2<f32>, WG>;
@@ -110,17 +391,22 @@ fn reduce_pair(value: vec2<f32>, tid: u32) -> vec2<f32> {
   tr0[tid] = value.x;
   tr1[tid] = value.y;
   workgroupBarrier();
-  var stride: u32 = WG / 2u;
   loop {
     if (stride == 0u) { break; }
-    if (tid < stride) {
       tr0[tid] = tr0[tid] + tr0[tid + stride];
       tr1[tid] = tr1[tid] + tr1[tid + stride];
     }
     stride = stride / 2u;
     workgroupBarrier();
   }
   let reduced = vec2<f32>(tr0[0], tr1[0]);
   workgroupBarrier();
   return reduced;
 }
@@ -134,27 +420,33 @@ fn main(
   @builtin(subgroup_id) sg_id: u32,
   @builtin(num_subgroups) num_sg: u32{% endif %}
 ) {
   let row = wg_id.x + wg_id.y * params.rowStride;
   if (row >= params.rows) {
     return;
   }
   let tid = lid.x;
-{% if packedBf16Embedding %}
-  let source_row = indices[row];
-{% if vec4 %}
   let base = row * HIDDEN_V;
 {% else %}
   let base = row * HIDDEN;
 {% endif %}
-{% elif vec4 and not scalarIo %}
-  let base = row * HIDDEN_V;
-{% else %}
-  let base = row * HIDDEN;
 {% endif %}
 {% if vec4 %}
-{% if scalarIo %}
-  let shift = f32(x[base]);
 {% else %}
   let shift = f32(x[base].x);
 {% endif %}
@@ -164,27 +456,15 @@ fn main(
   var acc = vec2<f32>(0.0, 0.0);
 {% if vec4 %}
-  for (var i = tid; i < HIDDEN_V; i = i + WG) {
-{% if packedBf16Embedding %}
-    let v = embedding_vec4(source_row, i);
-    embedding_out[base + i] = v;
-{% elif scalarIo %}
-    let v = load_vec4(base + i * 4u);
-{% else %}
-    let v = vec4<f32>(x[base + i]);
-{% endif %}
-    let d = v - vec4<f32>(shift);
-    acc.x = acc.x + d.x + d.y + d.z + d.w;
     acc.y = acc.y + dot(d, d);
   }
 {% else %}
   for (var i = tid; i < HIDDEN; i = i + WG) {
-{% if packedBf16Embedding %}
-    let v = embedding_scalar(source_row, i);
-    embedding_out[base + i] = v;
-{% else %}
     let v = f32(x[base + i]);
-{% endif %}
     let d = v - shift;
     acc.x = acc.x + d;
     acc.y = acc.y + d * d;
@@ -204,50 +484,20 @@ fn main(
   }
 {% 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 %}
-    var value = (v - vec4<f32>(row_mean)) * inv * vec4<f32>(scale[i]);
 {% if hasBias %}
-    value = value + vec4<f32>(bias[i]);
 {% endif %}
     y[idx] = {{ vecType }}(value);
   }
-{% 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 %}
     var value = (v - row_mean) * inv * f32(scale[i]);
 {% if hasBias %}
     value = value + f32(bias[i]);
@@ -256,3 +506,4 @@ fn main(
   }
 {% endif %}
 }

+{% set scalarIo = false %}
+{% set packedWidth = packedWidth | default(4) %}
+{% set packedF32 = "vec" ~ packedWidth ~ "<f32>" %}
 {% 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 }}
+{% set affineRowBroadcast = affineRowBroadcast | default(false) %}
+{% set batchRows = batchRows | default(1) %}
+{% set batchLanes = batchLanes | default(0) %}
+{% set xRowShape = xRowShape | default([]) %}
+{% set scaleRowShape = scaleRowShape | default([]) %}
+{% set biasRowShape = biasRowShape | default([]) %}
+{% if affineRowBroadcast %}
+{% macro offset_fn(fn_name, opShape, opRank, op_same, op_numel, outShape, outRank, out_numel) %}
+fn {{ fn_name }}({% if out_numel != 0 and op_numel != 1 %}out_index: u32{% endif %}) -> u32 {
+{% if out_numel == 0 %}
+  return 0u;
+{% elif op_numel == 1 %}
+  return 0u;
+{% elif op_same %}
+  return out_index;
+{% else %}
+  var offset = 0u;
+{% for axis in range(outRank) %}
+{% set op_axis = axis - (outRank - opRank) %}
+{% if op_axis >= 0 and opShape[op_axis] != 1 %}
+{% set c_stride = namespace(value=1) %}
+{% for j in range(axis + 1, outRank) %}
+{% set c_stride.value = c_stride.value * outShape[j] %}
+{% endfor %}
+{% set op_stride = namespace(value=1) %}
+{% for j in range(op_axis + 1, opRank) %}
+{% set op_stride.value = op_stride.value * opShape[j] %}
+{% endfor %}
+{% if c_stride.value == 1 %}
+  let coord{{ axis }} = out_index % {{ outShape[axis] }}u;
+{% else %}
+  let coord{{ axis }} = (out_index / {{ c_stride.value }}u) % {{ outShape[axis] }}u;
 {% endif %}
+{% if op_stride.value == 1 %}
+  offset = offset + coord{{ axis }};
+{% else %}
+  offset = offset + coord{{ axis }} * {{ op_stride.value }}u;
 {% endif %}
 {% endif %}
+{% endfor %}
+  return offset;
+{% endif %}
+}{% endmacro %}
+{% macro broadcast_offset_call(fn_name, opShape, outShape, out_index) %}
+{% set op_numel = namespace(value=1) %}
+{% for d in opShape %}
+{% set op_numel.value = op_numel.value * d %}
+{% endfor %}
+{% set out_numel = namespace(value=1) %}
+{% for d in outShape %}
+{% set out_numel.value = out_numel.value * d %}
+{% endfor %}
+{{ fn_name }}({% if out_numel.value != 0 and op_numel.value != 1 %}{{ out_index }}{% endif %}){% endmacro %}
+{% macro broadcast_offset_fn(fn_name, opShape, opRank, outShape, outRank) %}
+{% set op_numel = namespace(value=1) %}
+{% for d in opShape %}
+{% set op_numel.value = op_numel.value * d %}
+{% endfor %}
+{% set out_numel = namespace(value=1) %}
+{% for d in outShape %}
+{% set out_numel.value = out_numel.value * d %}
+{% endfor %}
+{% set op_same = namespace(value=(opRank == outRank)) %}
+{% if op_same.value %}
+{% for axis in range(outRank) %}
+{% if opShape[axis] != outShape[axis] %}
+{% set op_same.value = false %}
+{% endif %}
+{% endfor %}
+{% endif %}
+{{ offset_fn(fn_name, opShape, opRank, op_same.value, op_numel.value, outShape, outRank, out_numel.value) }}{% endmacro %}
+{{ broadcast_offset_fn("scale_row_offset", scaleRowShape, scaleRowShape | length, xRowShape, xRowShape | length) }}
+{% if hasBias %}
+{{ broadcast_offset_fn("bias_row_offset", biasRowShape, biasRowShape | length, xRowShape, xRowShape | length) }}
+{% endif %}
+{% endif %}
+{% 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 }};
+{% set halfWriteMean = halfWriteMean if halfWriteMean is defined else (writeStats and modeSpec == "layer") %}
+{% set halfWriteInv = halfWriteInv if halfWriteInv is defined else writeStats %}
+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;
+  var local_sum = 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"] %}
+    local_sum = pair_add(local_sum, vec2<f32>(v.{{ component }}, 0.0));
+{% endfor %}
+  }
+{% else %}
+  for (var i = tid; i < HIDDEN; i += WG) {
+    local_sum = pair_add(local_sum, vec2<f32>(load_value(base + i), 0.0));
   }
 {% endif %}
+  let mean = pair_div(reduce_pair(local_sum, {{ reduceArgs }}), f32(HIDDEN));
+  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 halfWriteMean %}
+  if (tid == 0u) { mean_out[row] = mean.x; }
+{% endif %}
+{% if halfWriteInv %}
+  if (tid == 0u) { inv_std_out[row] = inv.x; }
+{% endif %}
+{% 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]));
+{% if modeSpec == "layer" and hasBias %}
+    value = half_stage4(value + vec4<f32>(bias[i]));
+{% endif %}
+    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") }}]));
+{% if modeSpec == "layer" and hasBias %}
+    value = half_stage(value + f32(bias[{{ halfBiasOffset | default("i") }}]));
+{% endif %}
+    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.
+//
+// Shifted moments avoid cancellation from a large common offset; scaling uses
+// inverseSqrt(variance + EPSILON).
+const HIDDEN: u32 = {{ hidden }}u;
+{% if vec4 %}
+const HIDDEN_V: u32 = {{ hiddenVec }}u;
 {% endif %}
+const WG: u32 = {{ wg }}u;
+{% if batchRows > 1 %}
+const ROW_LANES: u32 = {{ batchLanes }}u;
+const ROW_BATCH: u32 = {{ batchRows }}u;
+{% endif %}
+const EPSILON: f32 = {{ epsilon }};
 {% if combineSubgroups %}
 var<workgroup> sg_partials: array<vec2<f32>, WG>;
   tr0[tid] = value.x;
   tr1[tid] = value.y;
   workgroupBarrier();
+  var stride: u32 = {% if batchRows > 1 %}ROW_LANES{% else %}WG{% endif %} / 2u;
   loop {
     if (stride == 0u) { break; }
+    if (tid{% if batchRows > 1 %} % ROW_LANES{% endif %} < stride) {
       tr0[tid] = tr0[tid] + tr0[tid + stride];
       tr1[tid] = tr1[tid] + tr1[tid + stride];
     }
     stride = stride / 2u;
     workgroupBarrier();
   }
+  {% if batchRows > 1 %}
+  let row_base = tid - tid % ROW_LANES;
+  let reduced = vec2<f32>(tr0[row_base], tr1[row_base]);
+{% else %}
   let reduced = vec2<f32>(tr0[0], tr1[0]);
+{% endif %}
   workgroupBarrier();
   return reduced;
 }
   @builtin(subgroup_id) sg_id: u32,
   @builtin(num_subgroups) num_sg: u32{% endif %}
 ) {
+{% if batchRows > 1 %}
+  let row = (wg_id.x + wg_id.y * params.rowStride) * ROW_BATCH + lid.x / ROW_LANES;
+  let row_active = row < params.rows;
+{% else %}
   let row = wg_id.x + wg_id.y * params.rowStride;
   if (row >= params.rows) {
     return;
   }
+{% endif %}
   let tid = lid.x;
+{% if vec4 and not scalarIo %}
   let base = row * HIDDEN_V;
 {% else %}
   let base = row * HIDDEN;
 {% endif %}
+{% if affineRowBroadcast %}
+  let scale_base = {{ broadcast_offset_call("scale_row_offset", scaleRowShape, xRowShape, "row") }} * HIDDEN_V;
+{% if hasBias %}
+  let bias_base = {{ broadcast_offset_call("bias_row_offset", biasRowShape, xRowShape, "row") }} * HIDDEN_V;
 {% endif %}
+{% endif %}
 {% if vec4 %}
+{% if batchRows > 1 %}
+  var shift = 0.0;
+  if (row_active) { shift = f32(x[base].x); }
 {% else %}
   let shift = f32(x[base].x);
 {% endif %}
   var acc = vec2<f32>(0.0, 0.0);
 {% if vec4 %}
+  for (var i = tid{% if batchRows > 1 %} % ROW_LANES{% endif %}; {% if batchRows > 1 %}row_active && {% endif %}i < HIDDEN_V; i = i + {% if batchRows > 1 %}ROW_LANES{% else %}WG{% endif %}) {
+    let v = {{ packedF32 }}(x[base + i]);
+    let d = v - {{ packedF32 }}(shift);
+    acc.x = acc.x + d.x + d.y{% if packedWidth == 4 %} + d.z + d.w{% endif %};
     acc.y = acc.y + dot(d, d);
   }
 {% else %}
   for (var i = tid; i < HIDDEN; i = i + WG) {
     let v = f32(x[base + i]);
     let d = v - shift;
     acc.x = acc.x + d;
     acc.y = acc.y + d * d;
   }
 {% endif %}
 {% if vec4 %}
+  for (var i = tid{% if batchRows > 1 %} % ROW_LANES{% endif %}; {% if batchRows > 1 %}row_active && {% endif %}i < HIDDEN_V; i = i + {% if batchRows > 1 %}ROW_LANES{% else %}WG{% endif %}) {
     let idx = base + i;
+    let v = {{ packedF32 }}(x[idx]);
+    var value = (v - {{ packedF32 }}(row_mean)) * inv * {{ packedF32 }}(scale[{% if affineRowBroadcast %}scale_base + {% endif %}i]);
 {% if hasBias %}
+    value = value + {{ packedF32 }}(bias[{% if affineRowBroadcast %}bias_base + {% endif %}i]);
 {% endif %}
     y[idx] = {{ vecType }}(value);
   }
 {% else %}
   for (var i = tid; i < HIDDEN; i = i + WG) {
     let idx = base + i;
     let v = f32(x[idx]);
     var value = (v - row_mean) * inv * f32(scale[i]);
 {% if hasBias %}
     value = value + f32(bias[i]);
   }
 {% endif %}
 }
+{% endif %}

build/webgpu/test.json CHANGED Viewed

@@ -4,6 +4,157 @@
     "onnx_backend_layer_normalization_input_x": [1.764052391052246, 0.40015721321105957, 0.978738009929657, 2.2408931255340576, 1.8675580024719238, -0.9772778749465942, 0.9500884413719177, -0.15135720372200012, -0.10321885347366333, 0.4105985164642334, 0.14404356479644775, 1.4542734622955322, 0.7610377073287964, 0.12167501449584961, 0.44386324286460876, 0.3336743414402008, 1.4940791130065918, -0.2051582634449005, 0.3130677044391632, -0.8540957570075989, -2.5529897212982178, 0.653618574142456, 0.8644362092018127, -0.7421650290489197, 2.269754648208618, -1.4543657302856445, 0.04575851559638977, -0.18718385696411133, 1.5327792167663574, 1.4693588018417358, 0.154947429895401, 0.37816253304481506, -0.8877857327461243, -1.980796456336975, -0.34791216254234314, 0.15634897351264954, 1.2302906513214111, 1.202379822731018, -0.38732680678367615, -0.302302747964859, -1.0485529899597168, -1.420017957687378, -1.7062702178955078, 1.950775384902954, -0.5096521973609924, -0.4380742907524109, -1.2527953386306763, 0.7774903774261475, -1.6138978004455566, -0.21274028718471527, -0.8954665660858154, 0.38690251111984253, -0.5108051300048828, -1.18063223361969, -0.02818222902715206, 0.4283318817615509, 0.06651721894741058, 0.30247190594673157, -0.6343221068382263, -0.3627411723136902, -0.6724604368209839, -0.35955315828323364, -0.8131462931632996, -1.7262825965881348, 0.17742614448070526, -0.4017809331417084, -1.630198359489441, 0.46278226375579834, -0.9072983860969543, 0.05194539576768875, 0.7290905714035034, 0.12898291647434235, 1.1394007205963135, -1.234825849533081, 0.4023416340351105, -0.6848101019859314, -0.8707971572875977, -0.5788496732711792, -0.3115525245666504, 0.056165341287851334, -1.1651498079299927, 0.9008265137672424, 0.4656624495983124, -1.5362436771392822, 1.4882521629333496, 1.895889163017273, 1.1787796020507812, -0.1799248307943344, -1.0707526206970215, 1.0544517040252686, -0.4031769335269928, 1.222445011138916, 0.2082749754190445, 0.9766390323638916, 0.3563663959503174, 0.7065731883049011, 0.01050002034753561, 1.7858705520629883, 0.12691208720207214, 0.4019893705844879, 1.8831506967544556, -1.3477590084075928, -1.2704850435256958, 0.969396710395813, -1.1731233596801758, 1.9436211585998535, -0.4136189818382263, -0.747454822063446, 1.922942042350769, 1.4805147647857666, 1.8675589561462402, 0.9060446619987488, -0.8612256646156311, 1.910064935684204, -0.26800337433815, 0.8024563789367676, 0.9472519755363464, -0.15501008927822113, 0.6140793561935425, 0.922206699848175]
   },
   "cases": [
     {
       "name": "subgroup_vec4_stats_no_bias_2x512",
       "attrs": { "epsilon": 0.00001, "axis": -1 },
@@ -1783,6 +1934,373 @@
         "mean": { "dtype": "float32", "shape": [2, 1], "tolerance": 0.000001 },
         "invStdDev": { "dtype": "float32", "shape": [2, 1], "tolerance": 0.00001 }
       }
     }
   ]
 }

     "onnx_backend_layer_normalization_input_x": [1.764052391052246, 0.40015721321105957, 0.978738009929657, 2.2408931255340576, 1.8675580024719238, -0.9772778749465942, 0.9500884413719177, -0.15135720372200012, -0.10321885347366333, 0.4105985164642334, 0.14404356479644775, 1.4542734622955322, 0.7610377073287964, 0.12167501449584961, 0.44386324286460876, 0.3336743414402008, 1.4940791130065918, -0.2051582634449005, 0.3130677044391632, -0.8540957570075989, -2.5529897212982178, 0.653618574142456, 0.8644362092018127, -0.7421650290489197, 2.269754648208618, -1.4543657302856445, 0.04575851559638977, -0.18718385696411133, 1.5327792167663574, 1.4693588018417358, 0.154947429895401, 0.37816253304481506, -0.8877857327461243, -1.980796456336975, -0.34791216254234314, 0.15634897351264954, 1.2302906513214111, 1.202379822731018, -0.38732680678367615, -0.302302747964859, -1.0485529899597168, -1.420017957687378, -1.7062702178955078, 1.950775384902954, -0.5096521973609924, -0.4380742907524109, -1.2527953386306763, 0.7774903774261475, -1.6138978004455566, -0.21274028718471527, -0.8954665660858154, 0.38690251111984253, -0.5108051300048828, -1.18063223361969, -0.02818222902715206, 0.4283318817615509, 0.06651721894741058, 0.30247190594673157, -0.6343221068382263, -0.3627411723136902, -0.6724604368209839, -0.35955315828323364, -0.8131462931632996, -1.7262825965881348, 0.17742614448070526, -0.4017809331417084, -1.630198359489441, 0.46278226375579834, -0.9072983860969543, 0.05194539576768875, 0.7290905714035034, 0.12898291647434235, 1.1394007205963135, -1.234825849533081, 0.4023416340351105, -0.6848101019859314, -0.8707971572875977, -0.5788496732711792, -0.3115525245666504, 0.056165341287851334, -1.1651498079299927, 0.9008265137672424, 0.4656624495983124, -1.5362436771392822, 1.4882521629333496, 1.895889163017273, 1.1787796020507812, -0.1799248307943344, -1.0707526206970215, 1.0544517040252686, -0.4031769335269928, 1.222445011138916, 0.2082749754190445, 0.9766390323638916, 0.3563663959503174, 0.7065731883049011, 0.01050002034753561, 1.7858705520629883, 0.12691208720207214, 0.4019893705844879, 1.8831506967544556, -1.3477590084075928, -1.2704850435256958, 0.969396710395813, -1.1731233596801758, 1.9436211585998535, -0.4136189818382263, -0.747454822063446, 1.922942042350769, 1.4805147647857666, 1.8675589561462402, 0.9060446619987488, -0.8612256646156311, 1.910064935684204, -0.26800337433815, 0.8024563789367676, 0.9472519755363464, -0.15501008927822113, 0.6140793561935425, 0.922206699848175]
   },
   "cases": [
+    {
+      "name": "half_normalized_midpoint_vec4",
+      "attrs": { "axis": -1, "epsilon": 0.00001 },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [1, 12],
+          "data": {
+            "kind": "values",
+            "values": [0.103759765625, -0.128662109375, 0.271484375, 0.06829833984375, -0.00689697265625, 0.10546875, 0.568359375, -0.1162109375, 1.0185546875, 0.52392578125, -0.31591796875, 0.36279296875]
+          }
+        },
+        "scale": {
+          "dtype": "float16",
+          "shape": [12],
+          "data": {
+            "kind": "values",
+            "values": [0.9931640625, 0.82421875, 1.1064453125, 0.751953125, 1.2080078125, 1.0224609375, 0.9033203125, 1.228515625, 1.1533203125, 0.77001953125, 1.04296875, 1.1240234375]
+          }
+        },
+        "b": {
+          "dtype": "float16",
+          "shape": [12],
+          "data": {
+            "kind": "values",
+            "values": [0.115234375, -0.0238189697265625, -0.06988525390625, -0.04443359375, -0.1641845703125, 0.0662841796875, 0.00609588623046875, 0.0999755859375, 0.2476806640625, -0.0048675537109375, 0.2154541015625, -0.18359375]
+          }
+        }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "float16",
+          "shape": [1, 12],
+          "data": {
+            "kind": "values",
+            "values": [-0.16845703125, -0.80224609375, 0.139892578125, -0.3349609375, -0.8876953125, -0.2208251953125, 0.9365234375, -1.017578125, 2.908203125, 0.69189453125, -1.322265625, 0.3203125]
+          },
+          "tolerance": 0.000002
+        }
+      },
+      "provenance": {
+        "source": "ONNX LayerNormalization typed normalization and scale stages",
+        "notes": "Half-normalized midpoint followed by affine rounding. Pinned with independent float64 statistics and explicit half stages; a float32-only reduction can cross the midpoint."
+      }
+    },
+    {
+      "name": "onnx17_half_affine_stages_scalar_broadcast",
+      "attrs": { "axis": -1, "epsilon": 0, "stash_type": 1 },
+      "inputs": {
+        "x": { "dtype": "float16", "shape": [1, 3], "data": { "kind": "values", "values": [0.3, 1.7, -0.9] } },
+        "scale": { "dtype": "float16", "shape": [], "data": { "kind": "values", "values": [1.3] } },
+        "b": { "dtype": "float16", "shape": [], "data": { "kind": "values", "values": [0.07] } }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "float16",
+          "shape": [1, 3],
+          "tolerance": 0.000002,
+          "data": { "kind": "values", "values": [-0.0115966796875, 1.701171875, -1.4794921875] }
+        }
+      },
+      "provenance": {
+        "source": "ONNX LayerNormalization-17 typed affine stages",
+        "notes": "Scalar scale and bias exercise both rank-zero broadcast offsets in the generic half shader. Exact half values independently match explicit primitives on ONNX Runtime's CPU provider and typed PyTorch CPU stages."
+      }
+    },
+    {
+      "name": "onnx17_half_affine_stages_scalar",
+      "attrs": { "axis": -1, "epsilon": 0, "stash_type": 1 },
+      "inputs": {
+        "x": { "dtype": "float16", "shape": [1, 3], "data": { "kind": "values", "values": [0.3, 1.7, -0.9] } },
+        "scale": { "dtype": "float16", "shape": [3], "data": { "kind": "values", "values": [1.3, 2.1, -0.7] } },
+        "b": { "dtype": "float16", "shape": [3], "data": { "kind": "values", "values": [0.07, -0.3, 0.9] } }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "float16",
+          "shape": [1, 3],
+          "tolerance": 0.000002,
+          "data": { "kind": "values", "values": [-0.0115966796875, 2.333984375, 1.734375] }
+        }
+      },
+      "provenance": {
+        "source": "ONNX LayerNormalization-17 function: Cast Normalized to T, then Mul and Add in T",
+        "notes": "Exact half stages independently checked with ONNX Runtime's CPU primitive graph and PyTorch's CPU explicit stages. Native fused LayerNorm differs at the first output; the spec is authoritative."
+      }
+    },
+    {
+      "name": "onnx17_half_affine_stages_vec4",
+      "attrs": { "axis": -1, "epsilon": 0, "stash_type": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [1, 12],
+          "data": { "kind": "values", "values": [0.3, 1.7, -0.9, 0.3, 1.7, -0.9, 0.3, 1.7, -0.9, 0.3, 1.7, -0.9] }
+        },
+        "scale": {
+          "dtype": "float16",
+          "shape": [12],
+          "data": { "kind": "values", "values": [1.3, 2.1, -0.7, 1.3, 2.1, -0.7, 1.3, 2.1, -0.7, 1.3, 2.1, -0.7] }
+        },
+        "b": {
+          "dtype": "float16",
+          "shape": [12],
+          "data": { "kind": "values", "values": [0.07, -0.3, 0.9, 0.07, -0.3, 0.9, 0.07, -0.3, 0.9, 0.07, -0.3, 0.9] }
+        }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "float16",
+          "shape": [1, 12],
+          "tolerance": 0.000002,
+          "data": {
+            "kind": "values",
+            "values": [-0.0115966796875, 2.333984375, 1.734375, -0.0115966796875, 2.333984375, 1.734375, -0.0115966796875, 2.333984375, 1.734375, -0.0115966796875, 2.333984375, 1.734375]
+          }
+        }
+      },
+      "provenance": {
+        "source": "ONNX LayerNormalization-17 typed affine stages",
+        "notes": "Four repetitions of the 3-element pattern from this file's onnx17_half_affine_stages_scalar case (12 elements total) load in vec4 groups; float16 scale and bias check half-precision product rounding and addition, with exact expected output values."
+      }
+    },
+    {
+      "name": "onnx17_half_affine_stages_suffix_broadcast",
+      "attrs": { "axis": 1, "epsilon": 0, "stash_type": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [2, 2, 3],
+          "data": { "kind": "values", "values": [0.3, 1.7, -0.9, 0.3, 1.7, -0.9, 0.3, 1.7, -0.9, 0.3, 1.7, -0.9] }
+        },
+        "scale": { "dtype": "float16", "shape": [3], "data": { "kind": "values", "values": [1.3, 2.1, -0.7] } },
+        "b": { "dtype": "float16", "shape": [3], "data": { "kind": "values", "values": [0.07, -0.3, 0.9] } }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "float16",
+          "shape": [2, 2, 3],
+          "tolerance": 0.000002,
+          "data": {
+            "kind": "values",
+            "values": [-0.0115966796875, 2.333984375, 1.734375, -0.0115966796875, 2.333984375, 1.734375, -0.0115966796875, 2.333984375, 1.734375, -0.0115966796875, 2.333984375, 1.734375]
+          }
+        }
+      },
+      "provenance": {
+        "source": "ONNX LayerNormalization-17 typed affine stages",
+        "notes": "Suffix-axis normalization with broadcast scale/bias exercises the generic shader on two independent rows."
+      }
+    },
     {
       "name": "subgroup_vec4_stats_no_bias_2x512",
       "attrs": { "epsilon": 0.00001, "axis": -1 },
         "mean": { "dtype": "float32", "shape": [2, 1], "tolerance": 0.000001 },
         "invStdDev": { "dtype": "float32", "shape": [2, 1], "tolerance": 0.00001 }
       }
+    },
+    {
+      "name": "broadcast_affine_vec4_full",
+      "attrs": { "axis": -1, "epsilon": 0.00001 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 64],
+          "data": { "kind": "cycle", "values": [-0.25, 0.125, 0.5, -0.0625, 1.125, -0.75, 2.0, -1.25, 0.25] }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [3, 64],
+          "data": { "kind": "cycle", "values": [0.25, 0.5, 1.0, 1.5, -0.5] }
+        },
+        "b": {
+          "dtype": "float32",
+          "shape": [3, 64],
+          "data": { "kind": "cycle", "values": [-0.125, 0.0, 0.25, 1.5, -0.75, 0.5, 0.125] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [3, 64], "tolerance": 0.000002, "relTolerance": 0.00001 } },
+      "provenance": {
+        "notes": "Packed broadcast-affine rows: independent outer indexing, shifted moments, and partial row groups are checked against the CPU reference without changing tolerance."
+      }
+    },
+    {
+      "name": "broadcast_affine_vec4_shifted",
+      "attrs": { "axis": -1, "epsilon": 0.00001 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 64],
+          "data": {
+            "kind": "cycle",
+            "values": [8191.75, 8192.125, 8192.5, 8191.9375, 8193.125, 8191.25, 8194.0, 8190.75, 8192.25]
+          }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [3, 64],
+          "data": { "kind": "cycle", "values": [0.25, 0.5, 1.0, 1.5, -0.5] }
+        },
+        "b": {
+          "dtype": "float32",
+          "shape": [3, 64],
+          "data": { "kind": "cycle", "values": [-0.125, 0.0, 0.25, 1.5, -0.75, 0.5, 0.125] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [3, 64], "tolerance": 0.000002, "relTolerance": 0.00001 } },
+      "provenance": {
+        "notes": "Packed broadcast-affine rows: independent outer indexing, shifted moments, and partial row groups are checked against the CPU reference without changing tolerance."
+      }
+    },
+    {
+      "name": "broadcast_affine_vec4_outer_mixed",
+      "attrs": { "axis": -1, "epsilon": 0.00001 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 3, 4, 64],
+          "data": { "kind": "cycle", "values": [-0.25, 0.125, 0.5, -0.0625, 1.125, -0.75, 2.0, -1.25, 0.25] }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [2, 1, 4, 64],
+          "data": { "kind": "cycle", "values": [0.25, 0.5, 1.0, 1.5, -0.5] }
+        },
+        "b": {
+          "dtype": "float32",
+          "shape": [1, 3, 1, 64],
+          "data": { "kind": "cycle", "values": [-0.125, 0.0, 0.25, 1.5, -0.75, 0.5, 0.125] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 3, 4, 64], "tolerance": 0.000002, "relTolerance": 0.00001 } },
+      "provenance": {
+        "notes": "Packed broadcast-affine rows: independent outer indexing, shifted moments, and partial row groups are checked against the CPU reference without changing tolerance."
+      }
+    },
+    {
+      "name": "broadcast_affine_vec4_shared_scale_full_bias",
+      "attrs": { "axis": -1, "epsilon": 0.00001 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 12],
+          "data": { "kind": "cycle", "values": [-0.25, 0.125, 0.5, -0.0625, 1.125, -0.75, 2.0, -1.25, 0.25] }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [12],
+          "data": { "kind": "cycle", "values": [0.25, 0.5, 1.0, 1.5, -0.5] }
+        },
+        "b": {
+          "dtype": "float32",
+          "shape": [3, 12],
+          "data": { "kind": "cycle", "values": [-0.125, 0.0, 0.25, 1.5, -0.75, 0.5, 0.125] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [3, 12], "tolerance": 0.000002, "relTolerance": 0.00001 } },
+      "provenance": {
+        "notes": "Packed broadcast-affine rows: independent outer indexing, shifted moments, and partial row groups are checked against the CPU reference without changing tolerance."
+      }
+    },
+    {
+      "name": "broadcast_affine_vec4_full_scale_no_bias",
+      "attrs": { "axis": -1, "epsilon": 0.00001 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 1024],
+          "data": { "kind": "cycle", "values": [-0.25, 0.125, 0.5, -0.0625, 1.125, -0.75, 2.0, -1.25, 0.25] }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [3, 1024],
+          "data": { "kind": "cycle", "values": [0.25, 0.5, 1.0, 1.5, -0.5] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [3, 1024], "tolerance": 0.000002, "relTolerance": 0.00001 } },
+      "provenance": {
+        "notes": "Packed broadcast-affine rows: independent outer indexing, shifted moments, and partial row groups are checked against the CPU reference without changing tolerance."
+      }
+    },
+    {
+      "name": "broadcast_affine_vec4_right_aligned",
+      "attrs": { "axis": -1, "epsilon": 0.00001 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 3, 8],
+          "data": { "kind": "cycle", "values": [-0.25, 0.125, 0.5, -0.0625, 1.125, -0.75, 2.0, -1.25, 0.25] }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [3, 8],
+          "data": { "kind": "cycle", "values": [0.25, 0.5, 1.0, 1.5, -0.5] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 3, 8], "tolerance": 0.000002, "relTolerance": 0.00001 } },
+      "provenance": {
+        "notes": "Packed broadcast-affine rows: independent outer indexing, shifted moments, and partial row groups are checked against the CPU reference without changing tolerance."
+      }
+    },
+    {
+      "name": "broadcast_affine_vec4_mixed_shifted",
+      "attrs": { "axis": -1, "epsilon": 0.00001 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 3, 8],
+          "data": {
+            "kind": "cycle",
+            "values": [39999.75, 40000.125, 40000.5, 39999.9375, 40001.125, 39999.25, 40002.0, 39998.75, 40000.25]
+          }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [2, 1, 8],
+          "data": { "kind": "cycle", "values": [0.25, 0.5, 1.0, 1.5, -0.5] }
+        },
+        "b": {
+          "dtype": "float32",
+          "shape": [1, 3, 8],
+          "data": { "kind": "cycle", "values": [-0.125, 0.0, 0.25, 1.5, -0.75, 0.5, 0.125] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 3, 8], "tolerance": 0.000002, "relTolerance": 0.00001 } },
+      "provenance": {
+        "notes": "Packed broadcast-affine rows: independent outer indexing, shifted moments, and partial row groups are checked against the CPU reference without changing tolerance."
+      }
+    },
+    {
+      "name": "broadcast_affine_vec4_full_scale_shared_bias",
+      "attrs": { "axis": -1, "epsilon": 0.00001 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 3, 8],
+          "data": { "kind": "cycle", "values": [-0.25, 0.125, 0.5, -0.0625, 1.125, -0.75, 2.0, -1.25, 0.25] }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [2, 3, 8],
+          "data": { "kind": "cycle", "values": [0.25, 0.5, 1.0, 1.5, -0.5] }
+        },
+        "b": {
+          "dtype": "float32",
+          "shape": [8],
+          "data": { "kind": "cycle", "values": [-0.125, 0.0, 0.25, 1.5, -0.75, 0.5, 0.125] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 3, 8], "tolerance": 0.000002, "relTolerance": 0.00001 } },
+      "provenance": {
+        "notes": "Packed broadcast-affine rows: independent outer indexing, shifted moments, and partial row groups are checked against the CPU reference without changing tolerance."
+      }
+    },
+    {
+      "name": "broadcast_affine_vec4_zero_variance",
+      "attrs": { "axis": -1, "epsilon": 0.00001 },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [2, 3, 8], "data": { "kind": "constant", "value": 40000.0 } },
+        "scale": {
+          "dtype": "float32",
+          "shape": [2, 1, 8],
+          "data": { "kind": "cycle", "values": [0.25, 0.5, 1.0, 1.5, -0.5] }
+        },
+        "b": {
+          "dtype": "float32",
+          "shape": [1, 3, 8],
+          "data": { "kind": "cycle", "values": [-0.125, 0.0, 0.25, 1.5, -0.75, 0.5, 0.125] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 3, 8], "tolerance": 0.000002, "relTolerance": 0.00001 } },
+      "provenance": {
+        "notes": "Packed broadcast-affine rows: independent outer indexing, shifted moments, and partial row groups are checked against the CPU reference without changing tolerance."
+      }
+    },
+    {
+      "name": "broadcast_affine_vec4_batch_tail",
+      "attrs": { "axis": -1, "epsilon": 0.00001 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [17, 64],
+          "data": { "kind": "cycle", "values": [-0.25, 0.125, 0.5, -0.0625, 1.125, -0.75, 2.0, -1.25, 0.25] }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [17, 64],
+          "data": { "kind": "cycle", "values": [0.25, 0.5, 1.0, 1.5, -0.5] }
+        },
+        "b": {
+          "dtype": "float32",
+          "shape": [17, 64],
+          "data": { "kind": "cycle", "values": [-0.125, 0.0, 0.25, 1.5, -0.75, 0.5, 0.125] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [17, 64], "tolerance": 0.000002, "relTolerance": 0.00001 } },
+      "provenance": {
+        "notes": "Packed broadcast-affine rows: independent outer indexing, shifted moments, and partial row groups are checked against the CPU reference without changing tolerance."
+      }
+    },
+    {
+      "name": "broadcast_affine_vec4_batch_tail_no_bias",
+      "attrs": { "axis": -1, "epsilon": 0.00001 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [17, 64],
+          "data": { "kind": "cycle", "values": [-0.25, 0.125, 0.5, -0.0625, 1.125, -0.75, 2.0, -1.25, 0.25] }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [17, 64],
+          "data": { "kind": "cycle", "values": [0.25, 0.5, 1.0, 1.5, -0.5] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [17, 64], "tolerance": 0.000002, "relTolerance": 0.00001 } },
+      "provenance": {
+        "notes": "Packed broadcast-affine rows: independent outer indexing, shifted moments, and partial row groups are checked against the CPU reference without changing tolerance."
+      }
+    },
+    {
+      "name": "broadcast_affine_vec4_wide_full_scale",
+      "attrs": { "axis": -1, "epsilon": 0.00001 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 4096],
+          "data": { "kind": "cycle", "values": [-0.25, 0.125, 0.5, -0.0625, 1.125, -0.75, 2.0, -1.25, 0.25] }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [3, 4096],
+          "data": { "kind": "cycle", "values": [0.25, 0.5, 1.0, 1.5, -0.5] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [3, 4096], "tolerance": 0.000002, "relTolerance": 0.00001 } },
+      "provenance": {
+        "notes": "Packed broadcast-affine rows: independent outer indexing, shifted moments, and partial row groups are checked against the CPU reference without changing tolerance."
+      }
+    },
+    {
+      "name": "f32_last_axis_vec2_scale_only_h126",
+      "attrs": { "axis": 1, "epsilon": 0.00001 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 126],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.13 }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [126],
+          "data": { "kind": "cycle", "values": [0.5, -1.0, 1.25, 2.0, -0.25] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [3, 126], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "f32_last_axis_vec2_scale_only_h514",
+      "attrs": { "axis": 1, "epsilon": 0.00001 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 514],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.13 }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [514],
+          "data": { "kind": "cycle", "values": [0.5, -1.0, 1.25, 2.0, -0.25] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [3, 514], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "f32_last_axis_vec2_scale_only_h1022",
+      "attrs": { "axis": 1, "epsilon": 0.00001 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 1022],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.13 }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [1022],
+          "data": { "kind": "cycle", "values": [0.5, -1.0, 1.25, 2.0, -0.25] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [3, 1022], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "f32_last_axis_vec2_scale_only_h4098",
+      "attrs": { "axis": 1, "epsilon": 0.00001 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 4098],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.13 }
+        },
+        "scale": {
+          "dtype": "float32",
+          "shape": [4098],
+          "data": { "kind": "cycle", "values": [0.5, -1.0, 1.25, 2.0, -0.25] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [3, 4098], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "f32_last_axis_vec2_shifted_bias_h6",
+      "attrs": { "axis": -1, "epsilon": 0.00001 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 6],
+          "data": {
+            "kind": "values",
+            "values": [7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 10000.0, 10000.125, 9999.875, 10000.25, 9999.75, 10000.0, -8.0, 2.0, 0.0, 7.0, -4.0, 3.0]
+          }
+        },
+        "scale": { "dtype": "float32", "shape": [6], "data": { "kind": "cycle", "values": [0.5, -1.0, 1.25] } },
+        "b": { "dtype": "float32", "shape": [6], "data": { "kind": "cycle", "values": [-0.25, 1.0, 0.75] } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [3, 6], "tolerance": 0.00001 } }
     }
   ]
 }