sync 2e7068faf55e

Browse files

Files changed (8) hide show

README.md +73 -0
build/webgpu/bench.json +73 -0
build/webgpu/bias-softmax-longrow-normalize.wgsl.jinja +59 -0
build/webgpu/bias-softmax-longrow-stats.wgsl.jinja +221 -0
build/webgpu/bias-softmax.wgsl.jinja +281 -0
build/webgpu/manifest.json +279 -0
build/webgpu/metadata.json +20 -0
build/webgpu/test.json +442 -0

README.md CHANGED Viewed

@@ -1,3 +1,76 @@
 ---
 license: apache-2.0
 ---

 ---
+library_name: kernels
 license: apache-2.0
+tags:
+- kernel
+- webgpu
+- wgsl
 ---
+# com.microsoft.BiasSoftmax
+`com.microsoft`  ·  ONNX Runtime contrib operator  ·  contrib since_version 1
+## Description
+Computes `softmax(data + bias)` over the flattened suffix beginning at `axis`. The required `is_inner_broadcast` attribute selects how bias rows are reused: consecutive groups for inner broadcast or cyclic groups for outer broadcast. This specializes the `softmax(scores + additive_mask)` pattern used by transformer attention. Float16 and float32 are supported; the schema's double type is not.
+See the [ONNX Runtime `BiasSoftmax` contrib-operator spec](https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.BiasSoftmax) for the reference semantics.
+## Inputs
+| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `data` | `data` | `T` | — | — | The input data tensor. | required |
+| `bias` | `bias` | `T` | — | — | The bias (or additive mask) tensor. Its element count must be an integral number of flattened softmax rows and that row count must divide the data row count. | required |
+## Outputs
+| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `output` | `output` | `T` | same as `data` | same as `data` | The output tensor; same shape as data. | required |
+## Attributes
+Attributes and default values (overridable per request):
+| Attribute | Default | Description |
+| --- | --- | --- |
+| `axis` | `1` | The axis from which softmax is applied; dimensions from `axis` onward are included in the softmax reduction. |
+| `is_inner_broadcast` | — | When 1, bias is broadcast across dimensions from `broadcast_axis` to `axis-1`; when 0, bias is broadcast across dimensions 0 to `broadcast_axis-1`. |
+## Type constraints
+| Variable | Allowed dtypes |
+| --- | --- |
+| `T` | `float32`, `float16` |
+## Files
+- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, provenance)
+- [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
+- [`test.json`](build/webgpu/test.json) — correctness cases
+- [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
+- [`bias-softmax-longrow-normalize.wgsl.jinja`](build/webgpu/bias-softmax-longrow-normalize.wgsl.jinja)
+- [`bias-softmax-longrow-stats.wgsl.jinja`](build/webgpu/bias-softmax-longrow-stats.wgsl.jinja)
+- [`bias-softmax.wgsl.jinja`](build/webgpu/bias-softmax.wgsl.jinja)
+## Use with `@huggingface/kernels`
+The loader derives every required output's shape and logical dtype from the manifest contract and this call.
+It then allocates the result tensors automatically.
+The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+Replace each `*Data` placeholder with a typed array containing the corresponding input data.
+```js
+import { getKernel } from "@huggingface/kernels";
+const kernel = await getKernel("webgpu-kernels/com.microsoft.BiasSoftmax", { version: 1 });
+const { output } = await kernel({
+  data: { data: dataData, shape: [1, 2, 2] },
+  bias: { data: biasData, shape: [1, 2, 2] },
+}, {
+  attrs: { is_inner_broadcast: 1 },
+});
+```

build/webgpu/bench.json ADDED Viewed

	@@ -0,0 +1,73 @@

+{
+  "op": "com.microsoft.BiasSoftmax",
+  "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256], "BLOCK_COLS": [1024, 2048, 4096, 8192] },
+  "cases": [
+    {
+      "name": "biassoftmax-f32-attn-32x512x512",
+      "preset": "smoke",
+      "vars": { "dtype": "float32" },
+      "attrs": { "axis": 2, "is_inner_broadcast": 0 },
+      "inputs": {
+        "data": { "shape": [32, 512, 512], "dtype": "float32", "dist": "normal", "seed": 510, "scale": 2 },
+        "bias": { "shape": [1, 512, 512], "dtype": "float32", "dist": "normal", "seed": 511, "scale": 1 }
+      },
+      "outputs": { "output": { "shape": [32, 512, 512], "dtype": "float32" } },
+      "bench": {
+        "primary": true,
+        "metrics": [{ "type": "bandwidth", "value": "(32 * 512 * 512 * 2 + 512 * 512) * dtypeBytes(args.dtype)" }]
+      }
+    },
+    {
+      "name": "biassoftmax-f32-suffix-64x16x256",
+      "preset": "stress",
+      "vars": { "dtype": "float32" },
+      "attrs": { "axis": 1, "is_inner_broadcast": 0 },
+      "inputs": {
+        "data": { "shape": [64, 16, 256], "dtype": "float32", "dist": "normal", "seed": 512, "scale": 2 },
+        "bias": { "shape": [1, 16, 256], "dtype": "float32", "dist": "normal", "seed": 513, "scale": 1 }
+      },
+      "outputs": { "output": { "shape": [64, 16, 256], "dtype": "float32" } },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "(64 * 16 * 256 * 2 + 16 * 256) * dtypeBytes(args.dtype)" }]
+      }
+    },
+    {
+      "name": "biassoftmax-f16-attn-32x512x512",
+      "preset": "model",
+      "vars": { "dtype": "float16" },
+      "attrs": { "axis": 2, "is_inner_broadcast": 0 },
+      "inputs": {
+        "data": { "shape": [32, 512, 512], "dtype": "float16", "dist": "normal", "seed": 514, "scale": 2 },
+        "bias": { "shape": [1, 512, 512], "dtype": "float16", "dist": "normal", "seed": 515, "scale": 1 }
+      },
+      "outputs": { "output": { "shape": [32, 512, 512], "dtype": "float16" } },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "(32 * 512 * 512 * 2 + 512 * 512) * dtypeBytes(args.dtype)" }]
+      }
+    },
+    {
+      "name": "biassoftmax-f32-lowocc-8x524288-axis1",
+      "preset": "stress",
+      "vars": { "dtype": "float32" },
+      "attrs": { "axis": 1, "is_inner_broadcast": 0 },
+      "inputs": {
+        "data": { "shape": [8, 524288], "dtype": "float32", "dist": "normal", "seed": 530, "scale": 2 },
+        "bias": { "shape": [1, 524288], "dtype": "float32", "dist": "normal", "seed": 531, "scale": 1 }
+      },
+      "outputs": { "output": { "shape": [8, 524288], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "(8 * 524288 * 2 + 524288) * dtypeBytes(args.dtype)" }] }
+    },
+    {
+      "name": "biassoftmax-f32-launchbound-262144x2-axis1",
+      "preset": "stress",
+      "vars": { "dtype": "float32" },
+      "attrs": { "axis": 1, "is_inner_broadcast": 0 },
+      "inputs": {
+        "data": { "shape": [262144, 2], "dtype": "float32", "dist": "normal", "seed": 532, "scale": 2 },
+        "bias": { "shape": [1, 2], "dtype": "float32", "dist": "normal", "seed": 533, "scale": 1 }
+      },
+      "outputs": { "output": { "shape": [262144, 2], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "(262144 * 2 * 2 + 2) * dtypeBytes(args.dtype)" }] }
+    }
+  ]
+}

build/webgpu/bias-softmax-longrow-normalize.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,59 @@

+// Parallel normalization for split long rows. One workgroup handles each
+// BLOCK_COLS slice and reuses the finished per-row statistics.
+{% if usesF16 %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
+const BLOCK_COLS: u32 = {{ tunables.BLOCK_COLS }}u;
+fn bias_index(index: u32) -> u32 {
+  let row = index / params.blockSize;
+  let column = index % params.blockSize;
+{% if source.isInnerBroadcast %}
+  return (row / {{ source.innerRepeat }}u) * params.blockSize + column;
+{% else %}
+  return (row % {{ source.biasBlockCount }}u) * params.blockSize + column;
+{% endif %}
+}
+// FLT_MAX, not -inf, as the online (m, d) accumulator init: merges must keep
+// `m - m` finite so an empty lane / all--inf row contributes the exact
+// accumulator identity (m, d) = (-FLT_MAX, 0). Operator epilogues interpret
+// a zero final denominator according to their public semantics. Using -inf
+// here changes +inf-row behavior.
+const FLT_MAX: f32 = 3.4028234663852886e38;
+fn is_finite_f32(value: f32) -> bool {
+  return select(false, value <= FLT_MAX, value >= -FLT_MAX);
+}
+// x - m that is exactly 0 when x equals a finite m, so exp(shifted) == 1
+// exactly at the row max. `x - x` on an infinite max is a legal fast-math
+// fold to 0, which would silently turn +inf rows finite — the explicit
+// equality test keeps the NaN propagation of the serial kernels.
+fn shifted_value(value: f32, maxValue: f32) -> f32 {
+  let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
+  return select(value - maxValue, 0.0, equalFiniteMax);
+}
+fn exp_shift(value: f32, maxValue: f32) -> f32 {
+  return exp(shifted_value(value, maxValue));
+}
+@compute @workgroup_size(WG)
+fn main(@builtin(workgroup_id) wid: vec3<u32>,
+        @builtin(local_invocation_id) lid: vec3<u32>) {
+  let block = wid.x;
+  let row = wid.y;
+  let rowBase = row * params.blockSize;
+  let colStart = block * BLOCK_COLS;
+  let colEnd = min(params.blockSize, colStart + BLOCK_COLS);
+  let m = rowMax[row];
+  let d = rowSum[row];
+  let inv = select(0.0, 1.0 / d, d > 0.0);
+  for (var col = colStart + lid.x; col < colEnd; col += WG) {
+    let i = rowBase + col;
+    let v = f32(data[i]) + f32(bias[bias_index(i)]);
+    output[i] = {{ scalar }}(exp_shift(v, m) * inv);
+  }
+}

build/webgpu/bias-softmax-longrow-stats.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,221 @@

+// Split statistics for very long contiguous BiasSoftmax rows. The block stage
+// computes independent online (maximum, sum-exp) pairs over BLOCK_COLS. The row
+// stage deterministically folds them into one result per row.
+{% if source.stage == "block" and usesF16 %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
+{% if source.stage == "block" %}
+const BLOCK_COLS: u32 = {{ tunables.BLOCK_COLS }}u;
+fn bias_index(index: u32) -> u32 {
+  let row = index / params.blockSize;
+  let column = index % params.blockSize;
+{% if source.isInnerBroadcast | default(false) %}
+  return (row / {{ source.innerRepeat | default(1) }}u) * params.blockSize + column;
+{% else %}
+  return (row % {{ source.biasBlockCount | default(1) }}u) * params.blockSize + column;
+{% endif %}
+}
+{% endif %}
+// FLT_MAX, not -inf, as the online (m, d) accumulator init: merges must keep
+// `m - m` finite so an empty lane / all--inf row contributes the exact
+// accumulator identity (m, d) = (-FLT_MAX, 0). Operator epilogues interpret
+// a zero final denominator according to their public semantics. Using -inf
+// here changes +inf-row behavior.
+const FLT_MAX: f32 = 3.4028234663852886e38;
+fn is_finite_f32(value: f32) -> bool {
+  return select(false, value <= FLT_MAX, value >= -FLT_MAX);
+}
+// x - m that is exactly 0 when x equals a finite m, so exp(shifted) == 1
+// exactly at the row max. `x - x` on an infinite max is a legal fast-math
+// fold to 0, which would silently turn +inf rows finite — the explicit
+// equality test keeps the NaN propagation of the serial kernels.
+fn shifted_value(value: f32, maxValue: f32) -> f32 {
+  let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
+  return select(value - maxValue, 0.0, equalFiniteMax);
+}
+fn exp_shift(value: f32, maxValue: f32) -> f32 {
+  return exp(shifted_value(value, maxValue));
+}
+// Workgroup-cooperative merge of per-thread online-softmax (m, d) partials:
+// mNew = max(m1, m2)
+// dNew = d1 * exp(m1 - mNew) + d2 * exp(m2 - mNew)
+// Both the subgroup and portable barrier-tree engines return the same merged
+// pair to every invocation. Repeated merges require a workgroup barrier between
+// calls before their shared partial storage is reused.
+{% set combineSubgroups = combineSubgroups is defined and combineSubgroups %}
+{% if combineSubgroups %}
+// Per-subgroup partials are published into a deterministic slot: the subgroup's
+// ordinal index within the workgroup (lidx / sgSize). The online (m, d) merge
+// is not float-associative, so thread 0 must fold partials in a fixed order.
+// Subgroups partition a workgroup into contiguous ordinal ranges on supported
+// backends, so the ordinal slot is unique per subgroup and every slot in
+// [0, subgroupCount) is written (each subgroup elects one leader).
+// Sized for the worst case of one partial per invocation.
+var<workgroup> partialM: array<f32, WG>;
+var<workgroup> partialD: array<f32, WG>;
+var<workgroup> combinedMD: vec2<f32>;
+// When the whole workgroup is one subgroup the subgroup reduce already covers
+// it (no barriers, no shared state); otherwise subgroup leaders publish
+// partials through shared memory and thread 0 folds them in ordinal order.
+fn combine_partials(m: f32, d: f32, lidx: u32, sgSize: u32) -> vec2<f32> {
+  let sgM = subgroupMax(m);
+  // A lane with no elements contributes d == 0 (exact identity). A +inf
+  // element made exp(inf - inf) = NaN stick in that lane's d; a NaN element
+  // landed in d via exp(NaN); both survive the merge and are detected by the
+  // code after the reduction.
+  let sgD = subgroupAdd(d * exp_shift(m, sgM));
+  if (sgSize == WG) {
+    return vec2<f32>(sgM, sgD);
+  }
+  let subgroupCount = (WG + sgSize - 1u) / sgSize;
+  // Pre-seed every fold slot with the (max, denom) identity. The fold below reads a
+  // fixed subgroupCount slots in ordinal order (for determinism), but a slot whose
+  // subgroup elects no leader this call — e.g. a fully out-of-window key tile in the
+  // flash-attention loop that re-uses this shared memory each iteration — would
+  // otherwise read stale shared memory. Identity makes such a slot a no-op.
+  // (max identity = -FLT_MAX, denom identity = 0.)
+  if (lidx < subgroupCount) {
+    partialM[lidx] = -FLT_MAX;
+    partialD[lidx] = 0.0;
+  }
+  workgroupBarrier();
+  if (subgroupElect()) {
+    let slot = lidx / sgSize;
+    partialM[slot] = sgM;
+    partialD[slot] = sgD;
+  }
+  workgroupBarrier();
+  if (lidx == 0u) {
+    var accM = -FLT_MAX;
+    var accD = 0.0;
+    for (var i = 0u; i < subgroupCount; i = i + 1u) {
+      let mNew = max(accM, partialM[i]);
+      accD = accD * exp_shift(accM, mNew) + partialD[i] * exp_shift(partialM[i], mNew);
+      accM = mNew;
+    }
+    combinedMD = vec2<f32>(accM, accD);
+  }
+  workgroupBarrier();
+  return combinedMD;
+}
+{% else %}
+{% set mdStreamed = mdStreams is defined %}
+{% set mdStreams = mdStreams if mdStreams is defined else 1 %}
+{% set mdExtent = "WG" if mdStreams == 1 else "WG * " ~ mdStreams ~ "u" %}
+var<workgroup> partialM: array<f32, {{ mdExtent }}>;
+var<workgroup> partialD: array<f32, {{ mdExtent }}>;
+{% if mdStreamed %}
+// In-place fold of {{ mdStreams }} streams. The caller stores its per-thread
+// partials into partialM/partialD first and reads the merged pair of stream s
+// from slot s * WG afterwards.
+fn combine_partials_streams(lidx: u32) {
+  workgroupBarrier();
+  var stride = WG / 2u;
+  loop {
+    if (stride == 0u) {
+      break;
+    }
+    if (lidx < stride) {
+{% for s in range(mdStreams) %}
+      {
+        let slot = {{ s }}u * WG + lidx;
+        let m1 = partialM[slot];
+        let d1 = partialD[slot];
+        let m2 = partialM[slot + stride];
+        let d2 = partialD[slot + stride];
+        let mNew = max(m1, m2);
+        partialD[slot] = d1 * exp_shift(m1, mNew) + d2 * exp_shift(m2, mNew);
+        partialM[slot] = mNew;
+      }
+{% endfor %}
+    }
+    workgroupBarrier();
+    stride = stride / 2u;
+  }
+}
+{% else %}
+fn combine_partials(m: f32, d: f32, lidx: u32) -> vec2<f32> {
+  partialM[lidx] = m;
+  partialD[lidx] = d;
+  workgroupBarrier();
+  var stride = WG / 2u;
+  loop {
+    if (stride == 0u) {
+      break;
+    }
+    if (lidx < stride) {
+      let m1 = partialM[lidx];
+      let d1 = partialD[lidx];
+      let m2 = partialM[lidx + stride];
+      let d2 = partialD[lidx + stride];
+      let mNew = max(m1, m2);
+      partialD[lidx] = d1 * exp_shift(m1, mNew) + d2 * exp_shift(m2, mNew);
+      partialM[lidx] = mNew;
+    }
+    workgroupBarrier();
+    stride = stride / 2u;
+  }
+  let merged = vec2<f32>(partialM[0], partialD[0]);
+  // Trailing barrier so back-to-back calls cannot race a next call's partial
+  // stores against this call's reads of slot 0.
+  workgroupBarrier();
+  return merged;
+}
+{% endif %}
+{% endif %}
+@compute @workgroup_size(WG)
+fn main(@builtin(workgroup_id) wid: vec3<u32>,
+        @builtin(local_invocation_id) lid: vec3<u32>) {
+{% if source.stage == "block" %}
+  let block = wid.x;
+  let row = wid.y;
+  let rowBase = row * params.blockSize;
+  let colStart = block * BLOCK_COLS;
+  let colEnd = min(params.blockSize, colStart + BLOCK_COLS);
+  var m = -FLT_MAX;
+  var d = 0.0;
+  for (var col = colStart + lid.x; col < colEnd; col += WG) {
+    let i = rowBase + col;
+    let v = f32(data[i]) + f32(bias[bias_index(i)]);
+    let mNew = max(m, v);
+    d = d * exp_shift(m, mNew) + exp_shift(v, mNew);
+    m = mNew;
+  }
+{% else %}
+  let row = wid.x;
+  let base = row * params.blocks;
+  var m = -FLT_MAX;
+  var d = 0.0;
+  for (var block = lid.x; block < params.blocks; block += WG) {
+    let bm = blockMax[base + block];
+    let bd = blockSum[base + block];
+    let mNew = max(m, bm);
+    d = d * exp_shift(m, mNew) + bd * exp_shift(bm, mNew);
+    m = mNew;
+  }
+{% endif %}
+  let merged = combine_partials(m, d, lid.x);
+  if (lid.x == 0u) {
+{% if source.stage == "block" %}
+    let out = row * params.blocks + block;
+    blockMax[out] = merged.x;
+    blockSum[out] = merged.y;
+{% else %}
+    rowMax[row] = merged.x;
+    rowSum[row] = merged.y;
+{% endif %}
+  }
+}

build/webgpu/bias-softmax.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,281 @@

+// BiasSoftmax computes softmax(data + bias) over suffix axes [axis, rank), not
+// only the last axis. A workgroup normally owns one softmax block; short rows
+// can instead be packed with one invocation per row. Inner broadcast reuses
+// one bias row across consecutive data rows; outer broadcast cycles bias rows.
+{% if usesF16 %}
+enable f16;
+{% endif %}
+{% if combineSubgroups %}
+enable subgroups;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
+// FLT_MAX, not -inf, as the online (m, d) accumulator init: merges must keep
+// `m - m` finite so an empty lane / all--inf row contributes the exact
+// accumulator identity (m, d) = (-FLT_MAX, 0). Operator epilogues interpret
+// a zero final denominator according to their public semantics. Using -inf
+// here changes +inf-row behavior.
+const FLT_MAX: f32 = 3.4028234663852886e38;
+fn is_finite_f32(value: f32) -> bool {
+  return select(false, value <= FLT_MAX, value >= -FLT_MAX);
+}
+// x - m that is exactly 0 when x equals a finite m, so exp(shifted) == 1
+// exactly at the row max. `x - x` on an infinite max is a legal fast-math
+// fold to 0, which would silently turn +inf rows finite — the explicit
+// equality test keeps the NaN propagation of the serial kernels.
+fn shifted_value(value: f32, maxValue: f32) -> f32 {
+  let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
+  return select(value - maxValue, 0.0, equalFiniteMax);
+}
+fn exp_shift(value: f32, maxValue: f32) -> f32 {
+  return exp(shifted_value(value, maxValue));
+}
+{% if not packedRows %}
+// Workgroup-cooperative merge of per-thread online-softmax (m, d) partials:
+// mNew = max(m1, m2)
+// dNew = d1 * exp(m1 - mNew) + d2 * exp(m2 - mNew)
+// Both the subgroup and portable barrier-tree engines return the same merged
+// pair to every invocation. Repeated merges require a workgroup barrier between
+// calls before their shared partial storage is reused.
+{% set combineSubgroups = combineSubgroups is defined and combineSubgroups %}
+{% if combineSubgroups %}
+// Per-subgroup partials are published into a deterministic slot: the subgroup's
+// ordinal index within the workgroup (lidx / sgSize). The online (m, d) merge
+// is not float-associative, so thread 0 must fold partials in a fixed order.
+// Subgroups partition a workgroup into contiguous ordinal ranges on supported
+// backends, so the ordinal slot is unique per subgroup and every slot in
+// [0, subgroupCount) is written (each subgroup elects one leader).
+// Sized for the worst case of one partial per invocation.
+var<workgroup> partialM: array<f32, WG>;
+var<workgroup> partialD: array<f32, WG>;
+var<workgroup> combinedMD: vec2<f32>;
+// When the whole workgroup is one subgroup the subgroup reduce already covers
+// it (no barriers, no shared state); otherwise subgroup leaders publish
+// partials through shared memory and thread 0 folds them in ordinal order.
+fn combine_partials(m: f32, d: f32, lidx: u32, sgSize: u32) -> vec2<f32> {
+  let sgM = subgroupMax(m);
+  // A lane with no elements contributes d == 0 (exact identity). A +inf
+  // element made exp(inf - inf) = NaN stick in that lane's d; a NaN element
+  // landed in d via exp(NaN); both survive the merge and are detected by the
+  // code after the reduction.
+  let sgD = subgroupAdd(d * exp_shift(m, sgM));
+  if (sgSize == WG) {
+    return vec2<f32>(sgM, sgD);
+  }
+  let subgroupCount = (WG + sgSize - 1u) / sgSize;
+  // Pre-seed every fold slot with the (max, denom) identity. The fold below reads a
+  // fixed subgroupCount slots in ordinal order (for determinism), but a slot whose
+  // subgroup elects no leader this call — e.g. a fully out-of-window key tile in the
+  // flash-attention loop that re-uses this shared memory each iteration — would
+  // otherwise read stale shared memory. Identity makes such a slot a no-op.
+  // (max identity = -FLT_MAX, denom identity = 0.)
+  if (lidx < subgroupCount) {
+    partialM[lidx] = -FLT_MAX;
+    partialD[lidx] = 0.0;
+  }
+  workgroupBarrier();
+  if (subgroupElect()) {
+    let slot = lidx / sgSize;
+    partialM[slot] = sgM;
+    partialD[slot] = sgD;
+  }
+  workgroupBarrier();
+  if (lidx == 0u) {
+    var accM = -FLT_MAX;
+    var accD = 0.0;
+    for (var i = 0u; i < subgroupCount; i = i + 1u) {
+      let mNew = max(accM, partialM[i]);
+      accD = accD * exp_shift(accM, mNew) + partialD[i] * exp_shift(partialM[i], mNew);
+      accM = mNew;
+    }
+    combinedMD = vec2<f32>(accM, accD);
+  }
+  workgroupBarrier();
+  return combinedMD;
+}
+{% else %}
+{% set mdStreamed = mdStreams is defined %}
+{% set mdStreams = mdStreams if mdStreams is defined else 1 %}
+{% set mdExtent = "WG" if mdStreams == 1 else "WG * " ~ mdStreams ~ "u" %}
+var<workgroup> partialM: array<f32, {{ mdExtent }}>;
+var<workgroup> partialD: array<f32, {{ mdExtent }}>;
+{% if mdStreamed %}
+// In-place fold of {{ mdStreams }} streams. The caller stores its per-thread
+// partials into partialM/partialD first and reads the merged pair of stream s
+// from slot s * WG afterwards.
+fn combine_partials_streams(lidx: u32) {
+  workgroupBarrier();
+  var stride = WG / 2u;
+  loop {
+    if (stride == 0u) {
+      break;
+    }
+    if (lidx < stride) {
+{% for s in range(mdStreams) %}
+      {
+        let slot = {{ s }}u * WG + lidx;
+        let m1 = partialM[slot];
+        let d1 = partialD[slot];
+        let m2 = partialM[slot + stride];
+        let d2 = partialD[slot + stride];
+        let mNew = max(m1, m2);
+        partialD[slot] = d1 * exp_shift(m1, mNew) + d2 * exp_shift(m2, mNew);
+        partialM[slot] = mNew;
+      }
+{% endfor %}
+    }
+    workgroupBarrier();
+    stride = stride / 2u;
+  }
+}
+{% else %}
+fn combine_partials(m: f32, d: f32, lidx: u32) -> vec2<f32> {
+  partialM[lidx] = m;
+  partialD[lidx] = d;
+  workgroupBarrier();
+  var stride = WG / 2u;
+  loop {
+    if (stride == 0u) {
+      break;
+    }
+    if (lidx < stride) {
+      let m1 = partialM[lidx];
+      let d1 = partialD[lidx];
+      let m2 = partialM[lidx + stride];
+      let d2 = partialD[lidx + stride];
+      let mNew = max(m1, m2);
+      partialD[lidx] = d1 * exp_shift(m1, mNew) + d2 * exp_shift(m2, mNew);
+      partialM[lidx] = mNew;
+    }
+    workgroupBarrier();
+    stride = stride / 2u;
+  }
+  let merged = vec2<f32>(partialM[0], partialD[0]);
+  // Trailing barrier so back-to-back calls cannot race a next call's partial
+  // stores against this call's reads of slot 0.
+  workgroupBarrier();
+  return merged;
+}
+{% endif %}
+{% endif %}
+{% endif %}
+fn bias_index(index: u32) -> u32 {
+  let row = index / params.blockSize;
+  let column = index % params.blockSize;
+{% if source.isInnerBroadcast %}
+  return (row / {{ source.innerRepeat }}u) * params.blockSize + column;
+{% else %}
+  return (row % {{ source.biasBlockCount }}u) * params.blockSize + column;
+{% endif %}
+}
+@compute @workgroup_size(WG)
+fn main(@builtin(workgroup_id) wid: vec3<u32>,
+        @builtin(num_workgroups) nwg: vec3<u32>,
+        @builtin(local_invocation_id) lid: vec3<u32>{% if combineSubgroups and not packedRows %},
+        @builtin(subgroup_size) sgSize: u32{% endif %}) {
+{% if packedRows %}
+  // One invocation owns a complete short row. Grid-striding keeps the dispatch
+  // one-dimensional and bounded even for very large batch counts.
+  for (var row = wid.x * WG + lid.x; row < params.batchCount; row += nwg.x * WG) {
+    let base = row * params.blockSize;
+{% if source.blockSize == 2 %}
+    let i0 = base;
+    let i1 = base + 1u;
+    let v0 = f32(data[i0]) + f32(bias[bias_index(i0)]);
+    let v1 = f32(data[i1]) + f32(bias[bias_index(i1)]);
+    // The overwhelmingly common two-element case is a logistic pair. Choose
+    // the non-positive exponent branch for stability and use one exp total.
+    // Non-finite values retain the generic kernel's established semantics.
+    if (is_finite_f32(v0) && is_finite_f32(v1)) {
+      if (v0 >= v1) {
+        let e = exp(v1 - v0);
+        let inv = 1.0 / (1.0 + e);
+        output[i0] = {{ scalar }}(inv);
+        output[i1] = {{ scalar }}(e * inv);
+      } else {
+        let e = exp(v0 - v1);
+        let inv = 1.0 / (1.0 + e);
+        output[i0] = {{ scalar }}(e * inv);
+        output[i1] = {{ scalar }}(inv);
+      }
+    } else {
+      var m = -FLT_MAX;
+      var denom = 0.0;
+      let newM0 = max(m, v0);
+      denom = denom * exp_shift(m, newM0) + exp_shift(v0, newM0);
+      m = newM0;
+      let newM1 = max(m, v1);
+      denom = denom * exp_shift(m, newM1) + exp_shift(v1, newM1);
+      m = newM1;
+      let inv = select(0.0, 1.0 / denom, denom > 0.0);
+      output[i0] = {{ scalar }}(exp_shift(v0, m) * inv);
+      output[i1] = {{ scalar }}(exp_shift(v1, m) * inv);
+    }
+{% else %}
+    // Short rows fit in private registers, so the normalization sweep does not
+    // reread data or bias from storage.
+    var values: array<f32, {{ source.blockSize }}>;
+    var m = -FLT_MAX;
+    var denom = 0.0;
+    for (var dd = 0u; dd < params.blockSize; dd += 1u) {
+      let i = base + dd;
+      let v = f32(data[i]) + f32(bias[bias_index(i)]);
+      values[dd] = v;
+      let new_m = max(m, v);
+      denom = denom * exp_shift(m, new_m) + exp_shift(v, new_m);
+      m = new_m;
+    }
+    let inv = select(0.0, 1.0 / denom, denom > 0.0);
+    for (var dd = 0u; dd < params.blockSize; dd += 1u) {
+      let i = base + dd;
+      output[i] = {{ scalar }}(exp_shift(values[dd], m) * inv);
+    }
+{% endif %}
+  }
+{% else %}
+  // 2D-folded row index: wid.y carries the high bits past the maxComputeWorkgroupsPerDimension
+  // workgroup-per-dimension dispatch limit. Reduces to wid.x when nwg.y == 1;
+  // the row >= params.batchCount guard drops the over-dispatched tail.
+  let row = wid.x + wid.y * nwg.x;
+  if (row >= params.batchCount) {
+    return;
+  }
+  let base = row * params.blockSize;
+  let tid = lid.x;
+  // Single online (m, d) stats pass over (data + bias), then one cooperative
+  // combine. The row is read once for stats and once for normalization.
+  var m = -FLT_MAX;
+  var denom = 0.0;
+  for (var dd = tid; dd < params.blockSize; dd = dd + WG) {
+    let i = base + dd;
+    let v = f32(data[i]) + f32(bias[bias_index(i)]);
+    let new_m = max(m, v);
+    denom = denom * exp_shift(m, new_m) + exp_shift(v, new_m);
+    m = new_m;
+  }
+{% if combineSubgroups %}
+  let md = combine_partials(m, denom, tid, sgSize);
+{% else %}
+  let md = combine_partials(m, denom, tid);
+{% endif %}
+  let rowMax = md.x;
+  let inv = select(0.0, 1.0 / md.y, md.y > 0.0);
+  for (var dd = tid; dd < params.blockSize; dd = dd + WG) {
+    let i = base + dd;
+    let v = f32(data[i]) + f32(bias[bias_index(i)]);
+    output[i] = {{ scalar }}(exp_shift(v, rowMax) * inv);
+  }
+{% endif %}
+}

build/webgpu/manifest.json ADDED Viewed

	@@ -0,0 +1,279 @@

+{
+  "domain": "com.microsoft",
+  "name": "BiasSoftmax",
+  "sinceVersion": 1,
+  "description": "Computes `softmax(data + bias)` over the flattened suffix beginning at `axis`. The required `is_inner_broadcast` attribute selects how bias rows are reused: consecutive groups for inner broadcast or cyclic groups for outer broadcast. This specializes the `softmax(scores + additive_mask)` pattern used by transformer attention. Float16 and float32 are supported; the schema's double type is not.",
+  "inputs": [
+    { "role": "data", "dtype": "T", "description": "The input data tensor." },
+    {
+      "role": "bias",
+      "dtype": "T",
+      "description": "The bias (or additive mask) tensor. Its element count must be an integral number of flattened softmax rows and that row count must divide the data row count."
+    }
+  ],
+  "outputs": [
+    {
+      "role": "output",
+      "dtype": "T",
+      "rank": "ranks.data",
+      "shape": "shapes.data",
+      "description": "The output tensor; same shape as data."
+    }
+  ],
+  "attributes": { "axis": 1 },
+  "attributeConstraints": { "is_inner_broadcast": { "required": true } },
+  "attributeDescriptions": {
+    "axis": "The axis from which softmax is applied; dimensions from `axis` onward are included in the softmax reduction.",
+    "is_inner_broadcast": "When 1, bias is broadcast across dimensions from `broadcast_axis` to `axis-1`; when 0, bias is broadcast across dimensions 0 to `broadcast_axis-1`."
+  },
+  "typeConstraints": { "T": ["float32", "float16"] },
+  "args": {
+    "data": { "kind": "tensor", "semantic": "data", "role": "input" },
+    "bias": { "kind": "tensor", "semantic": "bias", "role": "input" },
+    "output": { "kind": "tensor", "semantic": "output", "role": "output" }
+  },
+  "derive": {
+    "axisNorm": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.data",
+    "batchCount": "outer(shapes.data, axisNorm)",
+    "blockSize": "dim(shapes.data, axisNorm) * inner(shapes.data, axisNorm)",
+    "biasBlockCount": "numel(shapes.bias) / max(1, blockSize)",
+    "biasContract": "(numel(shapes.data) == 0 and numel(shapes.bias) == 0) or (blockSize > 0 and biasBlockCount > 0 and numel(shapes.bias) % blockSize == 0 and biasBlockCount <= batchCount and batchCount % biasBlockCount == 0)"
+  },
+  "tunables": { "WORKGROUP_SIZE": 256, "BLOCK_COLS": 2048 },
+  "bindingSets": {
+    "biasSoftmax": [
+      {
+        "name": "data",
+        "arg": "data",
+        "semantic": "data",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$scalar"
+      },
+      {
+        "name": "bias",
+        "arg": "bias",
+        "semantic": "bias",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$scalar"
+      },
+      {
+        "name": "output",
+        "arg": "output",
+        "semantic": "output",
+        "buffer": { "type": "storage" },
+        "elementType": "$scalar"
+      },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "blockSize", "type": "u32", "value": "blockSize" },
+            { "name": "batchCount", "type": "u32", "value": "batchCount" }
+          ]
+        }
+      }
+    ]
+  },
+  "variants": [
+    {
+      "id": "longrow_split",
+      "priority": 40,
+      "when": ["numel(shapes.data) == numel(shapes.output)", "ranks.data >= 1", "attrs.axis + ranks.data >= 0", "attrs.axis < ranks.data", "biasContract", "blockSize >= 65536", "batchCount > 0", "batchCount <= device.limits.maxComputeWorkgroupsPerDimension", "ceilDiv(blockSize, tunables.BLOCK_COLS) <= device.limits.maxComputeWorkgroupsPerDimension", "f16Ok(dtypes.T)"],
+      "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"", "combineSubgroups": false },
+      "intermediates": [
+        {
+          "id": "blockMax",
+          "dtype": "float32",
+          "shape": "[outer(shapes.data, axisNorm) * ceilDiv(dim(shapes.data, axisNorm) * inner(shapes.data, axisNorm), tunables.BLOCK_COLS)]"
+        },
+        {
+          "id": "blockSum",
+          "dtype": "float32",
+          "shape": "[outer(shapes.data, axisNorm) * ceilDiv(dim(shapes.data, axisNorm) * inner(shapes.data, axisNorm), tunables.BLOCK_COLS)]"
+        },
+        { "id": "rowMax", "dtype": "float32", "shape": "[outer(shapes.data, axisNorm)]" },
+        { "id": "rowSum", "dtype": "float32", "shape": "[outer(shapes.data, axisNorm)]" }
+      ],
+      "passes": [
+        {
+          "id": "block_stats",
+          "name": "BiasSoftmax.LongRowBlockStats",
+          "source": {
+            "shader": "bias-softmax-longrow-stats.wgsl.jinja",
+            "inputs": {
+              "stage": "\"block\"",
+              "isInnerBroadcast": "attrs.is_inner_broadcast != 0",
+              "biasBlockCount": "max(1, biasBlockCount)",
+              "innerRepeat": "max(1, batchCount / max(1, biasBlockCount))"
+            }
+          },
+          "bindings": [
+            {
+              "name": "data",
+              "arg": "data",
+              "semantic": "data",
+              "buffer": { "type": "read-only-storage" },
+              "elementType": "$scalar"
+            },
+            {
+              "name": "bias",
+              "arg": "bias",
+              "semantic": "bias",
+              "buffer": { "type": "read-only-storage" },
+              "elementType": "$scalar"
+            },
+            { "name": "blockMax", "semantic": "blockMax", "buffer": { "type": "storage" }, "elementType": "f32" },
+            { "name": "blockSum", "semantic": "blockSum", "buffer": { "type": "storage" }, "elementType": "f32" },
+            {
+              "name": "params",
+              "semantic": "kernel.params",
+              "buffer": { "type": "uniform" },
+              "struct": {
+                "name": "Params",
+                "fields": [
+                  { "name": "blockSize", "type": "u32", "value": "blockSize" },
+                  { "name": "blocks", "type": "u32", "value": "ceilDiv(blockSize, tunables.BLOCK_COLS)" }
+                ]
+              }
+            }
+          ],
+          "dispatch": { "x": "ceilDiv(blockSize, tunables.BLOCK_COLS)", "y": "batchCount" }
+        },
+        {
+          "id": "row_stats",
+          "name": "BiasSoftmax.LongRowStats",
+          "source": { "shader": "bias-softmax-longrow-stats.wgsl.jinja", "inputs": { "stage": "\"row\"" } },
+          "bindings": [
+            {
+              "name": "blockMax",
+              "semantic": "blockMax",
+              "buffer": { "type": "read-only-storage" },
+              "elementType": "f32"
+            },
+            {
+              "name": "blockSum",
+              "semantic": "blockSum",
+              "buffer": { "type": "read-only-storage" },
+              "elementType": "f32"
+            },
+            { "name": "rowMax", "semantic": "rowMax", "buffer": { "type": "storage" }, "elementType": "f32" },
+            { "name": "rowSum", "semantic": "rowSum", "buffer": { "type": "storage" }, "elementType": "f32" },
+            {
+              "name": "params",
+              "semantic": "kernel.params",
+              "buffer": { "type": "uniform" },
+              "struct": {
+                "name": "Params",
+                "fields": [{ "name": "blocks", "type": "u32", "value": "ceilDiv(blockSize, tunables.BLOCK_COLS)" }]
+              }
+            }
+          ],
+          "dispatch": { "x": "batchCount" }
+        },
+        {
+          "id": "normalize",
+          "name": "BiasSoftmax.LongRowNormalize",
+          "source": {
+            "shader": "bias-softmax-longrow-normalize.wgsl.jinja",
+            "inputs": {
+              "isInnerBroadcast": "attrs.is_inner_broadcast != 0",
+              "biasBlockCount": "max(1, biasBlockCount)",
+              "innerRepeat": "max(1, batchCount / max(1, biasBlockCount))"
+            }
+          },
+          "bindings": [
+            {
+              "name": "data",
+              "arg": "data",
+              "semantic": "data",
+              "buffer": { "type": "read-only-storage" },
+              "elementType": "$scalar"
+            },
+            {
+              "name": "bias",
+              "arg": "bias",
+              "semantic": "bias",
+              "buffer": { "type": "read-only-storage" },
+              "elementType": "$scalar"
+            },
+            { "name": "rowMax", "semantic": "rowMax", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
+            { "name": "rowSum", "semantic": "rowSum", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
+            {
+              "name": "output",
+              "arg": "output",
+              "semantic": "output",
+              "buffer": { "type": "storage" },
+              "elementType": "$scalar"
+            },
+            {
+              "name": "params",
+              "semantic": "kernel.params",
+              "buffer": { "type": "uniform" },
+              "struct": { "name": "Params", "fields": [{ "name": "blockSize", "type": "u32", "value": "blockSize" }] }
+            }
+          ],
+          "dispatch": { "x": "ceilDiv(blockSize, tunables.BLOCK_COLS)", "y": "batchCount" }
+        }
+      ]
+    },
+    {
+      "id": "packed_rows",
+      "priority": 30,
+      "when": ["numel(shapes.data) == numel(shapes.output)", "ranks.data >= 1", "attrs.axis + ranks.data >= 0", "attrs.axis < ranks.data", "biasContract", "blockSize > 0", "blockSize <= 8", "batchCount >= 64", "f16Ok(dtypes.T)"],
+      "constants": {
+        "scalar": "dtypes.T",
+        "usesF16": "dtypes.T == \"f16\"",
+        "combineSubgroups": false,
+        "packedRows": true
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "BiasSoftmax.PackedRows",
+          "source": {
+            "shader": "bias-softmax.wgsl.jinja",
+            "inputs": {
+              "blockSize": "blockSize",
+              "isInnerBroadcast": "attrs.is_inner_broadcast != 0",
+              "biasBlockCount": "max(1, biasBlockCount)",
+              "innerRepeat": "max(1, batchCount / max(1, biasBlockCount))"
+            }
+          },
+          "bindings": "biasSoftmax",
+          "dispatch": { "gridStride": "batchCount", "workgroupSize": "tunables.WORKGROUP_SIZE" }
+        }
+      ]
+    },
+    {
+      "id": "adaptive_row",
+      "priority": 10,
+      "when": ["numel(shapes.data) == numel(shapes.output)", "ranks.data >= 1", "attrs.axis + ranks.data >= 0", "attrs.axis < ranks.data", "biasContract", "numel(shapes.data) >= 0", "f16Ok(dtypes.T)"],
+      "constants": {
+        "packedRows": false,
+        "scalar": "dtypes.T",
+        "usesF16": "dtypes.T == \"f16\"",
+        "combineSubgroups": "device.features.has(\"subgroups\")"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "BiasSoftmax.AdaptiveRow",
+          "source": {
+            "shader": "bias-softmax.wgsl.jinja",
+            "inputs": {
+              "blockSize": "blockSize",
+              "isInnerBroadcast": "attrs.is_inner_broadcast != 0",
+              "biasBlockCount": "max(1, biasBlockCount)",
+              "innerRepeat": "max(1, batchCount / max(1, biasBlockCount))"
+            }
+          },
+          "bindings": "biasSoftmax",
+          "dispatch": { "workgroups": "batchCount" }
+        }
+      ]
+    }
+  ]
+}

build/webgpu/metadata.json ADDED Viewed

	@@ -0,0 +1,20 @@

+{
+  "name": "com.microsoft.BiasSoftmax",
+  "id": "_com_microsoft_biassoftmax_webgpu_d835fbf",
+  "version": 1,
+  "license": "Apache-2.0",
+  "backend": { "type": "webgpu" },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "bench.json": "byPgkRP7DSHGPJmC+gIf4Ataa/cGEumT0wZKQTNmAbc=",
+      "bias-softmax-longrow-normalize.wgsl.jinja": "as8qI0C8JI2Hv60Gubt6JYGw7b2viGi8HjugoaR6sIM=",
+      "bias-softmax-longrow-stats.wgsl.jinja": "etXAjT/0FmGx03+aKetnioDZrZuZPxryTqxTDp2d5nY=",
+      "bias-softmax.wgsl.jinja": "JCf3tYMWpCyp8PRZMRx8q52g8Kg5CdcX0qzWX6F1w5M=",
+      "manifest.json": "Hmp1oerHvNRx2b3UbkPVfUJUK+xDiuZOFGY5e8ej5h0=",
+      "test.json": "tPesDd/IQUsZCm/THfvKdxo+IMX0k7DLG/Ggay8AFnk="
+    }
+  },
+  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
+  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.BiasSoftmax" }
+}

build/webgpu/test.json ADDED Viewed

	@@ -0,0 +1,442 @@

+{
+  "op": "com.microsoft.BiasSoftmax",
+  "cases": [
+    {
+      "name": "inner_broadcast_consecutive_row_groups",
+      "provenance": {
+        "source": "onnxruntime/contrib_ops/cuda/math/bias_softmax_impl.cu",
+        "notes": "Distinguishes ORT's flattened inner-broadcast row grouping from both outer mode and NumPy right-aligned broadcasting."
+      },
+      "attrs": { "axis": -1, "is_inner_broadcast": 1 },
+      "inputs": {
+        "data": { "dtype": "float32", "shape": [2, 3, 2], "data": { "kind": "constant", "value": 0.0 } },
+        "bias": { "dtype": "float32", "shape": [2, 2], "data": { "kind": "values", "values": [0.0, 2.0, 2.0, 0.0] } }
+      },
+      "outputs": {
+        "output": {
+          "dtype": "float32",
+          "shape": [2, 3, 2],
+          "tolerance": 0.000001,
+          "relTolerance": 0.000001,
+          "data": {
+            "kind": "values",
+            "values": [0.119202922022, 0.880797077978, 0.119202922022, 0.880797077978, 0.119202922022, 0.880797077978, 0.880797077978, 0.119202922022, 0.880797077978, 0.119202922022, 0.880797077978, 0.119202922022]
+          }
+        }
+      }
+    },
+    {
+      "name": "outer_broadcast_cycles_bias_rows",
+      "provenance": {
+        "source": "onnxruntime/contrib_ops/cuda/math/bias_softmax_impl.cu",
+        "notes": "Uses the same non-NumPy-broadcastable shapes as the inner case and pins ORT's cyclic outer-broadcast row selection."
+      },
+      "attrs": { "axis": -1, "is_inner_broadcast": 0 },
+      "inputs": {
+        "data": { "dtype": "float32", "shape": [2, 3, 2], "data": { "kind": "constant", "value": 0.0 } },
+        "bias": { "dtype": "float32", "shape": [2, 2], "data": { "kind": "values", "values": [0.0, 2.0, 2.0, 0.0] } }
+      },
+      "outputs": {
+        "output": {
+          "dtype": "float32",
+          "shape": [2, 3, 2],
+          "tolerance": 0.000001,
+          "relTolerance": 0.000001,
+          "data": {
+            "kind": "values",
+            "values": [0.119202922022, 0.880797077978, 0.880797077978, 0.119202922022, 0.119202922022, 0.880797077978, 0.880797077978, 0.119202922022, 0.119202922022, 0.880797077978, 0.880797077978, 0.119202922022]
+          }
+        }
+      }
+    },
+    {
+      "name": "ort_inner_broadcast_full_suffix_softmax",
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/bias_softmax_op_test.cc",
+        "test": "BiasSoftmaxTest.InnerBroadcastFullBiasBatch",
+        "notes": "Small deterministic projection of ORT's suffix softmax semantics."
+      },
+      "attrs": { "axis": 1, "is_inner_broadcast": 1 },
+      "inputs": {
+        "data": { "dtype": "float32", "shape": [1, 2, 2], "data": { "kind": "values", "values": [0.0, 0.0, 0.0, 0.0] } },
+        "bias": { "dtype": "float32", "shape": [1, 2, 2], "data": { "kind": "values", "values": [0.0, 1.0, 2.0, 3.0] } }
+      },
+      "outputs": {
+        "output": {
+          "dtype": "float32",
+          "shape": [1, 2, 2],
+          "tolerance": 0.000001,
+          "relTolerance": 0.000001,
+          "data": { "kind": "values", "values": [0.03205860328, 0.087144318742, 0.23688281809, 0.643914259888] }
+        }
+      }
+    },
+    {
+      "name": "ort_inner_broadcast_empty_bias_batch",
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/bias_softmax_op_test.cc",
+        "test": "BiasSoftmaxTest.InnerBroadcastEmptyBiasBatch",
+        "notes": "Compact deterministic projection where one bias batch broadcasts across two input batches."
+      },
+      "attrs": { "axis": 1, "is_inner_broadcast": 1 },
+      "inputs": {
+        "data": {
+          "dtype": "float32",
+          "shape": [2, 2, 2],
+          "data": { "kind": "values", "values": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] }
+        },
+        "bias": { "dtype": "float32", "shape": [1, 2, 2], "data": { "kind": "values", "values": [0.0, 1.0, 2.0, 3.0] } }
+      },
+      "outputs": {
+        "output": {
+          "dtype": "float32",
+          "shape": [2, 2, 2],
+          "tolerance": 0.000001,
+          "relTolerance": 0.000001,
+          "data": {
+            "kind": "values",
+            "values": [0.03205860328, 0.087144318742, 0.23688281809, 0.643914259888, 0.03205860328, 0.087144318742, 0.23688281809, 0.643914259888]
+          }
+        }
+      }
+    },
+    {
+      "name": "ort_inner_broadcast_negative_axis",
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/bias_softmax_op_test.cc",
+        "test": "BiasSoftmaxTest.InnerBroadcastNegativeAxis",
+        "notes": "Small deterministic projection of ORT's negative-axis inner broadcast pattern."
+      },
+      "attrs": { "axis": -2, "is_inner_broadcast": 1 },
+      "inputs": {
+        "data": {
+          "dtype": "float32",
+          "shape": [2, 2, 2],
+          "data": { "kind": "values", "values": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] }
+        },
+        "bias": { "dtype": "float32", "shape": [1, 2, 2], "data": { "kind": "values", "values": [0.0, 1.0, 2.0, 3.0] } }
+      },
+      "outputs": {
+        "output": {
+          "dtype": "float32",
+          "shape": [2, 2, 2],
+          "tolerance": 0.000001,
+          "relTolerance": 0.000001,
+          "data": {
+            "kind": "values",
+            "values": [0.03205860328, 0.087144318742, 0.23688281809, 0.643914259888, 0.03205860328, 0.087144318742, 0.23688281809, 0.643914259888]
+          }
+        }
+      }
+    },
+    {
+      "name": "ort_outer_broadcast_negative_axis",
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/bias_softmax_op_test.cc",
+        "test": "BiasSoftmaxTest.OuterBroadcastNegativeAxis",
+        "notes": "Small deterministic projection of ORT's negative-axis outer broadcast pattern."
+      },
+      "attrs": { "axis": -1, "is_inner_broadcast": 0 },
+      "inputs": {
+        "data": {
+          "dtype": "float32",
+          "shape": [2, 2, 3],
+          "data": { "kind": "values", "values": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] }
+        },
+        "bias": {
+          "dtype": "float32",
+          "shape": [2, 3],
+          "data": { "kind": "values", "values": [0.0, 1.0, 2.0, 2.0, 1.0, 0.0] }
+        }
+      },
+      "outputs": {
+        "output": {
+          "dtype": "float32",
+          "shape": [2, 2, 3],
+          "tolerance": 0.000001,
+          "relTolerance": 0.000001,
+          "data": {
+            "kind": "values",
+            "values": [0.09003057317, 0.244728471055, 0.665240955775, 0.665240955775, 0.244728471055, 0.09003057317, 0.09003057317, 0.244728471055, 0.665240955775, 0.665240955775, 0.244728471055, 0.09003057317]
+          }
+        }
+      }
+    },
+    {
+      "name": "ort_outer_broadcast_full_bias_batch",
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/bias_softmax_op_test.cc",
+        "test": "BiasSoftmaxTest.OuterBroadcastFullBiasBatch",
+        "notes": "Compact deterministic projection where every input batch has its own full bias row."
+      },
+      "attrs": { "axis": -1, "is_inner_broadcast": 0 },
+      "inputs": {
+        "data": {
+          "dtype": "float32",
+          "shape": [2, 2, 3],
+          "data": { "kind": "values", "values": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] }
+        },
+        "bias": {
+          "dtype": "float32",
+          "shape": [2, 2, 3],
+          "data": { "kind": "values", "values": [0.0, 1.0, 2.0, 2.0, 1.0, 0.0, -1.0, 0.0, 1.0, 1.0, 0.0, -1.0] }
+        }
+      },
+      "outputs": {
+        "output": {
+          "dtype": "float32",
+          "shape": [2, 2, 3],
+          "tolerance": 0.000001,
+          "relTolerance": 0.000001,
+          "data": {
+            "kind": "values",
+            "values": [0.09003057317, 0.244728471055, 0.665240955775, 0.665240955775, 0.244728471055, 0.09003057317, 0.09003057317, 0.244728471055, 0.665240955775, 0.665240955775, 0.244728471055, 0.09003057317]
+          }
+        }
+      }
+    },
+    {
+      "name": "f32_large_gap_subnormal_tail_gpu_gap",
+      "skipGpu": {
+        "category": "permanent",
+        "reason": "Portable WGSL floating-point semantics do not guarantee preservation of the subnormal values required by this fixture. Backend evidence: WebGPU/Metal flushes subnormals to zero in f32; the 87.5-point logit gap makes the tail probability subnormal, which collapses to zero instead of ORT's positive subnormal value."
+      },
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/bias_softmax_op_test.cc",
+        "test": "BiasSoftmaxTest.InnerBroadcastFullBiasBatch",
+        "notes": "BiasSoftmax is a suffix softmax over data+bias; an 87.5-point gap leaves a valid positive subnormal probability tail in ORT CPU."
+      },
+      "attrs": { "axis": 1, "is_inner_broadcast": 1 },
+      "inputs": {
+        "data": { "dtype": "float32", "shape": [1, 2], "data": { "kind": "values", "values": [0.0, -87.5] } },
+        "bias": { "dtype": "float32", "shape": [1, 2], "data": { "kind": "values", "values": [0.0, 0.0] } }
+      },
+      "outputs": {
+        "output": {
+          "dtype": "float32",
+          "shape": [1, 2],
+          "tolerance": 2e-45,
+          "relTolerance": 0,
+          "data": { "kind": "values", "values": [1.0, 9.982351397596697e-39] }
+        }
+      }
+    },
+    {
+      "name": "ort_outer_broadcast_empty_bias_batch",
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/bias_softmax_op_test.cc",
+        "test": "BiasSoftmaxTest.OuterBroadcastEmptyBiasBatch",
+        "notes": "Compact deterministic projection where one outer-broadcast bias row is reused for every input batch."
+      },
+      "attrs": { "axis": -1, "is_inner_broadcast": 0 },
+      "inputs": {
+        "data": {
+          "dtype": "float32",
+          "shape": [2, 2, 3],
+          "data": { "kind": "values", "values": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] }
+        },
+        "bias": { "dtype": "float32", "shape": [1, 3], "data": { "kind": "values", "values": [0.0, 1.0, 2.0] } }
+      },
+      "outputs": {
+        "output": {
+          "dtype": "float32",
+          "shape": [2, 2, 3],
+          "tolerance": 0.000001,
+          "relTolerance": 0.000001,
+          "data": {
+            "kind": "values",
+            "values": [0.09003057317, 0.244728471055, 0.665240955775, 0.09003057317, 0.244728471055, 0.665240955775, 0.09003057317, 0.244728471055, 0.665240955775, 0.09003057317, 0.244728471055, 0.665240955775]
+          }
+        }
+      }
+    },
+    {
+      "name": "empty_zero_dim",
+      "attrs": { "axis": 1, "is_inner_broadcast": 1 },
+      "inputs": {
+        "data": { "dtype": "float32", "shape": [0, 2, 2], "data": { "kind": "values", "values": [] } },
+        "bias": { "dtype": "float32", "shape": [0, 2, 2], "data": { "kind": "values", "values": [] } }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [0, 2, 2], "tolerance": 0 } }
+    },
+    {
+      "name": "f16_default_axis_two_element_row",
+      "attrs": { "is_inner_broadcast": 1 },
+      "inputs": {
+        "data": { "dtype": "float16", "shape": [1, 2], "data": { "kind": "values", "values": [0.0, 0.0] } },
+        "bias": { "dtype": "float16", "shape": [1, 2], "data": { "kind": "values", "values": [0.0, 0.0] } }
+      },
+      "outputs": {
+        "output": {
+          "dtype": "float16",
+          "shape": [1, 2],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [0.5, 0.5] }
+        }
+      }
+    },
+    {
+      "name": "dispatch_cliff_rows_65537",
+      "provenance": {
+        "notes": "65,537 short rows cross the per-dimension workgroup limit. The selected packed route keeps the dispatch small and grid-strides; an all-variant run also retains coverage of the generic kernel's 2D-folded row reconstruction."
+      },
+      "attrs": { "axis": 1, "is_inner_broadcast": 0 },
+      "inputs": {
+        "data": {
+          "dtype": "float32",
+          "shape": [65537, 4],
+          "data": { "kind": "cycle", "values": [0.1, 0.2, 0.3, 0.4, 0.5] }
+        },
+        "bias": { "dtype": "float32", "shape": [65537, 4], "data": { "kind": "constant", "value": 0.0 } }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [65537, 4], "tolerance": 0.0001 } }
+    },
+    {
+      "name": "fully_masked_row_neg_inf_bias",
+      "provenance": {
+        "notes": "Fully additive-masked softmax row (bias = -Infinity over the whole suffix block). GPU (FLT_MAX-init online softmax), the TS reference, ai.onnx.Softmax, and ORT's WebGPU EP all emit 0 for a fully-masked row (a NaN would poison attention). Pinned to 0."
+      },
+      "attrs": { "axis": 1, "is_inner_broadcast": 0 },
+      "inputs": {
+        "data": { "dtype": "float32", "shape": [1, 4], "data": { "kind": "values", "values": [0.0, 0.0, 0.0, 0.0] } },
+        "bias": {
+          "dtype": "float32",
+          "shape": [1, 4],
+          "data": { "kind": "values", "values": ["-Infinity", "-Infinity", "-Infinity", "-Infinity"] }
+        }
+      },
+      "outputs": {
+        "output": {
+          "dtype": "float32",
+          "shape": [1, 4],
+          "tolerance": 0.000001,
+          "data": { "kind": "values", "values": [0.0, 0.0, 0.0, 0.0] }
+        }
+      }
+    },
+    {
+      "name": "axis0_full_reduce_single_block",
+      "provenance": {
+        "notes": "axis=0 => outer(data,0)=1 => one softmax block over the whole tensor (batchCount=1, blockSize=numel). ORT requires one complete flattened bias row, so bias has the full [4,6] shape."
+      },
+      "attrs": { "axis": 0, "is_inner_broadcast": 0 },
+      "inputs": {
+        "data": { "dtype": "float32", "shape": [4, 6], "data": { "kind": "linspace", "start": -3.0, "end": 3.0 } },
+        "bias": {
+          "dtype": "float32",
+          "shape": [4, 6],
+          "data": {
+            "kind": "values",
+            "values": [0.5, -0.5, 1.0, -1.0, 0.25, -0.25, 0.5, -0.5, 1.0, -1.0, 0.25, -0.25, 0.5, -0.5, 1.0, -1.0, 0.25, -0.25, 0.5, -0.5, 1.0, -1.0, 0.25, -0.25]
+          }
+        }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [4, 6], "tolerance": 0.000001, "relTolerance": 0.000001 } }
+    },
+    {
+      "name": "bias_fold_skip_middle_dim_axis1",
+      "provenance": {
+        "notes": "bias [1,3,4] broadcast over batch axis of data [2,3,4], axis=1 (suffix softmax over 3*4=12). Exercises the compiled bias_index() fold when a size-1 bias axis is skipped and the reduction spans multiple axes. Truth from the TS reference."
+      },
+      "attrs": { "axis": 1, "is_inner_broadcast": 0 },
+      "inputs": {
+        "data": { "dtype": "float32", "shape": [2, 3, 4], "data": { "kind": "linspace", "start": -2.0, "end": 2.0 } },
+        "bias": { "dtype": "float32", "shape": [1, 3, 4], "data": { "kind": "linspace", "start": -1.0, "end": 1.0 } }
+      },
+      "outputs": {
+        "output": { "dtype": "float32", "shape": [2, 3, 4], "tolerance": 0.000001, "relTolerance": 0.000001 }
+      }
+    },
+    {
+      "name": "many_two_element_rows_axis1",
+      "provenance": {
+        "notes": "Compact companion for bench case biassoftmax-f32-launchbound-262144x2-axis1: many tiny softmax rows stress the scalar-row fallback and folded dispatch accounting."
+      },
+      "attrs": { "axis": 1, "is_inner_broadcast": 0 },
+      "inputs": {
+        "data": {
+          "dtype": "float32",
+          "shape": [4096, 2],
+          "data": { "kind": "fillFloat32", "scale": 0.25, "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "bias": { "dtype": "float32", "shape": [1, 2], "data": { "kind": "values", "values": [0.125, -0.25] } }
+      },
+      "outputs": {
+        "output": { "dtype": "float32", "shape": [4096, 2], "tolerance": 0.000001, "relTolerance": 0.000001 }
+      }
+    },
+    {
+      "name": "longrow_split_inner_broadcast_axis1_4x65536_groups2",
+      "provenance": {
+        "notes": "A long split row with consecutive inner-broadcast groups checks that each group reads the correct bias row across the full reduction."
+      },
+      "attrs": { "axis": 1, "is_inner_broadcast": 1 },
+      "inputs": {
+        "data": {
+          "dtype": "float32",
+          "shape": [4, 65536],
+          "data": { "kind": "fillFloat32", "scale": 0.5, "sinStep": 0.0017, "cosStep": 0.0031 }
+        },
+        "bias": {
+          "dtype": "float32",
+          "shape": [2, 65536],
+          "data": { "kind": "fillFloat32", "scale": 0.25, "sinStep": 0.0011, "cosStep": 0.0023 }
+        }
+      },
+      "outputs": {
+        "output": { "dtype": "float32", "shape": [4, 65536], "tolerance": 0.000001, "relTolerance": 0.00001 }
+      }
+    },
+    {
+      "name": "longrow_split_f16_axis1_4x65536_groups2",
+      "provenance": {
+        "notes": "Exercises float16 storage on the split long-row path over 65536 elements. A -30-to-0 ramp concentrates probability mass so normal float16 outputs and underflowing tail zeros coexist, while the row statistics remain in float32 scratch."
+      },
+      "attrs": { "axis": 1, "is_inner_broadcast": 1 },
+      "inputs": {
+        "data": { "dtype": "float16", "shape": [4, 65536], "data": { "kind": "linspace", "start": -30.0, "end": 0.0 } },
+        "bias": { "dtype": "float16", "shape": [2, 65536], "data": { "kind": "constant", "value": 0.0 } }
+      },
+      "outputs": { "output": { "dtype": "float16", "shape": [4, 65536], "tolerance": 1e-7, "relTolerance": 0.01 } }
+    },
+    {
+      "name": "attn_rows_axis2_4x64x256",
+      "provenance": {
+        "notes": "Compact companion for bench case biassoftmax-f32-attn-32x512x512, preserving the attention-score layout with a broadcast bias over many axis-2 softmax rows."
+      },
+      "attrs": { "axis": 2, "is_inner_broadcast": 0 },
+      "inputs": {
+        "data": {
+          "dtype": "float32",
+          "shape": [4, 64, 256],
+          "data": { "kind": "fillFloat32", "scale": 0.5, "sinStep": 0.017, "cosStep": 0.031 }
+        },
+        "bias": {
+          "dtype": "float32",
+          "shape": [1, 64, 256],
+          "data": { "kind": "fillFloat32", "scale": 0.25, "sinStep": 0.011, "cosStep": 0.023 }
+        }
+      },
+      "outputs": {
+        "output": { "dtype": "float32", "shape": [4, 64, 256], "tolerance": 0.000001, "relTolerance": 0.000001 }
+      }
+    },
+    {
+      "name": "rank8_inner_broadcast_full_suffix_softmax",
+      "attrs": { "axis": 1, "is_inner_broadcast": 1 },
+      "inputs": {
+        "data": {
+          "dtype": "float32",
+          "shape": [1, 2, 1, 2, 1, 2, 2, 2],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.27 }
+        },
+        "bias": {
+          "dtype": "float32",
+          "shape": [1, 2, 1, 2, 1, 2, 2, 2],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.11 }
+        }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [1, 2, 1, 2, 1, 2, 2, 2], "tolerance": 0.000002 } }
+    }
+  ]
+}