sync 6fdf6301e2bb

Browse files

Files changed (10) hide show

README.md +17 -3
build/webgpu/bench.json +473 -1
build/webgpu/bias-softmax-longrow-normalize.wgsl.jinja +7 -6
build/webgpu/bias-softmax.wgsl.jinja +15 -48
build/webgpu/manifest.json +139 -26
build/webgpu/metadata.json +15 -10
build/webgpu/{bias-softmax-longrow-stats.wgsl.jinja → softmax-longrow-stats.wgsl.jinja} +31 -111
build/webgpu/softmax-online.wgsl.jinja +149 -0
build/webgpu/softmax-subgroup-rows.wgsl.jinja +109 -0
build/webgpu/test.json +440 -0

README.md CHANGED Viewed

@@ -44,20 +44,34 @@ Attributes and default values (overridable per request):
 | --- | --- |
 | `T` | `float32`, `float16` |
 ## Files
 - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
-- [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 - [`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`
 ```sh
-npm install --save-exact @huggingface/kernels@0.0.1-preview.2
 ```
 Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.

 | --- | --- |
 | `T` | `float32`, `float16` |
+## Implementation variants
+One implementation is selected per call from the device capabilities, the request shapes and the dtypes; these notes say what each one covers.
+- `online_subgroup_rows_vec4` — Subgroup-per-row fused bias and softmax with register-resident vec4 values, f32 addition and reduction, and the existing zero-output convention for fully masked rows.
+- `online_workgroup_vec4` — Vectorized workgroup online bias-softmax with f32 addition and reduction; portable fallback below the register-resident subgroup-row path.
+- `online_workgroup_scalar_vec4` — Logical vec4 online reduction with scalar storage for unaligned row strides. Invalid tail components contribute zero denominator and are never stored.
+## Device requirements
+Some implementation variants require `subgroups`. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.
 ## Files
 - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
+- [`bench.json`](build/webgpu/bench.json) — benchmark cases
 - [`bias-softmax-longrow-normalize.wgsl.jinja`](build/webgpu/bias-softmax-longrow-normalize.wgsl.jinja)
 - [`bias-softmax.wgsl.jinja`](build/webgpu/bias-softmax.wgsl.jinja)
+- [`softmax-longrow-stats.wgsl.jinja`](build/webgpu/softmax-longrow-stats.wgsl.jinja)
+- [`softmax-online.wgsl.jinja`](build/webgpu/softmax-online.wgsl.jinja)
+- [`softmax-subgroup-rows.wgsl.jinja`](build/webgpu/softmax-subgroup-rows.wgsl.jinja)
 ## Use with `@huggingface/kernels`
 ```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.3
 ```
 Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.

build/webgpu/bench.json CHANGED Viewed

@@ -1,5 +1,9 @@
 {
-  "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256], "BLOCK_COLS": [1024, 2048, 4096, 8192] },
   "cases": [
     {
       "name": "biassoftmax-f32-attn-32x512x512",
@@ -67,6 +71,474 @@
       },
       "outputs": { "output": { "shape": [262144, 2], "dtype": "float32" } },
       "bench": { "metrics": [{ "type": "bandwidth", "value": "(262144 * 2 * 2 + 2) * dtypeBytes(args.dtype)" }] }
     }
   ]
 }

 {
+  "tunableSpace": {
+    "WORKGROUP_SIZE": [64, 128, 256],
+    "BLOCK_COLS": [1024, 2048, 4096, 8192],
+    "ONLINE_WORKGROUP_SIZE": [32, 64, 128, 256]
+  },
   "cases": [
     {
       "name": "biassoftmax-f32-attn-32x512x512",
       },
       "outputs": { "output": { "shape": [262144, 2], "dtype": "float32" } },
       "bench": { "metrics": [{ "type": "bandwidth", "value": "(262144 * 2 * 2 + 2) * dtypeBytes(args.dtype)" }] }
+    },
+    {
+      "name": "reuse_float32_r4096_c513_bias16_wg256_inner0",
+      "preset": "stress",
+      "vars": { "dtype": "float32" },
+      "attrs": { "axis": 1, "is_inner_broadcast": 0 },
+      "tunables": { "WORKGROUP_SIZE": 256 },
+      "inputs": {
+        "data": { "shape": [4096, 513], "dtype": "float32", "dist": "normal", "seed": 601, "scale": 2 },
+        "bias": { "shape": [16, 513], "dtype": "float32", "dist": "normal", "seed": 602, "scale": 1 }
+      },
+      "outputs": { "output": { "shape": [4096, 513], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "(4096 * 513 * 2 + 16 * 513) * dtypeBytes(args.dtype)" }] }
+    },
+    {
+      "name": "reuse_float32_r1024_c4095_bias16_wg256_inner0",
+      "preset": "stress",
+      "vars": { "dtype": "float32" },
+      "attrs": { "axis": 1, "is_inner_broadcast": 0 },
+      "tunables": { "WORKGROUP_SIZE": 256 },
+      "inputs": {
+        "data": { "shape": [1024, 4095], "dtype": "float32", "dist": "normal", "seed": 601, "scale": 2 },
+        "bias": { "shape": [16, 4095], "dtype": "float32", "dist": "normal", "seed": 602, "scale": 1 }
+      },
+      "outputs": { "output": { "shape": [1024, 4095], "dtype": "float32" } },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "(1024 * 4095 * 2 + 16 * 4095) * dtypeBytes(args.dtype)" }]
+      }
+    },
+    {
+      "name": "reuse_float16_r4096_c513_bias16_wg256_inner0",
+      "preset": "stress",
+      "vars": { "dtype": "float16" },
+      "attrs": { "axis": 1, "is_inner_broadcast": 0 },
+      "tunables": { "WORKGROUP_SIZE": 256 },
+      "inputs": {
+        "data": { "shape": [4096, 513], "dtype": "float16", "dist": "normal", "seed": 601, "scale": 2 },
+        "bias": { "shape": [16, 513], "dtype": "float16", "dist": "normal", "seed": 602, "scale": 1 }
+      },
+      "outputs": { "output": { "shape": [4096, 513], "dtype": "float16" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "(4096 * 513 * 2 + 16 * 513) * dtypeBytes(args.dtype)" }] }
+    },
+    {
+      "name": "reuse_float16_r1024_c4095_bias16_wg256_inner0",
+      "preset": "stress",
+      "vars": { "dtype": "float16" },
+      "attrs": { "axis": 1, "is_inner_broadcast": 0 },
+      "tunables": { "WORKGROUP_SIZE": 256 },
+      "inputs": {
+        "data": { "shape": [1024, 4095], "dtype": "float16", "dist": "normal", "seed": 601, "scale": 2 },
+        "bias": { "shape": [16, 4095], "dtype": "float16", "dist": "normal", "seed": 602, "scale": 1 }
+      },
+      "outputs": { "output": { "shape": [1024, 4095], "dtype": "float16" } },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "(1024 * 4095 * 2 + 16 * 4095) * dtypeBytes(args.dtype)" }]
+      }
+    },
+    {
+      "name": "long_float32_r4_c65537_bias2_wg64_inner1",
+      "preset": "stress",
+      "vars": { "dtype": "float32" },
+      "attrs": { "axis": 1, "is_inner_broadcast": 1 },
+      "tunables": { "WORKGROUP_SIZE": 64 },
+      "inputs": {
+        "data": { "shape": [4, 65537], "dtype": "float32", "dist": "normal", "seed": 611, "scale": 2 },
+        "bias": { "shape": [2, 65537], "dtype": "float32", "dist": "normal", "seed": 612, "scale": 1 }
+      },
+      "outputs": { "output": { "shape": [4, 65537], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "(4 * 65537 * 2 + 2 * 65537) * dtypeBytes(args.dtype)" }] }
+    },
+    {
+      "name": "long_float32_r8_c131073_bias4_wg256_inner0",
+      "preset": "stress",
+      "vars": { "dtype": "float32" },
+      "attrs": { "axis": 1, "is_inner_broadcast": 0 },
+      "tunables": { "WORKGROUP_SIZE": 256 },
+      "inputs": {
+        "data": { "shape": [8, 131073], "dtype": "float32", "dist": "normal", "seed": 611, "scale": 2 },
+        "bias": { "shape": [4, 131073], "dtype": "float32", "dist": "normal", "seed": 612, "scale": 1 }
+      },
+      "outputs": { "output": { "shape": [8, 131073], "dtype": "float32" } },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "(8 * 131073 * 2 + 4 * 131073) * dtypeBytes(args.dtype)" }]
+      }
+    },
+    {
+      "name": "long_float16_r4_c65537_bias2_wg64_inner1",
+      "preset": "stress",
+      "vars": { "dtype": "float16" },
+      "attrs": { "axis": 1, "is_inner_broadcast": 1 },
+      "tunables": { "WORKGROUP_SIZE": 64 },
+      "inputs": {
+        "data": { "shape": [4, 65537], "dtype": "float16", "dist": "normal", "seed": 611, "scale": 2 },
+        "bias": { "shape": [2, 65537], "dtype": "float16", "dist": "normal", "seed": 612, "scale": 1 }
+      },
+      "outputs": { "output": { "shape": [4, 65537], "dtype": "float16" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "(4 * 65537 * 2 + 2 * 65537) * dtypeBytes(args.dtype)" }] }
+    },
+    {
+      "name": "long_float16_r8_c131073_bias4_wg256_inner0",
+      "preset": "stress",
+      "vars": { "dtype": "float16" },
+      "attrs": { "axis": 1, "is_inner_broadcast": 0 },
+      "tunables": { "WORKGROUP_SIZE": 256 },
+      "inputs": {
+        "data": { "shape": [8, 131073], "dtype": "float16", "dist": "normal", "seed": 611, "scale": 2 },
+        "bias": { "shape": [4, 131073], "dtype": "float16", "dist": "normal", "seed": 612, "scale": 1 }
+      },
+      "outputs": { "output": { "shape": [8, 131073], "dtype": "float16" } },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "(8 * 131073 * 2 + 4 * 131073) * dtypeBytes(args.dtype)" }]
+      }
+    },
+    {
+      "name": "subgroup-neighbor-float32-r64-c64-bias2-inner0",
+      "preset": "stress",
+      "attrs": { "axis": 1, "is_inner_broadcast": 0 },
+      "inputs": {
+        "data": { "shape": [64, 64], "dtype": "float32", "dist": "normal", "seed": 761, "scale": 2 },
+        "bias": { "shape": [2, 64], "dtype": "float32", "dist": "normal", "seed": 762, "scale": 1 }
+      },
+      "outputs": { "output": { "shape": [64, 64], "dtype": "float32" } },
+      "vars": { "dtype": "float32", "rows": 64, "cols": 64, "biasRows": 2 },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "(2 * args.rows * args.cols + args.biasRows * args.cols) * dtypeBytes(args.dtype)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "subgroup-neighbor-float32-r65-c256-bias5-inner1",
+      "preset": "stress",
+      "attrs": { "axis": 1, "is_inner_broadcast": 1 },
+      "inputs": {
+        "data": { "shape": [65, 256], "dtype": "float32", "dist": "normal", "seed": 761, "scale": 2 },
+        "bias": { "shape": [5, 256], "dtype": "float32", "dist": "normal", "seed": 762, "scale": 1 }
+      },
+      "outputs": { "output": { "shape": [65, 256], "dtype": "float32" } },
+      "vars": { "dtype": "float32", "rows": 65, "cols": 256, "biasRows": 5 },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "(2 * args.rows * args.cols + args.biasRows * args.cols) * dtypeBytes(args.dtype)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "subgroup-neighbor-float32-r2048-c128-bias16-inner0",
+      "preset": "stress",
+      "attrs": { "axis": 1, "is_inner_broadcast": 0 },
+      "inputs": {
+        "data": { "shape": [2048, 128], "dtype": "float32", "dist": "normal", "seed": 761, "scale": 2 },
+        "bias": { "shape": [16, 128], "dtype": "float32", "dist": "normal", "seed": 762, "scale": 1 }
+      },
+      "outputs": { "output": { "shape": [2048, 128], "dtype": "float32" } },
+      "vars": { "dtype": "float32", "rows": 2048, "cols": 128, "biasRows": 16 },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "(2 * args.rows * args.cols + args.biasRows * args.cols) * dtypeBytes(args.dtype)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "subgroup-neighbor-float32-r512-c1024-bias8-inner1",
+      "preset": "stress",
+      "attrs": { "axis": 1, "is_inner_broadcast": 1 },
+      "inputs": {
+        "data": { "shape": [512, 1024], "dtype": "float32", "dist": "normal", "seed": 761, "scale": 2 },
+        "bias": { "shape": [8, 1024], "dtype": "float32", "dist": "normal", "seed": 762, "scale": 1 }
+      },
+      "outputs": { "output": { "shape": [512, 1024], "dtype": "float32" } },
+      "vars": { "dtype": "float32", "rows": 512, "cols": 1024, "biasRows": 8 },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "(2 * args.rows * args.cols + args.biasRows * args.cols) * dtypeBytes(args.dtype)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "subgroup-neighbor-float32-r128-c1028-bias4-inner0",
+      "preset": "stress",
+      "attrs": { "axis": 1, "is_inner_broadcast": 0 },
+      "inputs": {
+        "data": { "shape": [128, 1028], "dtype": "float32", "dist": "normal", "seed": 761, "scale": 2 },
+        "bias": { "shape": [4, 1028], "dtype": "float32", "dist": "normal", "seed": 762, "scale": 1 }
+      },
+      "outputs": { "output": { "shape": [128, 1028], "dtype": "float32" } },
+      "vars": { "dtype": "float32", "rows": 128, "cols": 1028, "biasRows": 4 },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "(2 * args.rows * args.cols + args.biasRows * args.cols) * dtypeBytes(args.dtype)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "subgroup-neighbor-float32-r63-c512-bias3-inner1",
+      "preset": "stress",
+      "attrs": { "axis": 1, "is_inner_broadcast": 1 },
+      "inputs": {
+        "data": { "shape": [63, 512], "dtype": "float32", "dist": "normal", "seed": 761, "scale": 2 },
+        "bias": { "shape": [3, 512], "dtype": "float32", "dist": "normal", "seed": 762, "scale": 1 }
+      },
+      "outputs": { "output": { "shape": [63, 512], "dtype": "float32" } },
+      "vars": { "dtype": "float32", "rows": 63, "cols": 512, "biasRows": 3 },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "(2 * args.rows * args.cols + args.biasRows * args.cols) * dtypeBytes(args.dtype)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "subgroup-neighbor-float16-r64-c64-bias2-inner0",
+      "preset": "stress",
+      "attrs": { "axis": 1, "is_inner_broadcast": 0 },
+      "inputs": {
+        "data": { "shape": [64, 64], "dtype": "float16", "dist": "normal", "seed": 761, "scale": 2 },
+        "bias": { "shape": [2, 64], "dtype": "float16", "dist": "normal", "seed": 762, "scale": 1 }
+      },
+      "outputs": { "output": { "shape": [64, 64], "dtype": "float16" } },
+      "vars": { "dtype": "float16", "rows": 64, "cols": 64, "biasRows": 2 },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "(2 * args.rows * args.cols + args.biasRows * args.cols) * dtypeBytes(args.dtype)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "subgroup-neighbor-float16-r65-c256-bias5-inner1",
+      "preset": "stress",
+      "attrs": { "axis": 1, "is_inner_broadcast": 1 },
+      "inputs": {
+        "data": { "shape": [65, 256], "dtype": "float16", "dist": "normal", "seed": 761, "scale": 2 },
+        "bias": { "shape": [5, 256], "dtype": "float16", "dist": "normal", "seed": 762, "scale": 1 }
+      },
+      "outputs": { "output": { "shape": [65, 256], "dtype": "float16" } },
+      "vars": { "dtype": "float16", "rows": 65, "cols": 256, "biasRows": 5 },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "(2 * args.rows * args.cols + args.biasRows * args.cols) * dtypeBytes(args.dtype)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "subgroup-neighbor-float16-r2048-c128-bias16-inner0",
+      "preset": "stress",
+      "attrs": { "axis": 1, "is_inner_broadcast": 0 },
+      "inputs": {
+        "data": { "shape": [2048, 128], "dtype": "float16", "dist": "normal", "seed": 761, "scale": 2 },
+        "bias": { "shape": [16, 128], "dtype": "float16", "dist": "normal", "seed": 762, "scale": 1 }
+      },
+      "outputs": { "output": { "shape": [2048, 128], "dtype": "float16" } },
+      "vars": { "dtype": "float16", "rows": 2048, "cols": 128, "biasRows": 16 },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "(2 * args.rows * args.cols + args.biasRows * args.cols) * dtypeBytes(args.dtype)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "subgroup-neighbor-float16-r512-c1024-bias8-inner1",
+      "preset": "stress",
+      "attrs": { "axis": 1, "is_inner_broadcast": 1 },
+      "inputs": {
+        "data": { "shape": [512, 1024], "dtype": "float16", "dist": "normal", "seed": 761, "scale": 2 },
+        "bias": { "shape": [8, 1024], "dtype": "float16", "dist": "normal", "seed": 762, "scale": 1 }
+      },
+      "outputs": { "output": { "shape": [512, 1024], "dtype": "float16" } },
+      "vars": { "dtype": "float16", "rows": 512, "cols": 1024, "biasRows": 8 },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "(2 * args.rows * args.cols + args.biasRows * args.cols) * dtypeBytes(args.dtype)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "subgroup-neighbor-float16-r128-c1028-bias4-inner0",
+      "preset": "stress",
+      "attrs": { "axis": 1, "is_inner_broadcast": 0 },
+      "inputs": {
+        "data": { "shape": [128, 1028], "dtype": "float16", "dist": "normal", "seed": 761, "scale": 2 },
+        "bias": { "shape": [4, 1028], "dtype": "float16", "dist": "normal", "seed": 762, "scale": 1 }
+      },
+      "outputs": { "output": { "shape": [128, 1028], "dtype": "float16" } },
+      "vars": { "dtype": "float16", "rows": 128, "cols": 1028, "biasRows": 4 },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "(2 * args.rows * args.cols + args.biasRows * args.cols) * dtypeBytes(args.dtype)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "subgroup-neighbor-float16-r63-c512-bias3-inner1",
+      "preset": "stress",
+      "attrs": { "axis": 1, "is_inner_broadcast": 1 },
+      "inputs": {
+        "data": { "shape": [63, 512], "dtype": "float16", "dist": "normal", "seed": 761, "scale": 2 },
+        "bias": { "shape": [3, 512], "dtype": "float16", "dist": "normal", "seed": 762, "scale": 1 }
+      },
+      "outputs": { "output": { "shape": [63, 512], "dtype": "float16" } },
+      "vars": { "dtype": "float16", "rows": 63, "cols": 512, "biasRows": 3 },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "(2 * args.rows * args.cols + args.biasRows * args.cols) * dtypeBytes(args.dtype)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "workgroup-neighbor-float32-r65-c33-bias5-inner1",
+      "preset": "stress",
+      "attrs": { "axis": 1, "is_inner_broadcast": 1 },
+      "inputs": {
+        "data": { "dtype": "float32", "shape": [65, 33], "dist": "normal", "seed": 701, "scale": 2 },
+        "bias": { "dtype": "float32", "shape": [5, 33], "dist": "normal", "seed": 702, "scale": 1 }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [65, 33] } },
+      "bench": {
+        "metrics": [
+          { "type": "bandwidth", "value": "(args.rows * 2 + args.biasRows) * args.cols * dtypeBytes(args.dtype)" }
+        ]
+      },
+      "vars": { "dtype": "float32", "rows": 65, "cols": 33, "biasRows": 5 }
+    },
+    {
+      "name": "workgroup-neighbor-float32-r512-c257-bias8-inner0",
+      "preset": "stress",
+      "attrs": { "axis": 1, "is_inner_broadcast": 0 },
+      "inputs": {
+        "data": { "dtype": "float32", "shape": [512, 257], "dist": "normal", "seed": 701, "scale": 2 },
+        "bias": { "dtype": "float32", "shape": [8, 257], "dist": "normal", "seed": 702, "scale": 1 }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [512, 257] } },
+      "bench": {
+        "metrics": [
+          { "type": "bandwidth", "value": "(args.rows * 2 + args.biasRows) * args.cols * dtypeBytes(args.dtype)" }
+        ]
+      },
+      "vars": { "dtype": "float32", "rows": 512, "cols": 257, "biasRows": 8 }
+    },
+    {
+      "name": "workgroup-neighbor-float32-r128-c1025-bias4-inner1",
+      "preset": "stress",
+      "attrs": { "axis": 1, "is_inner_broadcast": 1 },
+      "inputs": {
+        "data": { "dtype": "float32", "shape": [128, 1025], "dist": "normal", "seed": 701, "scale": 2 },
+        "bias": { "dtype": "float32", "shape": [4, 1025], "dist": "normal", "seed": 702, "scale": 1 }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [128, 1025] } },
+      "bench": {
+        "metrics": [
+          { "type": "bandwidth", "value": "(args.rows * 2 + args.biasRows) * args.cols * dtypeBytes(args.dtype)" }
+        ]
+      },
+      "vars": { "dtype": "float32", "rows": 128, "cols": 1025, "biasRows": 4 }
+    },
+    {
+      "name": "workgroup-neighbor-float32-r64-c4097-bias2-inner0",
+      "preset": "stress",
+      "attrs": { "axis": 1, "is_inner_broadcast": 0 },
+      "inputs": {
+        "data": { "dtype": "float32", "shape": [64, 4097], "dist": "normal", "seed": 701, "scale": 2 },
+        "bias": { "dtype": "float32", "shape": [2, 4097], "dist": "normal", "seed": 702, "scale": 1 }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [64, 4097] } },
+      "bench": {
+        "metrics": [
+          { "type": "bandwidth", "value": "(args.rows * 2 + args.biasRows) * args.cols * dtypeBytes(args.dtype)" }
+        ]
+      },
+      "vars": { "dtype": "float32", "rows": 64, "cols": 4097, "biasRows": 2 }
+    },
+    {
+      "name": "workgroup-neighbor-float16-r65-c33-bias5-inner1",
+      "preset": "stress",
+      "attrs": { "axis": 1, "is_inner_broadcast": 1 },
+      "inputs": {
+        "data": { "dtype": "float16", "shape": [65, 33], "dist": "normal", "seed": 701, "scale": 2 },
+        "bias": { "dtype": "float16", "shape": [5, 33], "dist": "normal", "seed": 702, "scale": 1 }
+      },
+      "outputs": { "output": { "dtype": "float16", "shape": [65, 33] } },
+      "bench": {
+        "metrics": [
+          { "type": "bandwidth", "value": "(args.rows * 2 + args.biasRows) * args.cols * dtypeBytes(args.dtype)" }
+        ]
+      },
+      "vars": { "dtype": "float16", "rows": 65, "cols": 33, "biasRows": 5 }
+    },
+    {
+      "name": "workgroup-neighbor-float16-r512-c257-bias8-inner0",
+      "preset": "stress",
+      "attrs": { "axis": 1, "is_inner_broadcast": 0 },
+      "inputs": {
+        "data": { "dtype": "float16", "shape": [512, 257], "dist": "normal", "seed": 701, "scale": 2 },
+        "bias": { "dtype": "float16", "shape": [8, 257], "dist": "normal", "seed": 702, "scale": 1 }
+      },
+      "outputs": { "output": { "dtype": "float16", "shape": [512, 257] } },
+      "bench": {
+        "metrics": [
+          { "type": "bandwidth", "value": "(args.rows * 2 + args.biasRows) * args.cols * dtypeBytes(args.dtype)" }
+        ]
+      },
+      "vars": { "dtype": "float16", "rows": 512, "cols": 257, "biasRows": 8 }
+    },
+    {
+      "name": "workgroup-neighbor-float16-r128-c1025-bias4-inner1",
+      "preset": "stress",
+      "attrs": { "axis": 1, "is_inner_broadcast": 1 },
+      "inputs": {
+        "data": { "dtype": "float16", "shape": [128, 1025], "dist": "normal", "seed": 701, "scale": 2 },
+        "bias": { "dtype": "float16", "shape": [4, 1025], "dist": "normal", "seed": 702, "scale": 1 }
+      },
+      "outputs": { "output": { "dtype": "float16", "shape": [128, 1025] } },
+      "bench": {
+        "metrics": [
+          { "type": "bandwidth", "value": "(args.rows * 2 + args.biasRows) * args.cols * dtypeBytes(args.dtype)" }
+        ]
+      },
+      "vars": { "dtype": "float16", "rows": 128, "cols": 1025, "biasRows": 4 }
+    },
+    {
+      "name": "workgroup-neighbor-float16-r64-c4097-bias2-inner0",
+      "preset": "stress",
+      "attrs": { "axis": 1, "is_inner_broadcast": 0 },
+      "inputs": {
+        "data": { "dtype": "float16", "shape": [64, 4097], "dist": "normal", "seed": 701, "scale": 2 },
+        "bias": { "dtype": "float16", "shape": [2, 4097], "dist": "normal", "seed": 702, "scale": 1 }
+      },
+      "outputs": { "output": { "dtype": "float16", "shape": [64, 4097] } },
+      "bench": {
+        "metrics": [
+          { "type": "bandwidth", "value": "(args.rows * 2 + args.biasRows) * args.cols * dtypeBytes(args.dtype)" }
+        ]
+      },
+      "vars": { "dtype": "float16", "rows": 64, "cols": 4097, "biasRows": 2 }
     }
   ]
 }

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

@@ -4,13 +4,12 @@
 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 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
@@ -32,6 +31,7 @@ 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));
 }
@@ -42,6 +42,7 @@ fn main(@builtin(workgroup_id) wid: 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];
@@ -50,7 +51,7 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>,
   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);
   }
 }

 const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
 const BLOCK_COLS: u32 = {{ tunables.BLOCK_COLS }}u;
+// Flattened bias-row address shared by packed, adaptive, and split rows.
+fn bias_base(row: u32) -> u32 {
 {% if isInnerBroadcast %}
+  return (row / {{ innerRepeat }}u) * params.blockSize;
 {% else %}
+  return (row % {{ biasBlockCountSpec }}u) * params.blockSize;
 {% endif %}
 }
 // FLT_MAX, not -inf, as the online (m, d) accumulator init: merges must keep
   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));
 }
   let block = wid.x;
   let row = wid.y;
   let rowBase = row * params.blockSize;
+  let biasBase = bias_base(row);
   let colStart = block * BLOCK_COLS;
   let colEnd = min(params.blockSize, colStart + BLOCK_COLS);
   let m = rowMax[row];
   for (var col = colStart + lid.x; col < colEnd; col += WG) {
     let i = rowBase + col;
+    let v = f32(data[i]) + f32(bias[biasBase + col]);
     output[i] = {{ scalar }}(exp_shift(v, m) * inv);
   }
 }

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

@@ -27,6 +27,7 @@ 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));
 }
@@ -37,7 +38,6 @@ fn exp_shift(value: f32, maxValue: f32) -> f32 {
 // 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 %}
 // Cross-subgroup merge that assumes nothing about which invocations share a
 // subgroup or how many subgroups there are: each subgroup's elected lane
@@ -90,41 +90,9 @@ fn combine_partials(m: f32, d: f32, lidx: u32, sgSize: u32) -> vec2<f32> {
   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. 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;
-  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;
@@ -154,25 +122,22 @@ fn combine_partials(m: f32, d: f32, lidx: u32) -> vec2<f32> {
   return merged;
 }
 {% endif %}
-{% endif %}
 {% endif %}
-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 %}
@@ -180,11 +145,12 @@ 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 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.
@@ -221,7 +187,7 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>,
     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);
@@ -243,6 +209,7 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>,
     return;
   }
   let base = row * params.blockSize;
   let tid = lid.x;
   // Single online (m, d) stats pass over (data + bias), then one cooperative
@@ -251,7 +218,7 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>,
   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;
@@ -266,7 +233,7 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>,
   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 %}

   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));
 }
 // 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.
 {% 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
   return combinedMD;
 }
 {% else %}
+{% set mdExtent = "WG" %}
 var<workgroup> partialM: array<f32, {{ mdExtent }}>;
 var<workgroup> partialD: array<f32, {{ mdExtent }}>;
 fn combine_partials(m: f32, d: f32, lidx: u32) -> vec2<f32> {
   partialM[lidx] = m;
   return merged;
 }
 {% endif %}
 {% endif %}
+// Flattened bias-row address shared by packed, adaptive, and split rows.
+fn bias_base(row: u32) -> u32 {
 {% if isInnerBroadcast %}
+  return (row / {{ innerRepeat }}u) * params.blockSize;
 {% else %}
+  return (row % {{ biasBlockCountSpec }}u) * params.blockSize;
 {% 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;
+    let biasBase = bias_base(row);
 {% if blockSize == 2 %}
     let i0 = base;
     let i1 = base + 1u;
+    let v0 = f32(data[i0]) + f32(bias[biasBase]);
+    let v1 = f32(data[i1]) + f32(bias[biasBase + 1u]);
     // Compute a two-element row as a logistic pair, choosing the non-positive
     // exponent branch for stability. Non-finite rows use the reduction below.
     var denom = 0.0;
     for (var dd = 0u; dd < params.blockSize; dd += 1u) {
       let i = base + dd;
+      let v = f32(data[i]) + f32(bias[biasBase + dd]);
       values[dd] = v;
       let new_m = max(m, v);
       denom = denom * exp_shift(m, new_m) + exp_shift(v, new_m);
     return;
   }
   let base = row * params.blockSize;
+  let biasBase = bias_base(row);
   let tid = lid.x;
   // Single online (m, d) stats pass over (data + bias), then one cooperative
   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[biasBase + dd]);
     let new_m = max(m, v);
     denom = denom * exp_shift(m, new_m) + exp_shift(v, new_m);
     m = new_m;
   for (var dd = tid; dd < params.blockSize; dd = dd + WG) {
     let i = base + dd;
+    let v = f32(data[i]) + f32(bias[biasBase + dd]);
     output[i] = {{ scalar }}(exp_shift(v, rowMax) * inv);
   }
 {% endif %}

build/webgpu/manifest.json CHANGED Viewed

@@ -7,26 +7,149 @@
   "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)",
     "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)"
   },
   "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",
@@ -45,9 +168,10 @@
         {
           "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))"
@@ -55,8 +179,8 @@
           "bindings": [
             "data",
             "bias",
-            { "name": "blockMax", "buffer": "storage", "elementType": "f32" },
-            { "name": "blockSum", "buffer": "storage", "elementType": "f32" },
             {
               "name": "params",
               "struct": [
@@ -70,13 +194,13 @@
         {
           "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)" }]
@@ -109,12 +233,7 @@
       "id": "packed_rows",
       "priority": 30,
       "when": ["blockSize > 0", "blockSize <= 8", "batchCount >= 64"],
-      "derive": {
-        "scalar": "dtypes.T",
-        "usesF16": "dtypes.T == \"f16\"",
-        "combineSubgroups": false,
-        "packedRows": true
-      },
       "passes": [
         {
           "id": "main",
@@ -149,12 +268,7 @@
       "id": "adaptive_row",
       "priority": 10,
       "when": ["numel(shapes.data) >= 0"],
-      "derive": {
-        "packedRows": false,
-        "scalar": "dtypes.T",
-        "usesF16": "dtypes.T == \"f16\"",
-        "combineSubgroups": "device.features.has(\"subgroups\")"
-      },
       "passes": [
         {
           "id": "main",
@@ -177,8 +291,7 @@
               ]
             }
           ],
-          "dispatch": { "x": "min(batchCount, 65535)", "y": "ceilDiv(batchCount, 65535)", "z": 1 },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     }

   "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 },
+    "ONLINE_WORKGROUP_SIZE": { "default": 64 }
+  },
   "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)",
+    "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
+    "subgroupRowsWorkgroup": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
+    "subgroupRowsLaneVecs": "ceilDiv(blockSize / 4, device.adapterInfo.subgroupMinSize) if (device.features.has(\"subgroups\") and has(device.adapterInfo, \"subgroupMinSize\") and device.adapterInfo.subgroupMinSize > 0) else 0",
+    "biasBaseWorkgroup": "min(tunables.ONLINE_WORKGROUP_SIZE, deviceWorkgroupCap)",
+    "biasMaxWorkgroup": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
+    "biasWorkgroupSize": "min(biasMaxWorkgroup, pow2ceil(ceilDiv(blockSize, 4)), max(biasBaseWorkgroup, pow2ceil(ceilDiv(biasBaseWorkgroup * biasMaxWorkgroup, max(1, batchCount)))))"
   },
   "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": { "elementType": "$scalar" },
+    "bias": { "elementType": "$scalar" },
+    "output": { "elementType": "$scalar" }
   },
   "variants": [
+    {
+      "id": "online_subgroup_rows_vec4",
+      "priority": 20,
+      "when": ["blockSize % 4 == 0", "has(device.adapterInfo, \"subgroupMinSize\")", "device.adapterInfo.subgroupMinSize >= 16", "blockSize >= 64", "subgroupRowsLaneVecs >= 1", "subgroupRowsLaneVecs <= 8", "batchCount >= 64", "has(device.adapterInfo, \"subgroupMaxSize\")", "device.adapterInfo.subgroupMaxSize <= subgroupRowsWorkgroup", "pow2ceil(subgroupRowsWorkgroup) == subgroupRowsWorkgroup", "ceilDiv(ceilDiv(batchCount, subgroupRowsWorkgroup / device.adapterInfo.subgroupMaxSize), min(device.limits.maxComputeWorkgroupsPerDimension, 65535)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"],
+      "requires": { "features": ["subgroups"] },
+      "derive": {
+        "scalar": "dtypes.T",
+        "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
+        "workgroupSize": "subgroupRowsWorkgroup",
+        "vecsPerLane": "subgroupRowsLaneVecs",
+        "isInnerBroadcast": "attrs.is_inner_broadcast != 0",
+        "biasBlockCountSpec": "max(1, biasBlockCount)",
+        "innerRepeat": "max(1, batchCount / max(1, biasBlockCount))"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "BiasSoftmax.SubgroupRowsVec4",
+          "shader": "softmax-subgroup-rows.wgsl.jinja",
+          "derive": { "op": "\"biassoftmax\"" },
+          "bindings": [
+            { "arg": "data", "name": "x", "elementType": "$vectorScalar" },
+            { "arg": "bias", "name": "bias", "elementType": "$vectorScalar" },
+            { "arg": "output", "name": "y", "elementType": "$vectorScalar" },
+            {
+              "name": "params",
+              "struct": [
+                { "name": "rows", "type": "u32", "value": "batchCount" },
+                { "name": "vecCols", "type": "u32", "value": "blockSize / 4" }
+              ]
+            }
+          ],
+          "dispatch": {
+            "x": "min(ceilDiv(batchCount, subgroupRowsWorkgroup / device.adapterInfo.subgroupMaxSize), 65535)",
+            "y": "ceilDiv(ceilDiv(batchCount, subgroupRowsWorkgroup / device.adapterInfo.subgroupMaxSize), 65535)",
+            "z": 1
+          }
+        }
+      ]
+    },
+    {
+      "id": "online_workgroup_vec4",
+      "priority": 15,
+      "when": ["blockSize % 4 == 0", "blockSize > 32", "blockSize < 65536", "pow2ceil(biasWorkgroupSize) == biasWorkgroupSize", "ceilDiv(batchCount, min(device.limits.maxComputeWorkgroupsPerDimension, 65535)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"],
+      "derive": {
+        "scalar": "dtypes.T",
+        "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
+        "workgroupSize": "biasWorkgroupSize",
+        "combineSubgroups": false,
+        "isInnerBroadcast": "attrs.is_inner_broadcast != 0",
+        "biasBlockCountSpec": "max(1, biasBlockCount)",
+        "innerRepeat": "max(1, batchCount / max(1, biasBlockCount))",
+        "scalarVec4": false
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "BiasSoftmax.WorkgroupVec4",
+          "shader": "softmax-online.wgsl.jinja",
+          "derive": { "op": "\"biassoftmax\"", "useVec4": true },
+          "bindings": [
+            { "arg": "data", "name": "x", "elementType": "$vectorScalar" },
+            { "arg": "bias", "name": "bias", "elementType": "$vectorScalar" },
+            { "arg": "output", "name": "y", "elementType": "$vectorScalar" },
+            {
+              "name": "params",
+              "struct": [
+                { "name": "rows", "type": "u32", "value": "batchCount" },
+                { "name": "vecCols", "type": "u32", "value": "blockSize / 4" }
+              ]
+            }
+          ],
+          "dispatch": { "x": "min(batchCount, 65535)", "y": "ceilDiv(batchCount, 65535)", "z": 1 }
+        }
+      ]
+    },
+    {
+      "id": "online_workgroup_scalar_vec4",
+      "priority": 15,
+      "when": ["blockSize % 4 != 0", "blockSize > 32", "blockSize < 65536", "pow2ceil(biasWorkgroupSize) == biasWorkgroupSize", "ceilDiv(batchCount, min(device.limits.maxComputeWorkgroupsPerDimension, 65535)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"],
+      "derive": {
+        "scalar": "dtypes.T",
+        "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
+        "workgroupSize": "biasWorkgroupSize",
+        "combineSubgroups": false,
+        "isInnerBroadcast": "attrs.is_inner_broadcast != 0",
+        "biasBlockCountSpec": "max(1, biasBlockCount)",
+        "innerRepeat": "max(1, batchCount / max(1, biasBlockCount))",
+        "scalarVec4": true
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "BiasSoftmax.WorkgroupScalarVec4",
+          "shader": "softmax-online.wgsl.jinja",
+          "derive": { "op": "\"biassoftmax\"", "useVec4": true },
+          "bindings": [
+            { "arg": "data", "name": "x", "elementType": "$scalar" },
+            { "arg": "bias", "name": "bias", "elementType": "$scalar" },
+            { "arg": "output", "name": "y", "elementType": "$scalar" },
+            {
+              "name": "params",
+              "struct": [
+                { "name": "rows", "type": "u32", "value": "batchCount" },
+                { "name": "vecCols", "type": "u32", "value": "ceilDiv(blockSize, 4)" },
+                { "name": "cols", "type": "u32", "value": "blockSize" }
+              ]
+            }
+          ],
+          "dispatch": { "x": "min(batchCount, 65535)", "y": "ceilDiv(batchCount, 65535)", "z": 1 }
+        }
+      ]
+    },
     {
       "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", "combineSubgroups": false },
       "intermediates": [
         {
           "id": "blockMax",
         {
           "id": "block_stats",
           "name": "BiasSoftmax.LongRowBlockStats",
+          "shader": "softmax-longrow-stats.wgsl.jinja",
           "derive": {
             "stage": "\"block\"",
+            "biasRow": true,
             "isInnerBroadcast": "attrs.is_inner_broadcast != 0",
             "biasBlockCountSpec": "max(1, biasBlockCount)",
             "innerRepeat": "max(1, batchCount / max(1, biasBlockCount))"
           "bindings": [
             "data",
             "bias",
+            { "name": "blockMax", "elementType": "f32" },
+            { "name": "blockSum", "elementType": "f32" },
             {
               "name": "params",
               "struct": [
         {
           "id": "row_stats",
           "name": "BiasSoftmax.LongRowStats",
+          "shader": "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", "elementType": "f32" },
+            { "name": "rowSum", "elementType": "f32" },
             {
               "name": "params",
               "struct": [{ "name": "blocks", "type": "u32", "value": "ceilDiv(blockSize, tunables.BLOCK_COLS)" }]
       "id": "packed_rows",
       "priority": 30,
       "when": ["blockSize > 0", "blockSize <= 8", "batchCount >= 64"],
+      "derive": { "scalar": "dtypes.T", "combineSubgroups": false, "packedRows": true },
       "passes": [
         {
           "id": "main",
       "id": "adaptive_row",
       "priority": 10,
       "when": ["numel(shapes.data) >= 0"],
+      "derive": { "packedRows": false, "scalar": "dtypes.T", "combineSubgroups": "device.features.has(\"subgroups\")" },
       "passes": [
         {
           "id": "main",
               ]
             }
           ],
+          "dispatch": { "x": "min(batchCount, 65535)", "y": "ceilDiv(batchCount, 65535)", "z": 1 }
         }
       ]
     }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,25 +1,30 @@
 {
   "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"]
     }

 {
   "name": "com.microsoft.BiasSoftmax",
+  "id": "_com_microsoft_biassoftmax_webgpu_645de3f",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "ULVbXINWou5wlT+4e2N9gtCE7nF8AjENyJAZwBBGwP4=",
+      "bias-softmax-longrow-normalize.wgsl.jinja": "C4RjAERxuPnTg7GKB286RYMPx5HWXFd9LqKrUcAdtYs=",
+      "bias-softmax.wgsl.jinja": "ptSv4JQ/WGXTqUKNqFowI6g64po1WtcXEP0sjM2LHEk=",
+      "manifest.json": "P0QfnLqsi08AxrTCYJkPTTQ0nFFh24vGS/7rwAKoOMw=",
+      "softmax-longrow-stats.wgsl.jinja": "oj7lwGMXQNTAI9NWK3/PgXwLT2R3vtTEmh6HUw4BtLE=",
+      "softmax-online.wgsl.jinja": "N/W8ChLz0EbG2y9O85tk1rmZwc6CpvlQb5lwxgfvWks=",
+      "softmax-subgroup-rows.wgsl.jinja": "S9jb9PzuCJjMBpOcVe2hWmSupF9Nia9AA5OyaIlTwSk=",
+      "test.json": "Do/oRO/PgM8zp16pyq7EVawVJu1VThzud/P5lPfX+I8="
     }
   },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
   "webgpu": {
+    "manifestSpec": "2.1",
     "variants": {
+      "online_subgroup_rows_vec4": ["softmax-subgroup-rows.wgsl.jinja"],
+      "online_workgroup_vec4": ["softmax-online.wgsl.jinja"],
+      "online_workgroup_scalar_vec4": ["softmax-online.wgsl.jinja"],
+      "longrow_split": ["bias-softmax-longrow-normalize.wgsl.jinja", "softmax-longrow-stats.wgsl.jinja"],
       "packed_rows": ["bias-softmax.wgsl.jinja"],
       "adaptive_row": ["bias-softmax.wgsl.jinja"]
     }

build/webgpu/{bias-softmax-longrow-stats.wgsl.jinja → softmax-longrow-stats.wgsl.jinja} RENAMED Viewed

@@ -1,24 +1,29 @@
-// 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 %}
 // 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
@@ -38,6 +43,7 @@ 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));
 }
@@ -47,94 +53,9 @@ fn exp_shift(value: f32, maxValue: f32) -> f32 {
 // 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 %}
-// 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
-  // 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);
-  }
-  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);
-  }
-  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. 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;
-  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;
@@ -163,36 +84,35 @@ fn combine_partials(m: f32, d: f32, lidx: u32) -> vec2<f32> {
   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 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);

+// Online-softmax (maximum, sum-exp) statistics for very long contiguous
+// last-axis rows. The statistics are the same for softmax and log-softmax.
+// The block stage assigns one workgroup to each (block, row), merges per-thread
+// online pairs, and writes one partial pair. The row stage assigns one
+// workgroup per row and folds those block partials into rowMax and rowSum.
+{% set biasRow = biasRow | default(false) %}
+{% set isInnerBroadcast = isInnerBroadcast if isInnerBroadcast is defined else false %}
+{% set innerRepeat = innerRepeat if innerRepeat is defined else 1 %}
+{% set biasBlockCountSpec = biasBlockCountSpec if biasBlockCountSpec is defined else 1 %}
+{% set cols = "params.blockSize" if biasRow else "params.cols" %}
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
 {% if stage == "block" %}
 const BLOCK_COLS: u32 = {{ tunables.BLOCK_COLS }}u;
+{% if biasRow %}
+// Flattened bias-row address shared by packed, adaptive, and split rows.
+fn bias_base(row: u32) -> u32 {
+{% if isInnerBroadcast %}
+  return (row / {{ innerRepeat }}u) * params.blockSize;
 {% else %}
+  return (row % {{ biasBlockCountSpec }}u) * params.blockSize;
 {% endif %}
 }
 {% 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
   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));
 }
 // 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 mdExtent = "WG" %}
 var<workgroup> partialM: array<f32, {{ mdExtent }}>;
 var<workgroup> partialD: array<f32, {{ mdExtent }}>;
 fn combine_partials(m: f32, d: f32, lidx: u32) -> vec2<f32> {
   partialM[lidx] = m;
   workgroupBarrier();
   return merged;
 }
 @compute @workgroup_size(WG)
+fn main(@builtin(workgroup_id) wg: vec3<u32>,
         @builtin(local_invocation_id) lid: vec3<u32>) {
 {% if stage == "block" %}
+  let block = wg.x;
+  let row = wg.y;
+  let rowBase = row * {{ cols }};
+{% if biasRow %}
+  let biasBase = bias_base(row);
+{% endif %}
   let colStart = block * BLOCK_COLS;
+  let colEnd = min({{ cols }}, colStart + BLOCK_COLS);
   var m = -FLT_MAX;
   var d = 0.0;
+  for (var col = colStart + lid.x; col < colEnd; col = col + WG) {
+    let v = {{ "f32(data[rowBase + col]) + f32(bias[biasBase + col])" if biasRow else "f32(x[rowBase + col])" }};
     let mNew = max(m, v);
     d = d * exp_shift(m, mNew) + exp_shift(v, mNew);
     m = mNew;
   }
 {% else %}
+  let row = wg.x;
   let base = row * params.blocks;
   var m = -FLT_MAX;
   var d = 0.0;
+  for (var block = lid.x; block < params.blocks; block = block + WG) {
     let bm = blockMax[base + block];
     let bd = blockSum[base + block];
     let mNew = max(m, bm);

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

	@@ -0,0 +1,149 @@

+{{ env.wgsl.resourceDeclarations }}
+// Workgroup-level online normalization over one contiguous last-axis row per
+// workgroup. Each thread keeps an online (maximum, denominator) pair over its
+// strided chunk. The workgroup merges those pairs, then a fused second sweep
+// writes the normalized result. Vectorized rows use native vec4 storage when
+// aligned, or guarded scalar storage at unaligned row strides. Both paths retain
+// f32 bias and reduction math.
+//
+// `workgroupSize` fixes the merge-tree width for this specialization.
+// `scalarVec4` selects scalar storage for fused bias rows only.
+{% macro read_vec4(c) %}
+{% if scalarVec4 %}load_word(rowBase, biasBase, {{ c }}){% else %}vec4<f32>(x[rowBase + {{ c }}]) + vec4<f32>(bias[biasBase + {{ c }}]){% endif %}
+{% endmacro %}
+{% macro store_vec4(value, c) %}
+{% if scalarVec4 %}
+{% for i in range(4) %}
+    if ({{ c }} * 4u + {{ i }}u < params.cols) {
+      y[rowBase + {{ c }} * 4u + {{ i }}u] = {{ scalar }}({{ value }}[{{ i }}]);
+    }
+{% endfor %}
+{% else %}
+    y[rowBase + {{ c }}] = vec4<{{ scalar }}>({{ value }});
+{% endif %}
+{% endmacro %}
+const WG: u32 = {{ workgroupSize }}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));
+}
+// 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 mdExtent = "WG" %}
+var<workgroup> partialM: array<f32, {{ mdExtent }}>;
+var<workgroup> partialD: array<f32, {{ mdExtent }}>;
+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;
+}
+{% if useVec4 and scalarVec4 %}
+fn load_word(rowBase: u32, biasBase: u32, c: u32) -> vec4<f32> {
+  let col = c * 4u;
+  if (col + 3u < params.cols) {
+    return vec4<f32>({% for i in range(4) %}f32(x[rowBase + col + {{ i }}u]){% if not loop.last %}, {% endif %}{% endfor %}) + vec4<f32>({% for i in range(4) %}f32(bias[biasBase + col + {{ i }}u]){% if not loop.last %}, {% endif %}{% endfor %});
+  }
+  var v = vec4<f32>(-FLT_MAX);
+{% for i in range(4) %}
+  if (col + {{ i }}u < params.cols) {
+    v[{{ i }}] = f32(x[rowBase + col + {{ i }}u]) + f32(bias[biasBase + col + {{ i }}u]);
+  }
+{% endfor %}
+  return v;
+}
+{% endif %}
+@compute @workgroup_size({{ workgroupSize }})
+fn main(@builtin(workgroup_id) wid: vec3<u32>,
+        @builtin(local_invocation_id) lid: vec3<u32>) {
+  // 2D-folded row index: wid.y carries the high bits past the
+  // per-axis dispatch fold width (rows > the per-axis dispatch fold width).
+  let row = wid.x + wid.y * {{ DISPATCH_FOLD_WIDTH }}u;
+  if (row >= params.rows) {
+    return;
+  }
+  let rowBase = row * params.{% if scalarVec4 %}cols{% else %}vecCols{% endif %};
+{% if isInnerBroadcast %}
+  let biasBase = (row / {{ innerRepeat }}u) * params.{% if scalarVec4 %}cols{% else %}vecCols{% endif %};
+{% else %}
+  let biasBase = (row % {{ biasBlockCountSpec }}u) * params.{% if scalarVec4 %}cols{% else %}vecCols{% endif %};
+{% endif %}
+  // Thread-local online (m, d) over a strided chunk; f32 accumulation.
+  var m = -FLT_MAX;
+  var d = 0.0;
+  for (var c = lid.x; c < params.vecCols; c = c + WG) {
+    let v = {{ read_vec4("c") }};
+    let mNew = max(m, max(max(v.x, v.y), max(v.z, v.w)));
+{% if scalarVec4 %}
+    // Padding is an identity even when valid values equal -FLT_MAX.
+    // Its exponential must contribute zero, not exp(0).
+    let valid = vec4<u32>(c * 4u) + vec4<u32>(0u, 1u, 2u, 3u) < vec4<u32>(params.cols);
+    let terms = vec4<f32>(exp_shift(v.x, mNew), exp_shift(v.y, mNew), exp_shift(v.z, mNew), exp_shift(v.w, mNew));
+    let masked = select(vec4<f32>(0.0), terms, valid);
+    d = d * exp_shift(m, mNew) + masked.x + masked.y + masked.z + masked.w;
+{% else %}
+    d = d * exp_shift(m, mNew)
+      + exp_shift(v.x, mNew) + exp_shift(v.y, mNew) + exp_shift(v.z, mNew) + exp_shift(v.w, mNew);
+{% endif %}
+    m = mNew;
+  }
+  let rowMD = combine_partials(m, d, lid.x);
+  let finalM = rowMD.x;
+  let finalD = rowMD.y;
+  let inv = select(0.0, 1.0 / finalD, finalD > 0.0);
+  for (var c = lid.x; c < params.vecCols; c += WG) {
+    let v = {{ read_vec4("c") }};
+    let o = vec4<f32>(exp_shift(v.x, finalM), exp_shift(v.y, finalM), exp_shift(v.z, finalM), exp_shift(v.w, finalM)) * inv;
+{{ store_vec4("o", "c") }}
+  }
+}

build/webgpu/softmax-subgroup-rows.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,109 @@

+enable subgroups;
+{{ env.wgsl.resourceDeclarations }}
+// Subgroup-per-row online normalization over a contiguous last axis of
+// moderate width. One SUBGROUP owns each row and a workgroup carries as many
+// rows as it has subgroups, so the (max, denom) merge is one pair of subgroup
+// collectives: no workgroup memory, no barriers. Each lane keeps its vec4 words
+// in registers between the statistics sweep and the normalizing store, so the
+// row is read from memory once. Each lane forms an online (maximum,
+// denominator) pair, and subgroup collectives merge those pairs.
+const WG: u32 = {{ workgroupSize }}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));
+}
+var<workgroup> sgLaneClaims: atomic<u32>;
+struct SubgroupLogicalLanes {
+  ltid: u32,
+  ord: u32,
+  rank: u32,
+  count: u32,
+}
+fn subgroup_logical_lanes(rank: u32, count: u32) -> SubgroupLogicalLanes {
+  // Every lane performs the atomic so no collective follows a lane guard: a
+  // subgroup op after a closed `if (rank == 0u)` block is a reconvergence
+  // hazard, since every subgroup operation must be reached in uniform control
+  // flow. Only the lane-0 lane adds its subgroup's claim, every other lane
+  // adds 0 and discards its snapshot. Lane 0 is the lowest lane, which is the
+  // lane `subgroupBroadcastFirst` reads.
+  let ticket = atomicAdd(&sgLaneClaims, select(0u, count | (1u << 16u), rank == 0u));
+  let claim = subgroupBroadcastFirst(ticket);
+  return SubgroupLogicalLanes((claim & 0xffffu) + rank, claim >> 16u, rank, count);
+}
+@compute @workgroup_size({{ workgroupSize }})
+fn main(@builtin(workgroup_id) wid: vec3<u32>,
+        @builtin(subgroup_invocation_id) sgLaneId: u32,
+        @builtin(subgroup_size) sgSize: u32) {
+  // Rows tile the folded workgroup grid; the lanes of one subgroup share a row.
+  // Which subgroup owns which row, and which lane holds which columns, come from
+  // the logical lane coordinates, so no assumption is made about which
+  // invocations share a subgroup.
+  let L = subgroup_logical_lanes(sgLaneId, sgSize);
+  let sgLane = L.rank;
+  let rowsPerWorkgroup = WG / sgSize;
+  let row = (wid.x + wid.y * {{ DISPATCH_FOLD_WIDTH }}u) * rowsPerWorkgroup + L.ord;
+  // A subgroup past the last row runs the collectives on an empty row and
+  // stores nothing, so every collective stays in uniform control flow.
+  let rowValid = row < params.rows;
+  let vecLimit = select(0u, params.vecCols, rowValid);
+  let rowBase = row * params.vecCols;
+{% if isInnerBroadcast %}
+  let biasBase = (row / {{ innerRepeat }}u) * params.vecCols;
+{% else %}
+  let biasBase = (row % {{ biasBlockCountSpec }}u) * params.vecCols;
+{% endif %}
+  // Lane-local online (m, d) over the lane's strided words; f32 accumulation.
+  var m = -FLT_MAX;
+  var d = 0.0;
+{% for i in range(vecsPerLane) %}
+  var v{{ i }} = vec4<f32>(0.0);
+  let c{{ i }} = sgLane + {{ i }}u * sgSize;
+  if (c{{ i }} < vecLimit) {
+    v{{ i }} = vec4<f32>(x[rowBase + c{{ i }}]) + vec4<f32>(bias[biasBase + c{{ i }}]);
+    let mNew{{ i }} = max(m, max(max(v{{ i }}.x, v{{ i }}.y), max(v{{ i }}.z, v{{ i }}.w)));
+    d = d * exp_shift(m, mNew{{ i }})
+      + exp_shift(v{{ i }}.x, mNew{{ i }}) + exp_shift(v{{ i }}.y, mNew{{ i }})
+      + exp_shift(v{{ i }}.z, mNew{{ i }}) + exp_shift(v{{ i }}.w, mNew{{ i }});
+    m = mNew{{ i }};
+  }
+{% endfor %}
+  // Subgroup merge: a lane with no words contributes d == 0 (exact identity).
+  let finalM = subgroupMax(m);
+  let finalD = subgroupAdd(d * exp_shift(m, finalM));
+  // BiasSoftmax treats an entirely additive-masked row as zero.
+  let inv = select(0.0, 1.0 / finalD, finalD > 0.0);
+{% for i in range(vecsPerLane) %}
+  if (c{{ i }} < vecLimit) {
+    let o{{ i }} = vec4<f32>(
+      exp_shift(v{{ i }}.x, finalM), exp_shift(v{{ i }}.y, finalM),
+      exp_shift(v{{ i }}.z, finalM), exp_shift(v{{ i }}.w, finalM)) * inv;
+    y[rowBase + c{{ i }}] = vec4<{{ scalar }}>(o{{ i }});
+  }
+{% endfor %}
+}

build/webgpu/test.json CHANGED Viewed

@@ -1,4 +1,8 @@
 {
   "cases": [
     {
       "name": "inner_broadcast_consecutive_row_groups",
@@ -461,6 +465,442 @@
         }
       },
       "outputs": { "output": { "dtype": "float32", "shape": [1, 2, 1, 2, 1, 2, 2, 2], "tolerance": 0.000002 } }
     }
   ]
 }

 {
+  "fixtureArrays": {
+    "subgroup_rows_float32_nonfinite_bias_input_bias": [0.25, 0.25, 0.25, "Infinity", 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, "NaN", 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, "-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.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25],
+    "ragged_float32_inner_broadcast_nonfinite_tail_input_bias": [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, "NaN", "-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", "-Infinity", "-Infinity", "-Infinity"]
+  },
   "cases": [
     {
       "name": "inner_broadcast_consecutive_row_groups",
         }
       },
       "outputs": { "output": { "dtype": "float32", "shape": [1, 2, 1, 2, 1, 2, 2, 2], "tolerance": 0.000002 } }
+    },
+    {
+      "name": "bias_row_tail_float32_6x513_inner1_wg256",
+      "provenance": {
+        "notes": "Odd suffix lengths cross invocation and block boundaries. Distinct bias rows distinguish cyclic reuse from consecutive row groups."
+      },
+      "attrs": { "axis": 1, "is_inner_broadcast": 1 },
+      "tunables": { "WORKGROUP_SIZE": 256 },
+      "inputs": {
+        "data": {
+          "dtype": "float32",
+          "shape": [6, 513],
+          "data": { "kind": "fillFloat32", "scale": 0.5, "sinStep": 0.017, "cosStep": 0.031 }
+        },
+        "bias": {
+          "dtype": "float32",
+          "shape": [2, 513],
+          "data": { "kind": "fillFloat32", "scale": 0.25, "sinStep": 0.037, "cosStep": 0.053 }
+        }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [6, 513], "tolerance": 0.000001 } }
+    },
+    {
+      "name": "bias_row_tail_float32_6x4095_inner0_wg256",
+      "provenance": {
+        "notes": "Odd suffix lengths cross invocation and block boundaries. Distinct bias rows distinguish cyclic reuse from consecutive row groups."
+      },
+      "attrs": { "axis": 1, "is_inner_broadcast": 0 },
+      "tunables": { "WORKGROUP_SIZE": 256 },
+      "inputs": {
+        "data": {
+          "dtype": "float32",
+          "shape": [6, 4095],
+          "data": { "kind": "fillFloat32", "scale": 0.5, "sinStep": 0.017, "cosStep": 0.031 }
+        },
+        "bias": {
+          "dtype": "float32",
+          "shape": [3, 4095],
+          "data": { "kind": "fillFloat32", "scale": 0.25, "sinStep": 0.037, "cosStep": 0.053 }
+        }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [6, 4095], "tolerance": 1e-7 } }
+    },
+    {
+      "name": "bias_row_tail_float32_4x65537_inner1_wg64",
+      "provenance": {
+        "notes": "Odd suffix lengths cross invocation and block boundaries. Distinct bias rows distinguish cyclic reuse from consecutive row groups."
+      },
+      "attrs": { "axis": 1, "is_inner_broadcast": 1 },
+      "tunables": { "WORKGROUP_SIZE": 64 },
+      "inputs": {
+        "data": {
+          "dtype": "float32",
+          "shape": [4, 65537],
+          "data": { "kind": "fillFloat32", "scale": 0.5, "sinStep": 0.017, "cosStep": 0.031 }
+        },
+        "bias": {
+          "dtype": "float32",
+          "shape": [2, 65537],
+          "data": { "kind": "fillFloat32", "scale": 0.25, "sinStep": 0.037, "cosStep": 0.053 }
+        }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [4, 65537], "tolerance": 1e-8 } }
+    },
+    {
+      "name": "bias_row_tail_float32_4x131073_inner0_wg256",
+      "provenance": {
+        "notes": "Odd suffix lengths cross invocation and block boundaries. Distinct bias rows distinguish cyclic reuse from consecutive row groups."
+      },
+      "attrs": { "axis": 1, "is_inner_broadcast": 0 },
+      "tunables": { "WORKGROUP_SIZE": 256 },
+      "inputs": {
+        "data": {
+          "dtype": "float32",
+          "shape": [4, 131073],
+          "data": { "kind": "fillFloat32", "scale": 0.5, "sinStep": 0.017, "cosStep": 0.031 }
+        },
+        "bias": {
+          "dtype": "float32",
+          "shape": [2, 131073],
+          "data": { "kind": "fillFloat32", "scale": 0.25, "sinStep": 0.037, "cosStep": 0.053 }
+        }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [4, 131073], "tolerance": 1e-8 } }
+    },
+    {
+      "name": "bias_row_tail_float16_6x513_inner1_wg256",
+      "provenance": {
+        "notes": "Odd suffix lengths cross invocation and block boundaries. Distinct bias rows distinguish cyclic reuse from consecutive row groups."
+      },
+      "attrs": { "axis": 1, "is_inner_broadcast": 1 },
+      "tunables": { "WORKGROUP_SIZE": 256 },
+      "inputs": {
+        "data": {
+          "dtype": "float16",
+          "shape": [6, 513],
+          "data": { "kind": "fillFloat32", "scale": 0.5, "sinStep": 0.017, "cosStep": 0.031 }
+        },
+        "bias": {
+          "dtype": "float16",
+          "shape": [2, 513],
+          "data": { "kind": "fillFloat32", "scale": 0.25, "sinStep": 0.037, "cosStep": 0.053 }
+        }
+      },
+      "outputs": { "output": { "dtype": "float16", "shape": [6, 513], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "bias_row_tail_float16_6x4095_inner0_wg256",
+      "provenance": {
+        "notes": "Odd suffix lengths cross invocation and block boundaries. Distinct bias rows distinguish cyclic reuse from consecutive row groups."
+      },
+      "attrs": { "axis": 1, "is_inner_broadcast": 0 },
+      "tunables": { "WORKGROUP_SIZE": 256 },
+      "inputs": {
+        "data": {
+          "dtype": "float16",
+          "shape": [6, 4095],
+          "data": { "kind": "fillFloat32", "scale": 0.5, "sinStep": 0.017, "cosStep": 0.031 }
+        },
+        "bias": {
+          "dtype": "float16",
+          "shape": [3, 4095],
+          "data": { "kind": "fillFloat32", "scale": 0.25, "sinStep": 0.037, "cosStep": 0.053 }
+        }
+      },
+      "outputs": { "output": { "dtype": "float16", "shape": [6, 4095], "tolerance": 0.000001 } }
+    },
+    {
+      "name": "bias_row_tail_float16_4x65537_inner1_wg64",
+      "provenance": {
+        "notes": "Odd suffix lengths cross invocation and block boundaries. Distinct bias rows distinguish cyclic reuse from consecutive row groups."
+      },
+      "attrs": { "axis": 1, "is_inner_broadcast": 1 },
+      "tunables": { "WORKGROUP_SIZE": 64 },
+      "inputs": {
+        "data": {
+          "dtype": "float16",
+          "shape": [4, 65537],
+          "data": { "kind": "fillFloat32", "scale": 0.5, "sinStep": 0.017, "cosStep": 0.031 }
+        },
+        "bias": {
+          "dtype": "float16",
+          "shape": [2, 65537],
+          "data": { "kind": "fillFloat32", "scale": 0.25, "sinStep": 0.037, "cosStep": 0.053 }
+        }
+      },
+      "outputs": { "output": { "dtype": "float16", "shape": [4, 65537], "tolerance": 1e-7 } }
+    },
+    {
+      "name": "bias_row_tail_float16_4x131073_inner0_wg256",
+      "provenance": {
+        "notes": "Odd suffix lengths cross invocation and block boundaries. Distinct bias rows distinguish cyclic reuse from consecutive row groups."
+      },
+      "attrs": { "axis": 1, "is_inner_broadcast": 0 },
+      "tunables": { "WORKGROUP_SIZE": 256 },
+      "inputs": {
+        "data": {
+          "dtype": "float16",
+          "shape": [4, 131073],
+          "data": { "kind": "fillFloat32", "scale": 0.5, "sinStep": 0.017, "cosStep": 0.031 }
+        },
+        "bias": {
+          "dtype": "float16",
+          "shape": [2, 131073],
+          "data": { "kind": "fillFloat32", "scale": 0.25, "sinStep": 0.037, "cosStep": 0.053 }
+        }
+      },
+      "outputs": { "output": { "dtype": "float16", "shape": [4, 131073], "tolerance": 1e-7 } }
+    },
+    {
+      "name": "subgroup_rows_float32_65x256_bias5_inner1",
+      "provenance": {
+        "notes": "Subgroup-row route boundary: row tail or register budget boundary with flattened bias-row reuse. Values are added and accumulated in float32 before the output cast."
+      },
+      "attrs": { "axis": 1, "is_inner_broadcast": 1 },
+      "inputs": {
+        "data": {
+          "dtype": "float32",
+          "shape": [65, 256],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 2.0 }
+        },
+        "bias": {
+          "dtype": "float32",
+          "shape": [5, 256],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23, "scale": 1.0 }
+        }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [65, 256], "tolerance": 0.00001, "relTolerance": 0.0001 } }
+    },
+    {
+      "name": "subgroup_rows_float32_64x1024_bias4_inner0",
+      "provenance": {
+        "notes": "Subgroup-row route boundary: row tail or register budget boundary with flattened bias-row reuse. Values are added and accumulated in float32 before the output cast."
+      },
+      "attrs": { "axis": 1, "is_inner_broadcast": 0 },
+      "inputs": {
+        "data": {
+          "dtype": "float32",
+          "shape": [64, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 2.0 }
+        },
+        "bias": {
+          "dtype": "float32",
+          "shape": [4, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23, "scale": 1.0 }
+        }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [64, 1024], "tolerance": 0.00001, "relTolerance": 0.0001 } }
+    },
+    {
+      "name": "subgroup_rows_float32_nonfinite_bias",
+      "attrs": { "axis": 1, "is_inner_broadcast": 1 },
+      "inputs": {
+        "data": { "dtype": "float32", "shape": [64, 64], "data": { "kind": "constant", "value": 0.5 } },
+        "bias": {
+          "dtype": "float32",
+          "shape": [4, 64],
+          "data": {
+            "kind": "values",
+            "values": { "$ref": "#/fixtureArrays/subgroup_rows_float32_nonfinite_bias_input_bias" }
+          }
+        }
+      },
+      "outputs": {
+        "output": {
+          "dtype": "float32",
+          "shape": [64, 64],
+          "allowNaN": true,
+          "tolerance": 0.00001,
+          "relTolerance": 0.0001
+        }
+      },
+      "provenance": {
+        "notes": "Subgroup-row bias fusion must retain the operator semantics for positive infinity, a NaN bias entry, and fully negative-infinity masked rows; this fixture checks their positions as well as finite neighboring rows."
+      }
+    },
+    {
+      "name": "subgroup_rows_float16_65x256_bias5_inner1",
+      "provenance": {
+        "notes": "Subgroup-row route boundary: row tail or register budget boundary with flattened bias-row reuse. Values are added and accumulated in float32 before the output cast."
+      },
+      "attrs": { "axis": 1, "is_inner_broadcast": 1 },
+      "inputs": {
+        "data": {
+          "dtype": "float16",
+          "shape": [65, 256],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 2.0 }
+        },
+        "bias": {
+          "dtype": "float16",
+          "shape": [5, 256],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23, "scale": 1.0 }
+        }
+      },
+      "outputs": { "output": { "dtype": "float16", "shape": [65, 256], "tolerance": 0.001, "relTolerance": 0.002 } }
+    },
+    {
+      "name": "subgroup_rows_float16_64x1024_bias4_inner0",
+      "provenance": {
+        "notes": "Subgroup-row route boundary: row tail or register budget boundary with flattened bias-row reuse. Values are added and accumulated in float32 before the output cast."
+      },
+      "attrs": { "axis": 1, "is_inner_broadcast": 0 },
+      "inputs": {
+        "data": {
+          "dtype": "float16",
+          "shape": [64, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 2.0 }
+        },
+        "bias": {
+          "dtype": "float16",
+          "shape": [4, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23, "scale": 1.0 }
+        }
+      },
+      "outputs": { "output": { "dtype": "float16", "shape": [64, 1024], "tolerance": 0.001, "relTolerance": 0.002 } }
+    },
+    {
+      "name": "subgroup_rows_float16_nonfinite_bias",
+      "attrs": { "axis": 1, "is_inner_broadcast": 0 },
+      "inputs": {
+        "data": { "dtype": "float16", "shape": [64, 64], "data": { "kind": "constant", "value": 0.5 } },
+        "bias": {
+          "dtype": "float16",
+          "shape": [4, 64],
+          "data": {
+            "kind": "values",
+            "values": { "$ref": "#/fixtureArrays/subgroup_rows_float32_nonfinite_bias_input_bias" }
+          }
+        }
+      },
+      "outputs": {
+        "output": { "dtype": "float16", "shape": [64, 64], "allowNaN": true, "tolerance": 0.001, "relTolerance": 0.002 }
+      },
+      "provenance": {
+        "notes": "Subgroup-row bias fusion must retain the operator semantics for positive infinity, a NaN bias entry, and fully negative-infinity masked rows; this fixture checks their positions as well as finite neighboring rows."
+      }
+    },
+    {
+      "name": "ragged_float32_negative_limit_cols33",
+      "attrs": { "axis": 1, "is_inner_broadcast": 1 },
+      "inputs": {
+        "data": {
+          "dtype": "float32",
+          "shape": [3, 33],
+          "data": { "kind": "constant", "value": -3.4028234663852886e+38 }
+        },
+        "bias": { "dtype": "float32", "shape": [1, 33], "data": { "kind": "constant", "value": 0.0 } }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [3, 33], "tolerance": 1e-7, "relTolerance": 0.00001 } },
+      "provenance": {
+        "notes": "Invalid logical-vec4 tail components must contribute exactly zero denominator, including when valid data equals the negative finite-limit sentinel."
+      }
+    },
+    {
+      "name": "ragged_float32_negative_limit_cols65",
+      "attrs": { "axis": 1, "is_inner_broadcast": 1 },
+      "inputs": {
+        "data": {
+          "dtype": "float32",
+          "shape": [3, 65],
+          "data": { "kind": "constant", "value": -3.4028234663852886e+38 }
+        },
+        "bias": { "dtype": "float32", "shape": [1, 65], "data": { "kind": "constant", "value": 0.0 } }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [3, 65], "tolerance": 1e-7, "relTolerance": 0.00001 } },
+      "provenance": {
+        "notes": "Invalid logical-vec4 tail components must contribute exactly zero denominator, including when valid data equals the negative finite-limit sentinel."
+      }
+    },
+    {
+      "name": "ragged_float32_negative_limit_cols67",
+      "attrs": { "axis": 1, "is_inner_broadcast": 1 },
+      "inputs": {
+        "data": {
+          "dtype": "float32",
+          "shape": [3, 67],
+          "data": { "kind": "constant", "value": -3.4028234663852886e+38 }
+        },
+        "bias": { "dtype": "float32", "shape": [1, 67], "data": { "kind": "constant", "value": 0.0 } }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [3, 67], "tolerance": 1e-7, "relTolerance": 0.00001 } },
+      "provenance": {
+        "notes": "Invalid logical-vec4 tail components must contribute exactly zero denominator, including when valid data equals the negative finite-limit sentinel."
+      }
+    },
+    {
+      "name": "ragged_float32_inner_broadcast_nonfinite_tail",
+      "attrs": { "axis": 1, "is_inner_broadcast": 1 },
+      "inputs": {
+        "data": {
+          "dtype": "float32",
+          "shape": [6, 67],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 2.0 }
+        },
+        "bias": {
+          "dtype": "float32",
+          "shape": [2, 67],
+          "data": {
+            "kind": "values",
+            "values": { "$ref": "#/fixtureArrays/ragged_float32_inner_broadcast_nonfinite_tail_input_bias" }
+          }
+        }
+      },
+      "outputs": {
+        "output": { "dtype": "float32", "shape": [6, 67], "tolerance": 1e-7, "relTolerance": 0.00001, "allowNaN": true }
+      },
+      "provenance": {
+        "notes": "Consecutive groups reuse each ragged bias row: first group has a NaN in the last valid component, second group is fully additive-masked. Check exact NaN positions and zero rows."
+      }
+    },
+    {
+      "name": "ragged_float16_negative_limit_cols33",
+      "attrs": { "axis": 1, "is_inner_broadcast": 1 },
+      "inputs": {
+        "data": { "dtype": "float16", "shape": [3, 33], "data": { "kind": "constant", "value": -65504.0 } },
+        "bias": { "dtype": "float16", "shape": [1, 33], "data": { "kind": "constant", "value": 0.0 } }
+      },
+      "outputs": { "output": { "dtype": "float16", "shape": [3, 33], "tolerance": 0.00001, "relTolerance": 0.001 } },
+      "provenance": {
+        "notes": "Invalid logical-vec4 tail components must contribute exactly zero denominator, including when valid data equals the negative finite-limit sentinel."
+      }
+    },
+    {
+      "name": "ragged_float16_negative_limit_cols65",
+      "attrs": { "axis": 1, "is_inner_broadcast": 1 },
+      "inputs": {
+        "data": { "dtype": "float16", "shape": [3, 65], "data": { "kind": "constant", "value": -65504.0 } },
+        "bias": { "dtype": "float16", "shape": [1, 65], "data": { "kind": "constant", "value": 0.0 } }
+      },
+      "outputs": { "output": { "dtype": "float16", "shape": [3, 65], "tolerance": 0.00001, "relTolerance": 0.001 } },
+      "provenance": {
+        "notes": "Invalid logical-vec4 tail components must contribute exactly zero denominator, including when valid data equals the negative finite-limit sentinel."
+      }
+    },
+    {
+      "name": "ragged_float16_negative_limit_cols67",
+      "attrs": { "axis": 1, "is_inner_broadcast": 1 },
+      "inputs": {
+        "data": { "dtype": "float16", "shape": [3, 67], "data": { "kind": "constant", "value": -65504.0 } },
+        "bias": { "dtype": "float16", "shape": [1, 67], "data": { "kind": "constant", "value": 0.0 } }
+      },
+      "outputs": { "output": { "dtype": "float16", "shape": [3, 67], "tolerance": 0.00001, "relTolerance": 0.001 } },
+      "provenance": {
+        "notes": "Invalid logical-vec4 tail components must contribute exactly zero denominator, including when valid data equals the negative finite-limit sentinel."
+      }
+    },
+    {
+      "name": "ragged_float16_inner_broadcast_nonfinite_tail",
+      "attrs": { "axis": 1, "is_inner_broadcast": 1 },
+      "inputs": {
+        "data": {
+          "dtype": "float16",
+          "shape": [6, 67],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 2.0 }
+        },
+        "bias": {
+          "dtype": "float16",
+          "shape": [2, 67],
+          "data": {
+            "kind": "values",
+            "values": { "$ref": "#/fixtureArrays/ragged_float32_inner_broadcast_nonfinite_tail_input_bias" }
+          }
+        }
+      },
+      "outputs": {
+        "output": {
+          "dtype": "float16",
+          "shape": [6, 67],
+          "tolerance": 0.00001,
+          "relTolerance": 0.001,
+          "allowNaN": true
+        }
+      },
+      "provenance": {
+        "notes": "Consecutive groups reuse each ragged bias row: first group has a NaN in the last valid component, second group is fully additive-masked. Check exact NaN positions and zero rows."
+      }
     }
   ]
 }