sync 91d990483a17

Browse files

Files changed (8) hide show

README.md +15 -11
build/webgpu/bench.json +0 -1
build/webgpu/bias-softmax-longrow-normalize.wgsl.jinja +3 -6
build/webgpu/bias-softmax-longrow-stats.wgsl.jinja +32 -38
build/webgpu/bias-softmax.wgsl.jinja +38 -46
build/webgpu/manifest.json +88 -181
build/webgpu/metadata.json +16 -9
build/webgpu/test.json +38 -14

README.md CHANGED Viewed

@@ -18,16 +18,16 @@ See the [ONNX Runtime `BiasSoftmax` contrib-operator spec](https://github.com/mi
 ## 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
@@ -36,7 +36,7 @@ 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
@@ -46,7 +46,7 @@ Attributes and default values (overridable per request):
 ## Files
-- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
 - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
@@ -56,10 +56,14 @@ Attributes and default values (overridable per request):
 ## Use with `@huggingface/kernels`
-The loader derives every required output's shape and logical dtype from the manifest contract and this call.
-It then allocates the result tensors automatically.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

 ## Inputs
+| Name | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- |
+| `data` | `T` | — | — | The input data tensor. | required |
+| `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 | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- |
+| `output` | `T` | same as `data` | same as `data` | The output tensor; same shape as data. | required |
 ## Attributes
 | 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` | — | Let `data_row` and `bias_row` index flattened rows of length `product(data.shape[axis:])`. When 1, consecutive data-row groups reuse each bias row: `bias_row = floor(data_row / (data_row_count / bias_row_count))`. When 0, bias rows repeat cyclically: `bias_row = data_row % bias_row_count`. |
 ## Type constraints
 ## Files
+- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
 - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 ## Use with `@huggingface/kernels`
+```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.2
+```
+Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

build/webgpu/bench.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "com.microsoft.BiasSoftmax",
   "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256], "BLOCK_COLS": [1024, 2048, 4096, 8192] },
   "cases": [
     {

 {
   "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256], "BLOCK_COLS": [1024, 2048, 4096, 8192] },
   "cases": [
     {

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

@@ -1,8 +1,5 @@
 // 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;
@@ -10,10 +7,10 @@ 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

 // Parallel normalization for split long rows. One workgroup handles each
 // BLOCK_COLS slice and reuses the finished per-row statistics.
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
 fn bias_index(index: u32) -> u32 {
   let row = index / params.blockSize;
   let column = index % params.blockSize;
+{% if isInnerBroadcast %}
+  return (row / {{ innerRepeat }}u) * params.blockSize + column;
 {% else %}
+  return (row % {{ biasBlockCountSpec }}u) * params.blockSize + column;
 {% endif %}
 }
 // FLT_MAX, not -inf, as the online (m, d) accumulator init: merges must keep

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

@@ -1,21 +1,21 @@
 // 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 %}
@@ -49,20 +49,22 @@ fn exp_shift(value: f32, maxValue: f32) -> f32 {
 // 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
@@ -73,31 +75,24 @@ fn combine_partials(m: f32, d: f32, lidx: u32, sgSize: u32) -> vec2<f32> {
   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);
   }
@@ -112,9 +107,8 @@ 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;
@@ -176,7 +170,7 @@ fn combine_partials(m: f32, d: f32, lidx: u32) -> vec2<f32> {
 @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;
@@ -209,7 +203,7 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>,
   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;

 // 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 stage == "block" and usesF16 %}
 enable f16;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
+{% if 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 isInnerBroadcast | default(false) %}
+  return (row / {{ innerRepeat | default(1) }}u) * params.blockSize + column;
 {% else %}
+  return (row % {{ biasBlockCountSpec | default(1) }}u) * params.blockSize + column;
 {% endif %}
 }
 {% endif %}
 // calls before their shared partial storage is reused.
 {% set combineSubgroups = combineSubgroups is defined and combineSubgroups %}
 {% if combineSubgroups %}
+// Cross-subgroup merge that assumes nothing about which invocations share a
+// subgroup or how many subgroups there are: each subgroup's elected lane
+// publishes the subgroup pair in the slot at its OWN invocation index and sets
+// that index's bit in a workgroup bitmask; thread 0 then folds exactly the
+// published slots, in ascending index order (the online (m, d) merge is not
+// float-associative, so the order is fixed), and clears the mask for the next
+// call as it reads it. Workgroup memory starts zeroed, so the mask needs no
+// setup. Same three collectives as a single-subgroup reduce, two barriers.
 var<workgroup> partialM: array<f32, WG>;
 var<workgroup> partialD: array<f32, WG>;
+var<workgroup> leaderMask: array<atomic<u32>, (WG + 31u) / 32u>;
 var<workgroup> combinedMD: vec2<f32>;
 // When the whole workgroup is one subgroup the subgroup reduce already covers
+// it (no barriers, no shared state). `subgroup_size` is the size of the current
+// subgroup and uniform, so the test is exact and may guard the barriers below.
 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
   if (sgSize == WG) {
     return vec2<f32>(sgM, sgD);
   }
   if (subgroupElect()) {
+    partialM[lidx] = sgM;
+    partialD[lidx] = sgD;
+    atomicOr(&leaderMask[lidx / 32u], 1u << (lidx % 32u));
   }
   workgroupBarrier();
   if (lidx == 0u) {
     var accM = -FLT_MAX;
     var accD = 0.0;
+    for (var w = 0u; w < (WG + 31u) / 32u; w = w + 1u) {
+      var bits = atomicExchange(&leaderMask[w], 0u);
+      while (bits != 0u) {
+        let slot = w * 32u + firstTrailingBit(bits);
+        bits = bits & (bits - 1u);
+        let mNew = max(accM, partialM[slot]);
+        accD = accD * exp_shift(accM, mNew) + partialD[slot] * exp_shift(partialM[slot], mNew);
+        accM = mNew;
+      }
     }
     combinedMD = vec2<f32>(accM, accD);
   }
 var<workgroup> partialD: array<f32, {{ mdExtent }}>;
 {% if mdStreamed %}
+// In-place fold of {{ mdStreams }} streams. Input partials occupy
+// partialM/partialD; stream s returns its merged pair in slot s * WG.
 fn combine_partials_streams(lidx: u32) {
   workgroupBarrier();
   var stride = WG / 2u;
 @compute @workgroup_size(WG)
 fn main(@builtin(workgroup_id) wid: vec3<u32>,
         @builtin(local_invocation_id) lid: vec3<u32>) {
+{% if stage == "block" %}
   let block = wid.x;
   let row = wid.y;
   let rowBase = row * params.blockSize;
   let merged = combine_partials(m, d, lid.x);
   if (lid.x == 0u) {
+{% if stage == "block" %}
     let out = row * params.blocks + block;
     blockMax[out] = merged.x;
     blockSum[out] = merged.y;

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

@@ -2,9 +2,6 @@
 // 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 %}
@@ -42,20 +39,22 @@ fn exp_shift(value: f32, maxValue: f32) -> f32 {
 // 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
@@ -66,31 +65,24 @@ fn combine_partials(m: f32, d: f32, lidx: u32, sgSize: u32) -> vec2<f32> {
   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);
   }
@@ -105,9 +97,8 @@ 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;
@@ -170,16 +161,18 @@ fn combine_partials(m: f32, d: f32, lidx: u32) -> vec2<f32> {
 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 %}
@@ -187,15 +180,14 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>,
   // 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);
@@ -224,7 +216,7 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>,
 {% 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) {
@@ -243,10 +235,10 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>,
 {% 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;
   }

 // 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 combineSubgroups %}
 enable subgroups;
 {% endif %}
 // calls before their shared partial storage is reused.
 {% set combineSubgroups = combineSubgroups is defined and combineSubgroups %}
 {% if combineSubgroups %}
+// Cross-subgroup merge that assumes nothing about which invocations share a
+// subgroup or how many subgroups there are: each subgroup's elected lane
+// publishes the subgroup pair in the slot at its OWN invocation index and sets
+// that index's bit in a workgroup bitmask; thread 0 then folds exactly the
+// published slots, in ascending index order (the online (m, d) merge is not
+// float-associative, so the order is fixed), and clears the mask for the next
+// call as it reads it. Workgroup memory starts zeroed, so the mask needs no
+// setup. Same three collectives as a single-subgroup reduce, two barriers.
 var<workgroup> partialM: array<f32, WG>;
 var<workgroup> partialD: array<f32, WG>;
+var<workgroup> leaderMask: array<atomic<u32>, (WG + 31u) / 32u>;
 var<workgroup> combinedMD: vec2<f32>;
 // When the whole workgroup is one subgroup the subgroup reduce already covers
+// it (no barriers, no shared state). `subgroup_size` is the size of the current
+// subgroup and uniform, so the test is exact and may guard the barriers below.
 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
   if (sgSize == WG) {
     return vec2<f32>(sgM, sgD);
   }
   if (subgroupElect()) {
+    partialM[lidx] = sgM;
+    partialD[lidx] = sgD;
+    atomicOr(&leaderMask[lidx / 32u], 1u << (lidx % 32u));
   }
   workgroupBarrier();
   if (lidx == 0u) {
     var accM = -FLT_MAX;
     var accD = 0.0;
+    for (var w = 0u; w < (WG + 31u) / 32u; w = w + 1u) {
+      var bits = atomicExchange(&leaderMask[w], 0u);
+      while (bits != 0u) {
+        let slot = w * 32u + firstTrailingBit(bits);
+        bits = bits & (bits - 1u);
+        let mNew = max(accM, partialM[slot]);
+        accD = accD * exp_shift(accM, mNew) + partialD[slot] * exp_shift(partialM[slot], mNew);
+        accM = mNew;
+      }
     }
     combinedMD = vec2<f32>(accM, accD);
   }
 var<workgroup> partialD: array<f32, {{ mdExtent }}>;
 {% if mdStreamed %}
+// In-place fold of {{ mdStreams }} streams. Input partials occupy
+// partialM/partialD; stream s returns its merged pair in slot s * WG.
 fn combine_partials_streams(lidx: u32) {
   workgroupBarrier();
   var stride = WG / 2u;
 fn bias_index(index: u32) -> u32 {
   let row = index / params.blockSize;
   let column = index % params.blockSize;
+{% if isInnerBroadcast %}
+  return (row / {{ innerRepeat }}u) * params.blockSize + column;
 {% else %}
+  return (row % {{ biasBlockCountSpec }}u) * params.blockSize + column;
 {% endif %}
 }
 @compute @workgroup_size(WG)
 fn main(@builtin(workgroup_id) wid: vec3<u32>,
+{%- if packedRows %}
         @builtin(num_workgroups) nwg: vec3<u32>,
+{%- endif %}
         @builtin(local_invocation_id) lid: vec3<u32>{% if combineSubgroups and not packedRows %},
         @builtin(subgroup_size) sgSize: u32{% endif %}) {
 {% if packedRows %}
   // 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 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)]);
+    // Compute a two-element row as a logistic pair, choosing the non-positive
+    // exponent branch for stability. Non-finite rows use the reduction below.
     if (is_finite_f32(v0) && is_finite_f32(v1)) {
       if (v0 >= v1) {
         let e = exp(v1 - v0);
 {% else %}
     // Short rows fit in private registers, so the normalization sweep does not
     // reread data or bias from storage.
+    var values: array<f32, {{ blockSize }}>;
     var m = -FLT_MAX;
     var denom = 0.0;
     for (var dd = 0u; dd < params.blockSize; dd += 1u) {
 {% endif %}
   }
 {% else %}
+  // 2D-folded row index: wid.y carries the high bits once the row count passes the
+  // per-axis dispatch fold width, and is zero when it does not; the
+  // row >= params.batchCount guard drops the over-dispatched tail.
+  let row = wid.x + wid.y * {{ DISPATCH_FOLD_WIDTH }}u;
   if (row >= params.batchCount) {
     return;
   }

build/webgpu/manifest.json CHANGED Viewed

@@ -2,36 +2,12 @@
   "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)",
@@ -39,50 +15,18 @@
     "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",
@@ -101,43 +45,24 @@
         {
           "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" }
@@ -145,30 +70,16 @@
         {
           "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" }
@@ -176,44 +87,19 @@
         {
           "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" }
         }
@@ -222,8 +108,8 @@
     {
       "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,
@@ -233,25 +119,37 @@
         {
           "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\"",
@@ -261,17 +159,26 @@
         {
           "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" }
         }
       ]
     }

   "domain": "com.microsoft",
   "name": "BiasSoftmax",
   "sinceVersion": 1,
+  "inputs": { "data": { "dtype": "T" }, "bias": { "dtype": "T" } },
+  "outputs": { "output": { "dtype": "T", "rank": "ranks.data", "shape": "shapes.data" } },
+  "attributes": { "axis": { "default": 1 }, "is_inner_broadcast": {} },
   "attributeConstraints": { "is_inner_broadcast": { "required": true } },
   "typeConstraints": { "T": ["float32", "float16"] },
+  "tunables": { "WORKGROUP_SIZE": { "default": 256 }, "BLOCK_COLS": { "default": 2048 } },
   "derive": {
     "axisNorm": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.data",
     "batchCount": "outer(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)"
   },
+  "when": ["numel(shapes.data) == numel(shapes.output)", "ranks.data >= 1", "attrs.axis + ranks.data >= 0", "attrs.axis < ranks.data", "biasContract", "f16Ok(dtypes.T)"],
+  "bindings": {
+    "data": { "buffer": "read-only-storage", "elementType": "$scalar" },
+    "bias": { "buffer": "read-only-storage", "elementType": "$scalar" },
+    "output": { "buffer": "storage", "elementType": "$scalar" }
   },
   "variants": [
     {
       "id": "longrow_split",
       "priority": 40,
+      "when": ["blockSize >= 65536", "batchCount > 0", "batchCount <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceilDiv(blockSize, tunables.BLOCK_COLS) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"],
+      "derive": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"", "combineSubgroups": false },
       "intermediates": [
         {
           "id": "blockMax",
         {
           "id": "block_stats",
           "name": "BiasSoftmax.LongRowBlockStats",
+          "shader": "bias-softmax-longrow-stats.wgsl.jinja",
+          "derive": {
+            "stage": "\"block\"",
+            "isInnerBroadcast": "attrs.is_inner_broadcast != 0",
+            "biasBlockCountSpec": "max(1, biasBlockCount)",
+            "innerRepeat": "max(1, batchCount / max(1, biasBlockCount))"
           },
           "bindings": [
+            "data",
+            "bias",
+            { "name": "blockMax", "buffer": "storage", "elementType": "f32" },
+            { "name": "blockSum", "buffer": "storage", "elementType": "f32" },
             {
               "name": "params",
+              "struct": [
+                { "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",
+          "shader": "bias-softmax-longrow-stats.wgsl.jinja",
+          "derive": { "stage": "\"row\"" },
           "bindings": [
+            { "name": "blockMax", "buffer": "read-only-storage", "elementType": "f32" },
+            { "name": "blockSum", "buffer": "read-only-storage", "elementType": "f32" },
+            { "name": "rowMax", "buffer": "storage", "elementType": "f32" },
+            { "name": "rowSum", "buffer": "storage", "elementType": "f32" },
             {
               "name": "params",
+              "struct": [{ "name": "blocks", "type": "u32", "value": "ceilDiv(blockSize, tunables.BLOCK_COLS)" }]
             }
           ],
           "dispatch": { "x": "batchCount" }
         {
           "id": "normalize",
           "name": "BiasSoftmax.LongRowNormalize",
+          "shader": "bias-softmax-longrow-normalize.wgsl.jinja",
+          "derive": {
+            "isInnerBroadcast": "attrs.is_inner_broadcast != 0",
+            "biasBlockCountSpec": "max(1, biasBlockCount)",
+            "innerRepeat": "max(1, batchCount / max(1, biasBlockCount))"
           },
           "bindings": [
+            "data",
+            "bias",
+            { "name": "rowMax", "buffer": "read-only-storage", "elementType": "f32" },
+            { "name": "rowSum", "buffer": "read-only-storage", "elementType": "f32" },
+            "output",
+            { "name": "params", "struct": [{ "name": "blockSize", "type": "u32", "value": "blockSize" }] }
           ],
           "dispatch": { "x": "ceilDiv(blockSize, tunables.BLOCK_COLS)", "y": "batchCount" }
         }
     {
       "id": "packed_rows",
       "priority": 30,
+      "when": ["blockSize > 0", "blockSize <= 8", "batchCount >= 64"],
+      "derive": {
         "scalar": "dtypes.T",
         "usesF16": "dtypes.T == \"f16\"",
         "combineSubgroups": false,
         {
           "id": "main",
           "name": "BiasSoftmax.PackedRows",
+          "shader": "bias-softmax.wgsl.jinja",
+          "derive": {
+            "isInnerBroadcast": "attrs.is_inner_broadcast != 0",
+            "biasBlockCountSpec": "max(1, biasBlockCount)",
+            "innerRepeat": "max(1, batchCount / max(1, biasBlockCount))"
           },
+          "bindings": [
+            "data",
+            "bias",
+            "output",
+            {
+              "name": "params",
+              "struct": [
+                { "name": "blockSize", "type": "u32", "value": "blockSize" },
+                { "name": "batchCount", "type": "u32", "value": "batchCount" }
+              ]
+            }
+          ],
+          "dispatch": {
+            "x": "min(ceilDiv((batchCount), (tunables.WORKGROUP_SIZE)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
+            "y": 1,
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "adaptive_row",
       "priority": 10,
+      "when": ["numel(shapes.data) >= 0"],
+      "derive": {
         "packedRows": false,
         "scalar": "dtypes.T",
         "usesF16": "dtypes.T == \"f16\"",
         {
           "id": "main",
           "name": "BiasSoftmax.AdaptiveRow",
+          "shader": "bias-softmax.wgsl.jinja",
+          "derive": {
+            "isInnerBroadcast": "attrs.is_inner_broadcast != 0",
+            "biasBlockCountSpec": "max(1, biasBlockCount)",
+            "innerRepeat": "max(1, batchCount / max(1, biasBlockCount))"
           },
+          "bindings": [
+            "data",
+            "bias",
+            "output",
+            {
+              "name": "params",
+              "struct": [
+                { "name": "blockSize", "type": "u32", "value": "blockSize" },
+                { "name": "batchCount", "type": "u32", "value": "batchCount" }
+              ]
+            }
+          ],
+          "dispatch": { "x": "min(batchCount, 65535)", "y": "ceilDiv(batchCount, 65535)", "z": 1 },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,20 +1,27 @@
 {
   "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" }
 }

 {
   "name": "com.microsoft.BiasSoftmax",
+  "id": "_com_microsoft_biassoftmax_webgpu_cbcd64a",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "rgHjgRMHX0A/pGpGh4i7pvOQ59E3cUmsKVhIyuZBnvY=",
+      "bias-softmax-longrow-normalize.wgsl.jinja": "lX2Z/FpHVrFdQTOlt+JKYwA7eZAokjLCgluP7kssTkc=",
+      "bias-softmax-longrow-stats.wgsl.jinja": "SXOFecNckSGTdUceUizrqmBaFAeByEkUA0bei53gRw0=",
+      "bias-softmax.wgsl.jinja": "CR9eMAgk196zS8W9yMWQTLlqDcvrrzVdvHMe9NjPnO0=",
+      "manifest.json": "Byu1oYp1SQbvVzBg59JbyQIwhIJTedDlg9UrOMbePEU=",
+      "test.json": "B+sdyYMUEL3T1qIYkK3plwl3nTChHF7KfEa6LDITSdE="
     }
   },
+  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.0",
+    "variants": {
+      "longrow_split": ["bias-softmax-longrow-normalize.wgsl.jinja", "bias-softmax-longrow-stats.wgsl.jinja"],
+      "packed_rows": ["bias-softmax.wgsl.jinja"],
+      "adaptive_row": ["bias-softmax.wgsl.jinja"]
+    }
+  }
 }

build/webgpu/test.json CHANGED Viewed

@@ -1,11 +1,10 @@
 {
-  "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": {
@@ -29,7 +28,7 @@
       "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": {
@@ -54,7 +53,7 @@
       "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": {
@@ -105,7 +104,7 @@
       "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": {
@@ -134,7 +133,7 @@
       "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": {
@@ -199,12 +198,12 @@
       "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": {
@@ -278,7 +277,7 @@
     {
       "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": {
@@ -294,7 +293,7 @@
     {
       "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": {
@@ -314,10 +313,35 @@
         }
       }
     },
     {
       "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": {
@@ -336,7 +360,7 @@
     {
       "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": {
@@ -350,7 +374,7 @@
     {
       "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": {
@@ -402,7 +426,7 @@
     {
       "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": {

 {
   "cases": [
     {
       "name": "inner_broadcast_consecutive_row_groups",
       "provenance": {
         "source": "onnxruntime/contrib_ops/cuda/math/bias_softmax_impl.cu",
+        "notes": "Distinguishes flattened inner-broadcast row grouping from outer mode and right-aligned broadcasting."
       },
       "attrs": { "axis": -1, "is_inner_broadcast": 1 },
       "inputs": {
       "name": "outer_broadcast_cycles_bias_rows",
       "provenance": {
         "source": "onnxruntime/contrib_ops/cuda/math/bias_softmax_impl.cu",
+        "notes": "Non-broadcastable data and bias shapes exercise cyclic outer-broadcast row selection."
       },
       "attrs": { "axis": -1, "is_inner_broadcast": 0 },
       "inputs": {
       "provenance": {
         "source": "onnxruntime/test/contrib_ops/bias_softmax_op_test.cc",
         "test": "BiasSoftmaxTest.InnerBroadcastFullBiasBatch",
+        "notes": "Small deterministic suffix-softmax case."
       },
       "attrs": { "axis": 1, "is_inner_broadcast": 1 },
       "inputs": {
       "provenance": {
         "source": "onnxruntime/test/contrib_ops/bias_softmax_op_test.cc",
         "test": "BiasSoftmaxTest.InnerBroadcastNegativeAxis",
+        "notes": "Negative-axis inner-broadcast case."
       },
       "attrs": { "axis": -2, "is_inner_broadcast": 1 },
       "inputs": {
       "provenance": {
         "source": "onnxruntime/test/contrib_ops/bias_softmax_op_test.cc",
         "test": "BiasSoftmaxTest.OuterBroadcastNegativeAxis",
+        "notes": "Negative-axis outer-broadcast case."
       },
       "attrs": { "axis": -1, "is_inner_broadcast": 0 },
       "inputs": {
       "name": "f32_large_gap_subnormal_tail_gpu_gap",
       "skipGpu": {
         "category": "permanent",
+        "reason": "Portable WGSL floating-point semantics do not guarantee preservation of subnormal values. The 87.5-point logit gap requires a positive subnormal probability tail, which may be flushed to zero."
       },
       "provenance": {
         "source": "onnxruntime/test/contrib_ops/bias_softmax_op_test.cc",
         "test": "BiasSoftmaxTest.InnerBroadcastFullBiasBatch",
+        "notes": "An 87.5-point logit gap leaves a positive subnormal probability tail after suffix softmax over data+bias."
       },
       "attrs": { "axis": 1, "is_inner_broadcast": 1 },
       "inputs": {
     {
       "name": "dispatch_cliff_rows_65537",
       "provenance": {
+        "notes": "65,537 short rows cross the per-dimension workgroup limit. The packed route grid-strides over rows, while the generic route reconstructs row indices from a two-dimensional dispatch."
       },
       "attrs": { "axis": 1, "is_inner_broadcast": 0 },
       "inputs": {
     {
       "name": "fully_masked_row_neg_inf_bias",
       "provenance": {
+        "notes": "The additive bias is -Infinity across the entire suffix block. The expected result uses the package's all-zero convention for a fully masked row, preventing NaNs from propagating into attention."
       },
       "attrs": { "axis": 1, "is_inner_broadcast": 0 },
       "inputs": {
         }
       }
     },
+    {
+      "name": "attn_rows_axis3_row1_all_neg_inf_bias_2x8x4x64",
+      "provenance": {
+        "notes": "Data has shape [2, 8, 4, 64], and bias [1, 1, 4, 64] uses outer broadcasting. Bias query row 1 is -Infinity at every key, so 16 of the 64 softmax rows are fully masked and must be all zero; the other rows remain finite. This operator's zero-denominator rule differs from applying ONNX Softmax directly to data plus bias, which would produce NaN for the fully masked rows."
+      },
+      "attrs": { "axis": 3, "is_inner_broadcast": 0 },
+      "inputs": {
+        "data": {
+          "dtype": "float32",
+          "shape": [2, 8, 4, 64],
+          "data": { "kind": "fillFloat32", "scale": 0.5, "sinStep": 0.017, "cosStep": 0.031 }
+        },
+        "bias": {
+          "dtype": "float32",
+          "shape": [1, 1, 4, 64],
+          "data": {
+            "kind": "values",
+            "values": [0.1, 0.1906, 0.2752, 0.3492, 0.4087, 0.4504, 0.4717, 0.4715, 0.4493, 0.4061, 0.3436, 0.265, 0.1738, 0.0744, -0.0283, -0.1293, -0.2239, -0.3075, -0.376, -0.4264, -0.4563, -0.4646, -0.4512, -0.417, -0.3642, -0.2957, -0.2155, -0.1278, -0.0373, 0.051, 0.1326, 0.2033, 0.2593, 0.2978, 0.3171, 0.3164, 0.2959, 0.2571, 0.2022, 0.1347, 0.0584, -0.0222, -0.1024, -0.1777, -0.2435, -0.296, -0.3319, -0.349, -0.3458, -0.3223, -0.2793, -0.2187, -0.1434, -0.0571, 0.0358, 0.1307, 0.2226, 0.3069, 0.379, 0.4352, 0.4725, 0.4887, 0.4829, 0.4552, "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", 0.4895, 0.458, 0.4063, 0.3372, 0.2542, 0.1614, 0.0636, -0.0343, -0.1275, -0.2114, -0.282, -0.3358, -0.3705, -0.3848, -0.3781, -0.3515, -0.3066, -0.2464, -0.1744, -0.0948, -0.0122, 0.0686, 0.1431, 0.207, 0.2566, 0.2892, 0.3028, 0.2966, 0.2708, 0.2266, 0.1664, 0.0935, 0.0118, -0.0744, -0.1602, -0.2409, -0.3121, -0.3695, -0.41, -0.4309, -0.4307, -0.4092, -0.3669, -0.3056, -0.2282, -0.1382, -0.0401, 0.0616, 0.1619, 0.2557, 0.3387, 0.4067, 0.4564, 0.4856, 0.4928, 0.4779, 0.4417, 0.3864, 0.3147, 0.2305, 0.1381, 0.0423, -0.052, -0.1401, 0.271, 0.1834, 0.0902, -0.0037, -0.0937, -0.175, -0.2436, -0.2963, -0.3304, -0.3445, -0.3382, -0.3123, -0.2685, -0.2096, -0.139, -0.061, 0.0199, 0.0989, 0.1716, 0.2336, 0.2813, 0.3118, 0.3232, 0.3144, 0.2858, 0.2386, 0.1751, 0.0984, 0.0126, -0.078, -0.1685, -0.2541, -0.3303, -0.3928, -0.4382, -0.4639, -0.4683, -0.4508, -0.4119, -0.3534, -0.278, -0.1892, -0.0912, 0.0112, 0.1131, 0.2095, 0.2959, 0.3681, 0.4226, 0.4571, 0.4701, 0.4612, 0.4311, 0.3818, 0.316, 0.2374, 0.1503, 0.0592, -0.0309, -0.1153, -0.1897, -0.2502, -0.2939, -0.3186]
+          }
+        }
+      },
+      "outputs": {
+        "output": { "dtype": "float32", "shape": [2, 8, 4, 64], "tolerance": 0.000001, "relTolerance": 0.000001 }
+      }
+    },
     {
       "name": "axis0_full_reduce_single_block",
       "provenance": {
+        "notes": "With axis=0, one softmax block covers the whole tensor. The bias supplies one complete flattened row with shape [4,6]."
       },
       "attrs": { "axis": 0, "is_inner_broadcast": 0 },
       "inputs": {
     {
       "name": "bias_fold_skip_middle_dim_axis1",
       "provenance": {
+        "notes": "Bias [1,3,4] broadcasts across the batch axis of data [2,3,4]. axis=1 reduces each 12-element suffix while the size-one bias axis is skipped during broadcast indexing."
       },
       "attrs": { "axis": 1, "is_inner_broadcast": 0 },
       "inputs": {
     {
       "name": "many_two_element_rows_axis1",
       "provenance": {
+        "notes": "Many two-element softmax rows exercise the scalar-row fallback and folded dispatch accounting used by `biassoftmax-f32-launchbound-262144x2-axis1`."
       },
       "attrs": { "axis": 1, "is_inner_broadcast": 0 },
       "inputs": {
     {
       "name": "attn_rows_axis2_4x64x256",
       "provenance": {
+        "notes": "An attention-score layout with broadcast bias exercises many axis-2 softmax rows at a compact scale representative of `biassoftmax-f32-attn-32x512x512`."
       },
       "attrs": { "axis": 2, "is_inner_broadcast": 0 },
       "inputs": {