sync 2e7068faf55e

Browse files

Files changed (14) hide show

README.md +72 -0
build/webgpu/bench.json +151 -0
build/webgpu/manifest.json +571 -0
build/webgpu/metadata.json +26 -0
build/webgpu/softmax-longrow-normalize.wgsl.jinja +66 -0
build/webgpu/softmax-longrow-stats.wgsl.jinja +213 -0
build/webgpu/softmax-normalize.wgsl.jinja +123 -0
build/webgpu/softmax-online-packed-rows.wgsl.jinja +131 -0
build/webgpu/softmax-online.wgsl.jinja +252 -0
build/webgpu/softmax-row-stage-strided-vec4.wgsl.jinja +138 -0
build/webgpu/softmax-row-stage.wgsl.jinja +165 -0
build/webgpu/softmax-strided-online-lane.wgsl.jinja +116 -0
build/webgpu/softmax-strided-packed4-tail.wgsl.jinja +176 -0
build/webgpu/test.json +720 -0

README.md CHANGED Viewed

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

 ---
+library_name: kernels
 license: apache-2.0
+tags:
+- kernel
+- webgpu
+- wgsl
 ---
+# ai.onnx.LogSoftmax
+`ai.onnx`  ·  standard ONNX operator  ·  ONNX opset ≥ 13
+## Description
+Computes `log(softmax(input, axis))` along a single axis using a numerically stable shifted reduction. The output has the same shape as the input.
+See the [ONNX `LogSoftmax` spec](https://onnx.ai/onnx/operators/onnx__LogSoftmax.html) for the reference semantics.
+## Inputs
+| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `input` | `x` | `T` | — | — | The input tensor of rank >= 1. | required |
+## Outputs
+| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `output` | `y` | `T` | same as `input` | same as `input` | The log-softmax values; same shape as the input. | required |
+## Attributes
+Default values (overridable per request):
+| Attribute | Default | Description |
+| --- | --- | --- |
+| `axis` | `-1` | The axis along which log-softmax is computed. Negative values count from the end; the default `-1` operates over the last dimension. Accepted range is `[-r, r-1]` where `r` is the input rank. |
+## Type constraints
+| Variable | Allowed dtypes |
+| --- | --- |
+| `T` | `float32`, `float16` |
+## Files
+- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, provenance)
+- [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
+- [`test.json`](build/webgpu/test.json) — correctness cases
+- [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
+- [`softmax-longrow-normalize.wgsl.jinja`](build/webgpu/softmax-longrow-normalize.wgsl.jinja)
+- [`softmax-longrow-stats.wgsl.jinja`](build/webgpu/softmax-longrow-stats.wgsl.jinja)
+- [`softmax-normalize.wgsl.jinja`](build/webgpu/softmax-normalize.wgsl.jinja)
+- [`softmax-online-packed-rows.wgsl.jinja`](build/webgpu/softmax-online-packed-rows.wgsl.jinja)
+- [`softmax-online.wgsl.jinja`](build/webgpu/softmax-online.wgsl.jinja)
+- [`softmax-row-stage-strided-vec4.wgsl.jinja`](build/webgpu/softmax-row-stage-strided-vec4.wgsl.jinja)
+- [`softmax-row-stage.wgsl.jinja`](build/webgpu/softmax-row-stage.wgsl.jinja)
+- [`softmax-strided-online-lane.wgsl.jinja`](build/webgpu/softmax-strided-online-lane.wgsl.jinja)
+- [`softmax-strided-packed4-tail.wgsl.jinja`](build/webgpu/softmax-strided-packed4-tail.wgsl.jinja)
+## Use with `@huggingface/kernels`
+The loader derives every required output's shape and logical dtype from the manifest contract and this call.
+It then allocates the result tensors automatically.
+The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+Replace each `*Data` placeholder with a typed array containing the corresponding input data.
+```js
+import { getKernel } from "@huggingface/kernels";
+const kernel = await getKernel("webgpu-kernels/ai.onnx.LogSoftmax", { version: 1 });
+const { y } = await kernel({ x: { data: xData, shape: [1, 3] } });
+```

build/webgpu/bench.json ADDED Viewed

	@@ -0,0 +1,151 @@

+{
+  "op": "ai.onnx.LogSoftmax",
+  "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] },
+  "cases": [
+    {
+      "name": "attention_scores",
+      "inputs": { "x": { "dtype": "float32", "shape": [32, 128] } },
+      "outputs": { "y": { "dtype": "float32", "shape": [32, 128] } }
+    },
+    {
+      "name": "logsoftmax-f32-1024x2048",
+      "preset": "smoke",
+      "vars": { "dtype": "float32", "rows": 1024, "cols": 2048 },
+      "attrs": { "axis": 1 },
+      "inputs": { "x": { "shape": [1024, 2048], "dtype": "float32", "dist": "normal", "seed": 211, "scale": 2 } },
+      "outputs": { "y": { "shape": [1024, 2048], "dtype": "float32" } },
+      "bench": {
+        "primary": true,
+        "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * dtypeBytes(args.dtype) * 2" }]
+      }
+    },
+    {
+      "name": "logsoftmax-f32-16384x512",
+      "preset": "stress",
+      "vars": { "dtype": "float32", "rows": 16384, "cols": 512 },
+      "attrs": { "axis": 1 },
+      "inputs": { "x": { "shape": [16384, 512], "dtype": "float32", "dist": "normal", "seed": 212, "scale": 2 } },
+      "outputs": { "y": { "shape": [16384, 512], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * dtypeBytes(args.dtype) * 2" }] }
+    },
+    {
+      "name": "logsoftmax-logits-1x262144",
+      "preset": "model",
+      "vars": { "dtype": "float32", "rows": 1, "cols": 262144 },
+      "attrs": { "axis": 1 },
+      "inputs": { "x": { "shape": [1, 262144], "dtype": "float32", "dist": "normal", "seed": 213, "scale": 4 } },
+      "outputs": { "y": { "shape": [1, 262144], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * dtypeBytes(args.dtype) * 2" }] }
+    },
+    {
+      "name": "logsoftmax-f32-strided-axis1-64x1024x64",
+      "preset": "stress",
+      "vars": { "dtype": "float32", "rows": 64, "mid": 1024, "inner": 64 },
+      "attrs": { "axis": 1 },
+      "inputs": { "x": { "shape": [64, 1024, 64], "dtype": "float32", "dist": "normal", "seed": 311, "scale": 2 } },
+      "outputs": { "y": { "shape": [64, 1024, 64], "dtype": "float32" } },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "args.rows * args.mid * args.inner * dtypeBytes(args.dtype) * 2" }]
+      }
+    },
+    {
+      "name": "logsoftmax-f16-1024x2048-axis1",
+      "preset": "stress",
+      "vars": { "dtype": "float16", "rows": 1024, "cols": 2048 },
+      "attrs": { "axis": 1 },
+      "inputs": { "x": { "shape": [1024, 2048], "dtype": "float16", "dist": "normal", "seed": 312, "scale": 2 } },
+      "outputs": { "y": { "shape": [1024, 2048], "dtype": "float16" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * dtypeBytes(args.dtype) * 2" }] }
+    },
+    {
+      "name": "logsoftmax-f32-align-cliff-1024x2046",
+      "preset": "stress",
+      "vars": { "dtype": "float32", "rows": 1024, "cols": 2046 },
+      "attrs": { "axis": 1 },
+      "inputs": { "x": { "shape": [1024, 2046], "dtype": "float32", "dist": "normal", "seed": 313, "scale": 2 } },
+      "outputs": { "y": { "shape": [1024, 2046], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * dtypeBytes(args.dtype) * 2" }] }
+    },
+    {
+      "name": "logsoftmax-f32-tallwide-70000x2048",
+      "preset": "model",
+      "vars": { "dtype": "float32", "rows": 70000, "cols": 2048 },
+      "attrs": { "axis": 1 },
+      "inputs": { "x": { "shape": [70000, 2048], "dtype": "float32", "dist": "normal", "seed": 314, "scale": 2 } },
+      "outputs": { "y": { "shape": [70000, 2048], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * dtypeBytes(args.dtype) * 2" }] }
+    },
+    {
+      "name": "logsoftmax-f32-nosg-3pass-occupancy-131072x32",
+      "preset": "stress",
+      "vars": { "dtype": "float32", "rows": 131072, "cols": 32 },
+      "attrs": { "axis": 1 },
+      "inputs": { "x": { "shape": [131072, 32], "dtype": "float32", "dist": "normal", "seed": 411, "scale": 2 } },
+      "outputs": { "y": { "shape": [131072, 32], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * dtypeBytes(args.dtype) * 2" }] }
+    },
+    {
+      "name": "logsoftmax-f32-strided-axis1-64x2048x130-packed4-control",
+      "preset": "stress",
+      "provenance": {
+        "source": "strided-axis selector crossover",
+        "notes": "Below-crossover control for the large-inner tail pathology. At 8,320 independent strided rows the cooperative packed-4 reducer still beats the serial-lane capacity path."
+      },
+      "vars": { "dtype": "float32", "outer": 64, "axisDim": 2048, "inner": 130 },
+      "attrs": { "axis": 1 },
+      "inputs": { "x": { "shape": [64, 2048, 130], "dtype": "float32", "dist": "normal", "seed": 414, "scale": 2 } },
+      "outputs": { "y": { "shape": [64, 2048, 130], "dtype": "float32", "dist": "empty" } },
+      "bench": {
+        "metrics": [
+          { "type": "bandwidth", "value": "args.outer * args.axisDim * args.inner * dtypeBytes(args.dtype) * 2" }
+        ]
+      }
+    },
+    {
+      "name": "logsoftmax-f32-strided-axis1-96x2048x130-crossover",
+      "preset": "stress",
+      "provenance": {
+        "source": "strided-axis selector crossover",
+        "notes": "Above-crossover probe for the large-inner tail pathology. At 12,480 independent strided rows the serial-lane capacity path overtakes the cooperative packed-4 reducer."
+      },
+      "vars": { "dtype": "float32", "outer": 96, "axisDim": 2048, "inner": 130 },
+      "attrs": { "axis": 1 },
+      "inputs": { "x": { "shape": [96, 2048, 130], "dtype": "float32", "dist": "normal", "seed": 413, "scale": 2 } },
+      "outputs": { "y": { "shape": [96, 2048, 130], "dtype": "float32", "dist": "empty" } },
+      "bench": {
+        "metrics": [
+          { "type": "bandwidth", "value": "args.outer * args.axisDim * args.inner * dtypeBytes(args.dtype) * 2" }
+        ]
+      }
+    },
+    {
+      "name": "logsoftmax-f32-strided-axis1-120x2048x130-large-inner",
+      "preset": "stress",
+      "provenance": {
+        "source": "authored for variant coverage",
+        "notes": "Inner=130 exercises the two-element tail of the scalar-aligned packed-4 online reducer while sharing one workgroup across four adjacent rows."
+      },
+      "vars": { "dtype": "float32", "outer": 120, "axisDim": 2048, "inner": 130 },
+      "attrs": { "axis": 1 },
+      "inputs": { "x": { "shape": [120, 2048, 130], "dtype": "float32", "dist": "normal", "seed": 412, "scale": 2 } },
+      "outputs": { "y": { "shape": [120, 2048, 130], "dtype": "float32", "dist": "empty" } },
+      "bench": {
+        "metrics": [
+          { "type": "bandwidth", "value": "args.outer * args.axisDim * args.inner * dtypeBytes(args.dtype) * 2" }
+        ]
+      }
+    },
+    {
+      "name": "logsoftmax-f32-8192x512-midwidth",
+      "provenance": {
+        "notes": "Covers the 512-column midpoint between the 128- and 2048-column cases and probes the width term in online_wg_vec4's selector."
+      },
+      "preset": "smoke",
+      "vars": { "dtype": "float32", "rows": 8192, "cols": 512 },
+      "attrs": { "axis": 1 },
+      "inputs": { "x": { "shape": [8192, 512], "dtype": "float32", "dist": "normal", "seed": 215, "scale": 2 } },
+      "outputs": { "y": { "shape": [8192, 512], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "args.rows * args.cols * dtypeBytes(args.dtype) * 2" }] }
+    }
+  ]
+}

build/webgpu/manifest.json ADDED Viewed

	@@ -0,0 +1,571 @@

+{
+  "domain": "ai.onnx",
+  "name": "LogSoftmax",
+  "sinceVersion": 13,
+  "description": "Computes `log(softmax(input, axis))` along a single axis using a numerically stable shifted reduction. The output has the same shape as the input.",
+  "inputs": [{ "role": "input", "dtype": "T", "description": "The input tensor of rank >= 1." }],
+  "outputs": [
+    {
+      "role": "output",
+      "dtype": "T",
+      "rank": "ranks.input",
+      "description": "The log-softmax values; same shape as the input.",
+      "shape": "shapes.input"
+    }
+  ],
+  "attributes": { "axis": -1 },
+  "attributeDescriptions": {
+    "axis": "The axis along which log-softmax is computed. Negative values count from the end; the default `-1` operates over the last dimension. Accepted range is `[-r, r-1]` where `r` is the input rank."
+  },
+  "typeConstraints": { "T": ["float32", "float16"] },
+  "args": {
+    "x": { "kind": "tensor", "semantic": "input", "role": "input" },
+    "y": { "kind": "tensor", "semantic": "output", "role": "output" }
+  },
+  "tunables": {
+    "WORKGROUP_SIZE": 256,
+    "BLOCK_COLS": 1024,
+    "LONGROW_SPLIT_MIN_COLS": 65536,
+    "STRIDED_CAPACITY_MIN_ROWS": 12288
+  },
+  "derive": {
+    "storageBufferLimit": "min(device.limits.maxStorageBufferBindingSize, device.limits.maxBufferSize)",
+    "axis": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x",
+    "lastCols": "dim(shapes.x, ranks.x - 1)",
+    "lastRows": "rows(shapes.x, ranks.x - 1)",
+    "axisInner": "inner(shapes.x, axis)",
+    "axisOuter": "outer(shapes.x, axis)",
+    "axisRows": "axisOuter * axisInner",
+    "longrowBlocks": "ceilDiv(lastCols, tunables.BLOCK_COLS)",
+    "lastRowScratchFits": "4 * lastRows <= storageBufferLimit",
+    "axisRowScratchFits": "4 * axisRows <= storageBufferLimit",
+    "longrowBlockScratchFits": "4 * lastRows * longrowBlocks <= storageBufferLimit",
+    "lastDispatchFits": "ceilDiv(lastRows, device.limits.maxComputeWorkgroupsPerDimension) <= device.limits.maxComputeWorkgroupsPerDimension",
+    "axisDispatchFits": "ceilDiv(axisRows, device.limits.maxComputeWorkgroupsPerDimension) <= device.limits.maxComputeWorkgroupsPerDimension",
+    "axisThreadDispatchFits": "ceilDiv(ceilDiv(axisRows, tunables.WORKGROUP_SIZE), device.limits.maxComputeWorkgroupsPerDimension) <= device.limits.maxComputeWorkgroupsPerDimension",
+    "lastAxisOk": "numel(shapes.x) == numel(shapes.y) and ranks.x >= 1 and (attrs.axis == -1 or attrs.axis == ranks.x - 1) and f16Ok(dtypes.T)",
+    "stridedAxisOk": "numel(shapes.x) == numel(shapes.y) and ranks.x >= 2 and attrs.axis + ranks.x >= 0 and attrs.axis < ranks.x and not (attrs.axis == -1 or attrs.axis == ranks.x - 1)"
+  },
+  "bindingSets": {
+    "longrowBlockStats": [
+      { "name": "x", "arg": "x", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
+      { "name": "blockMax", "semantic": "blockMax", "buffer": { "type": "storage" }, "elementType": "f32" },
+      { "name": "blockSum", "semantic": "blockSum", "buffer": { "type": "storage" }, "elementType": "f32" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "cols", "type": "u32", "value": "lastCols" },
+            { "name": "blocks", "type": "u32", "value": "ceilDiv(lastCols, tunables.BLOCK_COLS)" }
+          ]
+        }
+      }
+    ],
+    "longrowRowStats": [
+      { "name": "blockMax", "semantic": "blockMax", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
+      { "name": "blockSum", "semantic": "blockSum", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
+      { "name": "rowMax", "semantic": "rowMax", "buffer": { "type": "storage" }, "elementType": "f32" },
+      { "name": "rowSum", "semantic": "rowSum", "buffer": { "type": "storage" }, "elementType": "f32" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [{ "name": "blocks", "type": "u32", "value": "ceilDiv(lastCols, tunables.BLOCK_COLS)" }]
+        }
+      }
+    ],
+    "longrowNormalize": [
+      { "name": "x", "arg": "x", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
+      { "name": "rowMax", "semantic": "rowMax", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
+      { "name": "rowSum", "semantic": "rowSum", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
+      { "name": "y", "arg": "y", "buffer": { "type": "storage" }, "elementType": "$T" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": { "name": "Params", "fields": [{ "name": "cols", "type": "u32", "value": "lastCols" }] }
+      }
+    ],
+    "online_vec4": [
+      {
+        "name": "x",
+        "arg": "x",
+        "semantic": "X",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$vectorScalar"
+      },
+      { "name": "y", "arg": "y", "buffer": { "type": "storage" }, "elementType": "$vectorScalar" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "rows", "type": "u32", "value": "lastRows" },
+            { "name": "vecCols", "type": "u32", "value": "lastCols / 4" }
+          ]
+        }
+      }
+    ],
+    "online_scalar": [
+      {
+        "name": "x",
+        "arg": "x",
+        "semantic": "input",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$scalar"
+      },
+      { "name": "y", "arg": "y", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "$scalar" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "rows", "type": "u32", "value": "lastRows" },
+            { "name": "cols", "type": "u32", "value": "lastCols" }
+          ]
+        }
+      }
+    ],
+    "stableReduceMax": [
+      { "name": "x", "arg": "x", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
+      { "name": "rowMax", "semantic": "rowMax", "buffer": { "type": "storage" }, "elementType": "f32" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "rows", "type": "u32", "value": "lastRows" },
+            { "name": "cols", "type": "u32", "value": "lastCols" }
+          ]
+        }
+      }
+    ],
+    "stableExpSum": [
+      { "name": "x", "arg": "x", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
+      { "name": "rowMax", "semantic": "rowMax", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
+      { "name": "y", "arg": "y", "buffer": { "type": "storage" }, "elementType": "$T" },
+      { "name": "rowSum", "semantic": "rowSum", "buffer": { "type": "storage" }, "elementType": "f32" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "rows", "type": "u32", "value": "lastRows" },
+            { "name": "cols", "type": "u32", "value": "lastCols" }
+          ]
+        }
+      }
+    ],
+    "normalize": [
+      { "name": "x", "arg": "x", "semantic": "input", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
+      { "name": "rowMax", "semantic": "rowMax", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
+      { "name": "rowSum", "semantic": "rowSum", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
+      { "name": "y", "arg": "y", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "$T" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.x)" }] }
+      }
+    ],
+    "stableNormalize": [
+      { "name": "x", "arg": "x", "semantic": "input", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
+      { "name": "rowMax", "semantic": "rowMax", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
+      { "name": "rowSum", "semantic": "rowSum", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
+      { "name": "y", "arg": "y", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "$T" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "count", "type": "u32", "value": "numel(shapes.x)" },
+            { "name": "cols", "type": "u32", "value": "lastCols" }
+          ]
+        }
+      }
+    ],
+    "stridedScalarReduce": [
+      { "name": "x", "arg": "x", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
+      { "name": "rowMax", "semantic": "rowMax", "buffer": { "type": "storage" }, "elementType": "f32" }
+    ],
+    "stridedScalarExpSum": [
+      { "name": "x", "arg": "x", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
+      { "name": "rowMax", "semantic": "rowMax", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
+      { "name": "y", "arg": "y", "buffer": { "type": "storage" }, "elementType": "$T" },
+      { "name": "rowSum", "semantic": "rowSum", "buffer": { "type": "storage" }, "elementType": "f32" }
+    ],
+    "stridedOnline": [
+      { "name": "x", "arg": "x", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
+      { "name": "y", "arg": "y", "buffer": { "type": "storage" }, "elementType": "$T" }
+    ]
+  },
+  "variants": [
+    {
+      "id": "longrow_split",
+      "priority": 30,
+      "when": ["lastCols >= tunables.LONGROW_SPLIT_MIN_COLS", "rows(shapes.x, ranks.x - 1) <= device.limits.maxComputeWorkgroupsPerDimension", "ceilDiv(dim(shapes.x, ranks.x - 1), tunables.BLOCK_COLS) <= device.limits.maxComputeWorkgroupsPerDimension", "lastRowScratchFits", "longrowBlockScratchFits", "lastAxisOk"],
+      "intermediates": [
+        {
+          "id": "blockMax",
+          "dtype": "float32",
+          "shape": "[rows(shapes.x, ranks.x - 1) * ceilDiv(dim(shapes.x, ranks.x - 1), tunables.BLOCK_COLS)]"
+        },
+        {
+          "id": "blockSum",
+          "dtype": "float32",
+          "shape": "[rows(shapes.x, ranks.x - 1) * ceilDiv(dim(shapes.x, ranks.x - 1), tunables.BLOCK_COLS)]"
+        },
+        { "id": "rowMax", "dtype": "float32", "shape": "[rows(shapes.x, ranks.x - 1)]" },
+        { "id": "rowSum", "dtype": "float32", "shape": "[rows(shapes.x, ranks.x - 1)]" }
+      ],
+      "passes": [
+        {
+          "id": "block_stats",
+          "name": "LogSoftmax.LongRowBlockStats",
+          "source": { "shader": "softmax-longrow-stats.wgsl.jinja", "inputs": { "stage": "\"block\"" } },
+          "bindings": "longrowBlockStats",
+          "dispatch": {
+            "x": "ceilDiv(dim(shapes.x, ranks.x - 1), tunables.BLOCK_COLS)",
+            "y": "rows(shapes.x, ranks.x - 1)"
+          }
+        },
+        {
+          "id": "row_stats",
+          "name": "LogSoftmax.LongRowStats",
+          "source": { "shader": "softmax-longrow-stats.wgsl.jinja", "inputs": { "stage": "\"row\"" } },
+          "bindings": "longrowRowStats",
+          "dispatch": { "x": "rows(shapes.x, ranks.x - 1)" }
+        },
+        {
+          "id": "normalize",
+          "name": "LogSoftmax.LongRowNormalize",
+          "source": { "shader": "softmax-longrow-normalize.wgsl.jinja", "inputs": { "op": "\"logsoftmax\"" } },
+          "bindings": "longrowNormalize",
+          "dispatch": {
+            "x": "ceilDiv(dim(shapes.x, ranks.x - 1), tunables.BLOCK_COLS)",
+            "y": "rows(shapes.x, ranks.x - 1)"
+          }
+        }
+      ]
+    },
+    {
+      "id": "online_wg_vec4",
+      "priority": 16,
+      "when": ["(device.features.has(\"subgroups\") and dim(shapes.x, ranks.x - 1) >= 1024) or (not device.features.has(\"subgroups\") and dim(shapes.x, ranks.x - 1) > 32)", "dim(shapes.x, ranks.x - 1) % 4 == 0", "ceilDiv(rows(shapes.x, ranks.x - 1), device.limits.maxComputeWorkgroupsPerDimension) <= device.limits.maxComputeWorkgroupsPerDimension", "lastAxisOk"],
+      "constants": {
+        "combineSubgroups": "device.features.has(\"subgroups\") and dim(shapes.x, ranks.x - 1) >= 1024",
+        "scalar": "dtypes.T",
+        "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
+        "usesF16": "dtypes.T == \"f16\""
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "LogSoftmax.OnlineWorkgroupVec4",
+          "source": { "shader": "softmax-online.wgsl.jinja", "inputs": { "op": "\"logsoftmax\"", "useVec4": "true" } },
+          "bindings": "online_vec4",
+          "dispatch": { "workgroups": "rows(shapes.x, ranks.x - 1)" }
+        }
+      ]
+    },
+    {
+      "id": "online_wg",
+      "priority": 15,
+      "when": ["(device.features.has(\"subgroups\") and dim(shapes.x, ranks.x - 1) > 0) or dim(shapes.x, ranks.x - 1) > 32", "ceilDiv(rows(shapes.x, ranks.x - 1), device.limits.maxComputeWorkgroupsPerDimension) <= device.limits.maxComputeWorkgroupsPerDimension", "lastAxisOk"],
+      "constants": {
+        "combineSubgroups": "device.features.has(\"subgroups\")",
+        "scalar": "dtypes.T",
+        "usesF16": "dtypes.T == \"f16\""
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "LogSoftmax.OnlineWorkgroup",
+          "source": { "shader": "softmax-online.wgsl.jinja", "inputs": { "op": "\"logsoftmax\"", "useVec4": "false" } },
+          "bindings": "online_scalar",
+          "dispatch": { "workgroups": "rows(shapes.x, ranks.x - 1)" }
+        }
+      ]
+    },
+    {
+      "id": "online_packed_rows_vec4_nosg",
+      "priority": 17,
+      "when": ["dim(shapes.x, ranks.x - 1) > 0", "dim(shapes.x, ranks.x - 1) <= 32", "dim(shapes.x, ranks.x - 1) % 4 == 0", "rows(shapes.x, ranks.x - 1) > 32", "lastAxisOk"],
+      "constants": {
+        "scalar": "dtypes.T",
+        "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
+        "usesF16": "dtypes.T == \"f16\"",
+        "lanesPerRow": 8,
+        "rowsPerWorkgroup": 32
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "LogSoftmax.OnlinePackedRowsVec4NoSg",
+          "source": {
+            "shader": "softmax-online-packed-rows.wgsl.jinja",
+            "inputs": {
+              "op": "\"logsoftmax\"",
+              "lanesPerRow": "constants.lanesPerRow",
+              "rowsPerWorkgroup": "constants.rowsPerWorkgroup"
+            }
+          },
+          "bindings": "online_vec4",
+          "dispatch": { "gridStride": "ceilDiv(rows(shapes.x, ranks.x - 1), constants.rowsPerWorkgroup)" }
+        }
+      ]
+    },
+    {
+      "id": "online_local_vec4_nosg",
+      "priority": 16,
+      "tunables": { "WORKGROUP_SIZE": 8 },
+      "when": ["dim(shapes.x, ranks.x - 1) > 0", "dim(shapes.x, ranks.x - 1) <= 32", "dim(shapes.x, ranks.x - 1) % 4 == 0", "ceilDiv(rows(shapes.x, ranks.x - 1), device.limits.maxComputeWorkgroupsPerDimension) <= device.limits.maxComputeWorkgroupsPerDimension", "lastAxisOk"],
+      "constants": {
+        "combineSubgroups": false,
+        "scalar": "dtypes.T",
+        "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
+        "usesF16": "dtypes.T == \"f16\""
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "LogSoftmax.OnlineLocalVec4NoSg",
+          "source": { "shader": "softmax-online.wgsl.jinja", "inputs": { "op": "\"logsoftmax\"", "useVec4": "true" } },
+          "bindings": "online_vec4",
+          "dispatch": { "workgroups": "rows(shapes.x, ranks.x - 1)" }
+        }
+      ]
+    },
+    {
+      "id": "online_local_nosg",
+      "priority": 14,
+      "tunables": { "WORKGROUP_SIZE": 32 },
+      "when": ["dim(shapes.x, ranks.x - 1) > 0", "dim(shapes.x, ranks.x - 1) <= 32", "ceilDiv(rows(shapes.x, ranks.x - 1), device.limits.maxComputeWorkgroupsPerDimension) <= device.limits.maxComputeWorkgroupsPerDimension", "lastAxisOk"],
+      "constants": { "combineSubgroups": false, "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "LogSoftmax.OnlineLocalNoSg",
+          "source": { "shader": "softmax-online.wgsl.jinja", "inputs": { "op": "\"logsoftmax\"", "useVec4": "false" } },
+          "bindings": "online_scalar",
+          "dispatch": { "workgroups": "rows(shapes.x, ranks.x - 1)" }
+        }
+      ]
+    },
+    {
+      "id": "stable_3pass",
+      "priority": 8,
+      "when": ["lastAxisOk", "lastDispatchFits", "lastRowScratchFits"],
+      "intermediates": [
+        { "id": "rowMax", "dtype": "float32", "shape": "[numel(shapes.x) / dim(shapes.x, ranks.x - 1)]" },
+        { "id": "rowSum", "dtype": "float32", "shape": "[numel(shapes.x) / dim(shapes.x, ranks.x - 1)]" }
+      ],
+      "passes": [
+        {
+          "id": "reduce_max",
+          "name": "LogSoftmax.ReduceMax",
+          "source": { "shader": "softmax-row-stage.wgsl.jinja", "inputs": { "stage": "\"max\"", "strided": "false" } },
+          "bindings": "stableReduceMax",
+          "dispatch": { "workgroups": "numel(shapes.x) / dim(shapes.x, ranks.x - 1)" }
+        },
+        {
+          "id": "exp_sum",
+          "name": "LogSoftmax.ExpSum",
+          "source": {
+            "shader": "softmax-row-stage.wgsl.jinja",
+            "inputs": { "stage": "\"exp_sum\"", "strided": "false" }
+          },
+          "bindings": "stableExpSum",
+          "dispatch": { "workgroups": "numel(shapes.x) / dim(shapes.x, ranks.x - 1)" }
+        },
+        {
+          "id": "normalize",
+          "name": "LogSoftmax.Normalize",
+          "source": {
+            "shader": "softmax-normalize.wgsl.jinja",
+            "inputs": { "strided": "false", "op": "\"logsoftmax\"" }
+          },
+          "bindings": "stableNormalize",
+          "dispatch": { "threads": "numel(shapes.x)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
+        }
+      ]
+    },
+    {
+      "id": "strided_3pass_vec4",
+      "priority": 14,
+      "derive": { "axis": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x" },
+      "when": ["dtypes.T == \"f32\"", "stridedAxisOk", "inner(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x) % 4 == 0", "axisDispatchFits", "axisRowScratchFits"],
+      "intermediates": [
+        { "id": "rowMax", "dtype": "float32", "shape": "[numel(shapes.x) / dim(shapes.x, axis)]" },
+        { "id": "rowSum", "dtype": "float32", "shape": "[numel(shapes.x) / dim(shapes.x, axis)]" }
+      ],
+      "passes": [
+        {
+          "id": "reduce_max",
+          "name": "LogSoftmax.StridedReduceMaxVec4",
+          "source": {
+            "shader": "softmax-row-stage-strided-vec4.wgsl.jinja",
+            "inputs": { "stage": "\"max\"", "xShape": "shapes.x", "axis": "axis" }
+          },
+          "bindings": [
+            {
+              "name": "x",
+              "arg": "x",
+              "semantic": "input",
+              "buffer": { "type": "read-only-storage" },
+              "elementType": "vec4<f32>"
+            },
+            { "name": "rowMax", "semantic": "rowMax", "buffer": { "type": "storage" }, "elementType": "vec4<f32>" }
+          ],
+          "dispatch": { "workgroups": "numel(shapes.x) / dim(shapes.x, axis) / 4" }
+        },
+        {
+          "id": "exp_sum",
+          "name": "LogSoftmax.StridedExpSumVec4",
+          "source": {
+            "shader": "softmax-row-stage-strided-vec4.wgsl.jinja",
+            "inputs": { "stage": "\"exp_sum\"", "xShape": "shapes.x", "axis": "axis" }
+          },
+          "bindings": [
+            {
+              "name": "x",
+              "arg": "x",
+              "semantic": "input",
+              "buffer": { "type": "read-only-storage" },
+              "elementType": "vec4<f32>"
+            },
+            {
+              "name": "rowMax",
+              "semantic": "rowMax",
+              "buffer": { "type": "read-only-storage" },
+              "elementType": "vec4<f32>"
+            },
+            {
+              "name": "y",
+              "arg": "y",
+              "semantic": "output",
+              "buffer": { "type": "storage" },
+              "elementType": "vec4<f32>"
+            },
+            { "name": "rowSum", "semantic": "rowSum", "buffer": { "type": "storage" }, "elementType": "vec4<f32>" }
+          ],
+          "dispatch": { "workgroups": "numel(shapes.x) / dim(shapes.x, axis) / 4" }
+        },
+        {
+          "id": "normalize",
+          "name": "LogSoftmax.StridedNormalize",
+          "source": {
+            "shader": "softmax-normalize.wgsl.jinja",
+            "inputs": { "strided": "true", "xShape": "shapes.x", "axis": "axis", "op": "\"logsoftmax\"" }
+          },
+          "bindings": "normalize",
+          "dispatch": { "threads": "numel(shapes.x)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
+        }
+      ]
+    },
+    {
+      "id": "strided_online_packed4_tail",
+      "priority": 18,
+      "derive": { "axis": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x" },
+      "when": ["dtypes.T == \"f32\"", "stridedAxisOk", "inner(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x) >= 4", "inner(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x) % 4 != 0", "dim(shapes.x, axis) < 2048 or axisRows < tunables.STRIDED_CAPACITY_MIN_ROWS", "axisDispatchFits"],
+      "passes": [
+        {
+          "id": "main",
+          "name": "LogSoftmax.StridedOnlinePacked4Tail",
+          "source": {
+            "shader": "softmax-strided-packed4-tail.wgsl.jinja",
+            "inputs": { "xShape": "shapes.x", "axis": "axis" }
+          },
+          "bindings": [
+            {
+              "name": "x",
+              "arg": "x",
+              "semantic": "input",
+              "buffer": { "type": "read-only-storage" },
+              "elementType": "f32"
+            },
+            { "name": "y", "arg": "y", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "f32" }
+          ],
+          "dispatch": {
+            "workgroups": "numel(shapes.x) / dim(shapes.x, axis) / inner(shapes.x, axis) * ceilDiv(inner(shapes.x, axis), 4)"
+          }
+        }
+      ]
+    },
+    {
+      "id": "strided_3pass",
+      "priority": 12,
+      "derive": { "axis": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x" },
+      "demoteWhen": ["dtypes.T == \"f32\" and inner(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x) >= 4 and inner(shapes.x, attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x) % 4 != 0"],
+      "when": ["stridedAxisOk", "f16Ok(dtypes.T)", "axisDispatchFits", "axisRowScratchFits"],
+      "intermediates": [
+        { "id": "rowMax", "dtype": "float32", "shape": "[numel(shapes.x) / dim(shapes.x, axis)]" },
+        { "id": "rowSum", "dtype": "float32", "shape": "[numel(shapes.x) / dim(shapes.x, axis)]" }
+      ],
+      "passes": [
+        {
+          "id": "reduce_max",
+          "name": "LogSoftmax.StridedReduceMax",
+          "source": {
+            "shader": "softmax-row-stage.wgsl.jinja",
+            "inputs": { "stage": "\"max\"", "strided": "true", "xShape": "shapes.x", "axis": "axis" }
+          },
+          "bindings": "stridedScalarReduce",
+          "dispatch": { "workgroups": "numel(shapes.x) / dim(shapes.x, axis)" }
+        },
+        {
+          "id": "exp_sum",
+          "name": "LogSoftmax.StridedExpSum",
+          "source": {
+            "shader": "softmax-row-stage.wgsl.jinja",
+            "inputs": { "stage": "\"exp_sum\"", "strided": "true", "xShape": "shapes.x", "axis": "axis" }
+          },
+          "bindings": "stridedScalarExpSum",
+          "dispatch": { "workgroups": "numel(shapes.x) / dim(shapes.x, axis)" }
+        },
+        {
+          "id": "normalize",
+          "name": "LogSoftmax.StridedNormalize",
+          "source": {
+            "shader": "softmax-normalize.wgsl.jinja",
+            "inputs": { "strided": "true", "xShape": "shapes.x", "axis": "axis", "op": "\"logsoftmax\"" }
+          },
+          "bindings": "normalize",
+          "dispatch": { "threads": "numel(shapes.x)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
+        }
+      ]
+    },
+    {
+      "id": "strided_online_capacity_fallback",
+      "priority": -1,
+      "when": ["stridedAxisOk", "f16Ok(dtypes.T)", "axisThreadDispatchFits"],
+      "passes": [
+        {
+          "id": "main",
+          "name": "LogSoftmax.StridedOnlineCapacityFallback",
+          "source": {
+            "shader": "softmax-strided-online-lane.wgsl.jinja",
+            "inputs": { "xShape": "shapes.x", "axis": "axis", "op": "\"logsoftmax\"" }
+          },
+          "bindings": "stridedOnline",
+          "dispatch": { "threads": "axisRows", "workgroupSize": "tunables.WORKGROUP_SIZE" }
+        }
+      ]
+    }
+  ]
+}

build/webgpu/metadata.json ADDED Viewed

	@@ -0,0 +1,26 @@

+{
+  "name": "ai.onnx.LogSoftmax",
+  "id": "_ai_onnx_logsoftmax_webgpu_0804fe4",
+  "version": 1,
+  "license": "Apache-2.0",
+  "backend": { "type": "webgpu" },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "bench.json": "rqFaX5XmJwQz0QQ1ET5FRMCuYEMotdcOJR+z6Y57a/0=",
+      "manifest.json": "QDbHBdv1drmn3+Mqxdqx24xatM7c38mQxNZrDzDPmM8=",
+      "softmax-longrow-normalize.wgsl.jinja": "AamVdoPDaWPQ3sk5loxtcuNu7oftC5J7hsPwLRDA7M0=",
+      "softmax-longrow-stats.wgsl.jinja": "K0PI+9cmRgjJsOOcBRyugN0Dm7yxBy0xvyPamEpCgxo=",
+      "softmax-normalize.wgsl.jinja": "eekHyttLyBk2X/mG54dUS2kIc5KycR3grH6000pqh1A=",
+      "softmax-online-packed-rows.wgsl.jinja": "94Xmz8Mtj/ji7XpIQUpsDdO6uShKJtWknzEI0Er+sT8=",
+      "softmax-online.wgsl.jinja": "qL/179d83HTFuyCUkfHqE0MBbtFMOEl8lzAH7aWmUXA=",
+      "softmax-row-stage-strided-vec4.wgsl.jinja": "l1byT6xoLl6Q7kSPd/uws6X5NsS2PwlSv2TBMCioVZA=",
+      "softmax-row-stage.wgsl.jinja": "dxlM6y02DwXeLeavHDiq70h0NCRZgz+bniCltyEgRX4=",
+      "softmax-strided-online-lane.wgsl.jinja": "BTLdB6LamPLNJsNmsllObl/9Y/Cx1EbuwdqeYpcppXg=",
+      "softmax-strided-packed4-tail.wgsl.jinja": "z7+TYeybaE/yXds/ODe4wMyosZeIrCtTfUmZsA2eIXo=",
+      "test.json": "uvZ0xakkYLGTptboBDjo8khdaQQ2h+QWiu4Z/wLWmfQ="
+    }
+  },
+  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
+  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.LogSoftmax" }
+}

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

	@@ -0,0 +1,66 @@

+// Final normalization pass for long-row split softmax. Each workgroup handles
+// one block of a row and re-reads the input against the completed row maximum
+// and denominator. The epilogue supports softmax and log-softmax.
+{% if T == "f16" %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
+const BLOCK_COLS: u32 = {{ tunables.BLOCK_COLS }}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);
+}
+// Bit-exact NaN test: `v != v` can be folded to false by fast-math backends,
+// and a constant NaN bit pattern is rejected at WGSL const-evaluation.
+fn is_nan_f32(value: f32) -> bool {
+  let bits = bitcast<u32>(value);
+  return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
+}
+@compute @workgroup_size(WG)
+fn main(@builtin(workgroup_id) wg: vec3<u32>,
+        @builtin(local_invocation_id) lid: vec3<u32>) {
+  let block = wg.x;
+  let row = wg.y;
+  let rowBase = row * params.cols;
+  let colStart = block * BLOCK_COLS;
+  let colEnd = min(params.cols, colStart + BLOCK_COLS);
+  let m = rowMax[row];
+  let d = rowSum[row];
+  // NaN row (any NaN input, or +inf max) -> all-NaN row, with the quiet NaN
+  // derived from the runtime value (a constant NaN bit pattern is rejected at
+  // WGSL const-evaluation).
+  let rowIsNan = m > FLT_MAX || is_nan_f32(d) || d == 0.0;
+  let nanOut = bitcast<f32>(bitcast<u32>(d) | 0x7fc00000u);
+  let logD = log(d);
+  for (var col = colStart + lid.x; col < colEnd; col = col + WG) {
+    let v = f32(x[rowBase + col]);
+    // (x - m) - log(d). x == m gives exactly 0 - log(d) (IEEE subtraction of
+    // equal finite values is +0), matching the serial kernel's shift trick.
+    var out = shifted_value(v, m) - logD;
+    if (rowIsNan) {
+      out = nanOut;
+    }
+    y[rowBase + col] = {{ T }}(out);
+  }
+}

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

	@@ -0,0 +1,213 @@

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

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

	@@ -0,0 +1,123 @@

+// Element-parallel pass 3 of three-pass softmax/log-softmax. Strided-axis mode
+// changes only how an element finds its row. Softmax divides the exponential
+// values already in y by the row sum. Log-softmax rereads x against the finished
+// row statistics and writes (x - rowMax) - log(rowSum), so entries whose
+// exponentials underflow remain finite.
+{% if T == "f16" %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+{% if source.strided %}
+{% set stridedAxisUses = ["row_for_offset"] %}
+// Compiled indexing for row-wise reductions over a non-last axis. A row is one
+// slice along the reduction axis, indexed as outer * INNER_SIZE + inner.
+{% set reduce_size = source.xShape[source.axis] %}
+{% set inner_size = namespace(value=1) %}
+{% for i in range(source.axis + 1, source.xShape | length) %}
+{% set inner_size.value = inner_size.value * source.xShape[i] %}
+{% endfor %}
+{% set outer_size = namespace(value=1) %}
+{% for i in range(source.axis) %}
+{% set outer_size.value = outer_size.value * source.xShape[i] %}
+{% endfor %}
+{% set usesAll = stridedAxisUses is not defined %}
+{% set usesElementOffset = usesAll or "element_offset" in stridedAxisUses %}
+{% set usesRowForOffset = usesAll or "row_for_offset" in stridedAxisUses %}
+{% set usesReduceSize = usesAll or usesElementOffset or "REDUCE_SIZE" in stridedAxisUses %}
+{% set usesInnerSize = usesAll or usesElementOffset or "INNER_SIZE" in stridedAxisUses %}
+{% set usesStridedRows = usesAll or "STRIDED_ROWS" in stridedAxisUses %}
+{% set usesInnerSizeSafe = usesAll or usesElementOffset or usesRowForOffset or "INNER_SIZE_SAFE" in stridedAxisUses %}
+{% set usesAxisStrideSafe = usesAll or usesRowForOffset or "AXIS_STRIDE_SAFE" in stridedAxisUses %}
+{% if usesReduceSize %}
+const REDUCE_SIZE: u32 = {{ reduce_size }}u;
+{% endif %}
+{% if usesInnerSize %}
+const INNER_SIZE: u32 = {{ inner_size.value }}u;
+{% endif %}
+{% if usesStridedRows %}
+// Total one-per-row workgroup count (product of every non-axis dim). Lets the
+// row-wise reduce kernels 2D-fold their dispatch past the maxComputeWorkgroupsPerDimension limit and guard
+// the over-dispatched tail without a runtime `params` uniform.
+const STRIDED_ROWS: u32 = {{ outer_size.value * inner_size.value }}u;
+{% endif %}
+{% if usesInnerSizeSafe %}
+const INNER_SIZE_SAFE: u32 = {{ inner_size.value if inner_size.value > 0 else 1 }}u;
+{% endif %}
+{% if usesAxisStrideSafe %}
+const AXIS_STRIDE_SAFE: u32 = {{ (reduce_size * inner_size.value) if (reduce_size * inner_size.value) > 0 else 1 }}u;
+{% endif %}
+{% if usesElementOffset %}
+// Flat offset of element k (along the reduce axis) of `row`.
+fn element_offset(row: u32, k: u32) -> u32 {
+  let outer = row / INNER_SIZE_SAFE;
+  let inner = row % INNER_SIZE_SAFE;
+  return outer * REDUCE_SIZE * INNER_SIZE + k * INNER_SIZE + inner;
+}
+{% endif %}
+{% if usesRowForOffset %}
+// Row that owns the flat element `offset` (inverse of element_offset over
+// the non-axis coordinates).
+fn row_for_offset(offset: u32) -> u32 {
+  let outer = offset / AXIS_STRIDE_SAFE;
+  let inner = offset % INNER_SIZE_SAFE;
+  return outer * INNER_SIZE_SAFE + inner;
+}
+{%- endif -%}
+{% endif %}
+// FLT_MAX, not -inf, as the online (m, d) accumulator init: merges must keep
+// `m - m` finite so an empty lane / all--inf row contributes the exact
+// accumulator identity (m, d) = (-FLT_MAX, 0). Operator epilogues interpret
+// a zero final denominator according to their public semantics. Using -inf
+// here changes +inf-row behavior.
+const FLT_MAX: f32 = 3.4028234663852886e38;
+fn is_finite_f32(value: f32) -> bool {
+  return select(false, value <= FLT_MAX, value >= -FLT_MAX);
+}
+// x - m that is exactly 0 when x equals a finite m, so exp(shifted) == 1
+// exactly at the row max. `x - x` on an infinite max is a legal fast-math
+// fold to 0, which would silently turn +inf rows finite — the explicit
+// equality test keeps the NaN propagation of the serial kernels.
+fn shifted_value(value: f32, maxValue: f32) -> f32 {
+  let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
+  return select(value - maxValue, 0.0, equalFiniteMax);
+}
+// Bit-exact NaN test: `v != v` can be folded to false by fast-math backends,
+// and a constant NaN bit pattern is rejected at WGSL const-evaluation.
+fn is_nan_f32(value: f32) -> bool {
+  let bits = bitcast<u32>(value);
+  return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
+}
+@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid : vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
+  // 2D-folded flat index: gid.y carries the high bits past the
+  // maxComputeWorkgroupsPerDimension dispatch limit (outputs > 16.7M elements).
+  let i = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
+  if (i >= params.count) {
+    return;
+  }
+{% if source.strided %}
+  let row = row_for_offset(i);
+{% else %}
+  let row = i / params.cols;
+{% endif %}
+  let m = rowMax[row];
+  let d = rowSum[row];
+  // NaN row (any NaN input, or +inf max) -> all-NaN row, quiet NaN derived from
+  // the runtime value (a constant NaN bit pattern is rejected at const-eval).
+  let rowIsNan = m > FLT_MAX || is_nan_f32(d) || d == 0.0;
+  let nanOut = bitcast<f32>(bitcast<u32>(d) | 0x7fc00000u);
+  // (x - m) - log(d). x == m gives exactly 0 - log(d) (IEEE subtraction of equal
+  // finite values is +0), matching the serial kernel's shift trick.
+  var out = shifted_value(f32(x[i]), m) - log(d);
+  if (rowIsNan) {
+    out = nanOut;
+  }
+  y[i] = {{ T }}(out);
+}

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

	@@ -0,0 +1,131 @@

+// Portable many-short-row online softmax / log-softmax. A 256-thread
+// workgroup is partitioned into independent fixed-width row groups (normally
+// 32 rows x 8 vec4 lanes). This preserves fully coalesced 32-column traffic
+// while one workgroup covers many short rows.
+{% if usesF16 %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+const LANES_PER_ROW: u32 = {{ source.lanesPerRow }}u;
+const ROWS_PER_WORKGROUP: u32 = {{ source.rowsPerWorkgroup }}u;
+const WG: u32 = LANES_PER_ROW * ROWS_PER_WORKGROUP;
+// 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));
+}
+// Bit-exact NaN test: `v != v` can be folded to false by fast-math backends,
+// and a constant NaN bit pattern is rejected at WGSL const-evaluation.
+fn is_nan_f32(value: f32) -> bool {
+  let bits = bitcast<u32>(value);
+  return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
+}
+// Each logical row owns one contiguous LANES_PER_ROW slice. All reductions
+// are workgroup-portable; no subgroup width or subgroup feature is assumed.
+var<workgroup> partialM: array<f32, WG>;
+var<workgroup> partialD: array<f32, WG>;
+@compute @workgroup_size({{ source.lanesPerRow * source.rowsPerWorkgroup }})
+fn main(
+  @builtin(workgroup_id) wid: vec3<u32>,
+  @builtin(num_workgroups) nwg: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>,
+) {
+  let lane = lid.x % LANES_PER_ROW;
+  let localRow = lid.x / LANES_PER_ROW;
+  let partialBase = localRow * LANES_PER_ROW;
+  let rowGroups = (params.rows + ROWS_PER_WORKGROUP - 1u) / ROWS_PER_WORKGROUP;
+  // A clamped grid-stride dispatch avoids rectangular over-dispatch once the
+  // row-group count exceeds one dispatch dimension while still covering every
+  // row of arbitrarily tall tensors.
+  for (var rowGroup = wid.x; rowGroup < rowGroups; rowGroup += nwg.x) {
+    let row = rowGroup * ROWS_PER_WORKGROUP + localRow;
+    let validRow = row < params.rows;
+    var m = -FLT_MAX;
+    var d = 0.0;
+    if (validRow) {
+      let rowBase = row * params.vecCols;
+      for (var c = lane; c < params.vecCols; c += LANES_PER_ROW) {
+        let v = vec4<f32>(x[rowBase + c]);
+        let mNew = max(m, max(max(v.x, v.y), max(v.z, v.w)));
+        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);
+        m = mNew;
+      }
+    }
+    let partialIndex = partialBase + lane;
+    partialM[partialIndex] = m;
+    partialD[partialIndex] = d;
+    workgroupBarrier();
+    var stride = LANES_PER_ROW / 2u;
+    loop {
+      if (stride == 0u) {
+        break;
+      }
+      if (lane < stride) {
+        let left = partialBase + lane;
+        let right = left + stride;
+        let m1 = partialM[left];
+        let d1 = partialD[left];
+        let m2 = partialM[right];
+        let d2 = partialD[right];
+        let mNew = max(m1, m2);
+        partialD[left] = d1 * exp_shift(m1, mNew) + d2 * exp_shift(m2, mNew);
+        partialM[left] = mNew;
+      }
+      workgroupBarrier();
+      stride /= 2u;
+    }
+    // The last reduction barrier publishes slot zero to every lane. Copy the
+    // result to private values before any lane can start the next row group.
+    let finalM = partialM[partialBase];
+    let finalD = partialD[partialBase];
+    let dBits = bitcast<u32>(finalD);
+    let rowIsNan = finalM > FLT_MAX || is_nan_f32(finalD) || finalD == 0.0;
+    let nanOut = bitcast<f32>(dBits | 0x7fc00000u);
+    if (validRow) {
+      let rowBase = row * params.vecCols;
+      let logD = log(finalD);
+      for (var c = lane; c < params.vecCols; c += LANES_PER_ROW) {
+        let v = vec4<f32>(x[rowBase + c]);
+        var o = vec4<f32>(
+          shifted_value(v.x, finalM),
+          shifted_value(v.y, finalM),
+          shifted_value(v.z, finalM),
+          shifted_value(v.w, finalM)) - logD;
+        if (rowIsNan) {
+          o = vec4<f32>(nanOut);
+        }
+        y[rowBase + c] = vec4<{{ scalar }}>(o);
+      }
+    }
+  }
+}

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

	@@ -0,0 +1,252 @@

+{% if usesF16 %}
+enable f16;
+{% endif %}
+{% if combineSubgroups %}
+enable subgroups;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+// Workgroup-level online softmax / log-softmax over one contiguous last-axis
+// row per workgroup. Each thread keeps an online (max, denom) pair over its
+// strided chunk, partials merge with the online-softmax rule (see
+// the shared online-softmax combiner) and a fused second sweep writes
+// the result. The epilogue is either softmax or log-softmax, and vec4 row access
+// is used when rows are vectorized.
+const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
+// FLT_MAX, not -inf, as the online (m, d) accumulator init: merges must keep
+// `m - m` finite so an empty lane / all--inf row contributes the exact
+// accumulator identity (m, d) = (-FLT_MAX, 0). Operator epilogues interpret
+// a zero final denominator according to their public semantics. Using -inf
+// here changes +inf-row behavior.
+const FLT_MAX: f32 = 3.4028234663852886e38;
+fn is_finite_f32(value: f32) -> bool {
+  return select(false, value <= FLT_MAX, value >= -FLT_MAX);
+}
+// x - m that is exactly 0 when x equals a finite m, so exp(shifted) == 1
+// exactly at the row max. `x - x` on an infinite max is a legal fast-math
+// fold to 0, which would silently turn +inf rows finite — the explicit
+// equality test keeps the NaN propagation of the serial kernels.
+fn shifted_value(value: f32, maxValue: f32) -> f32 {
+  let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
+  return select(value - maxValue, 0.0, equalFiniteMax);
+}
+fn exp_shift(value: f32, maxValue: f32) -> f32 {
+  return exp(shifted_value(value, maxValue));
+}
+// Bit-exact NaN test: `v != v` can be folded to false by fast-math backends,
+// and a constant NaN bit pattern is rejected at WGSL const-evaluation.
+fn is_nan_f32(value: f32) -> bool {
+  let bits = bitcast<u32>(value);
+  return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
+}
+// Workgroup-cooperative merge of per-thread online-softmax (m, d) partials:
+// mNew = max(m1, m2)
+// dNew = d1 * exp(m1 - mNew) + d2 * exp(m2 - mNew)
+// Both the subgroup and portable barrier-tree engines return the same merged
+// pair to every invocation. Repeated merges require a workgroup barrier between
+// calls before their shared partial storage is reused.
+{% set combineSubgroups = combineSubgroups is defined and combineSubgroups %}
+{% if combineSubgroups %}
+// Per-subgroup partials are published into a deterministic slot: the subgroup's
+// ordinal index within the workgroup (lidx / sgSize). The online (m, d) merge
+// is not float-associative, so thread 0 must fold partials in a fixed order.
+// Subgroups partition a workgroup into contiguous ordinal ranges on supported
+// backends, so the ordinal slot is unique per subgroup and every slot in
+// [0, subgroupCount) is written (each subgroup elects one leader).
+// Sized for the worst case of one partial per invocation.
+var<workgroup> partialM: array<f32, WG>;
+var<workgroup> partialD: array<f32, WG>;
+var<workgroup> combinedMD: vec2<f32>;
+// When the whole workgroup is one subgroup the subgroup reduce already covers
+// it (no barriers, no shared state); otherwise subgroup leaders publish
+// partials through shared memory and thread 0 folds them in ordinal order.
+fn combine_partials(m: f32, d: f32, lidx: u32, sgSize: u32) -> vec2<f32> {
+  let sgM = subgroupMax(m);
+  // A lane with no elements contributes d == 0 (exact identity). A +inf
+  // element made exp(inf - inf) = NaN stick in that lane's d; a NaN element
+  // landed in d via exp(NaN); both survive the merge and are detected by the
+  // code after the reduction.
+  let sgD = subgroupAdd(d * exp_shift(m, sgM));
+  if (sgSize == WG) {
+    return vec2<f32>(sgM, sgD);
+  }
+  let subgroupCount = (WG + sgSize - 1u) / sgSize;
+  // Pre-seed every fold slot with the (max, denom) identity. The fold below reads a
+  // fixed subgroupCount slots in ordinal order (for determinism), but a slot whose
+  // subgroup elects no leader this call — e.g. a fully out-of-window key tile in the
+  // flash-attention loop that re-uses this shared memory each iteration — would
+  // otherwise read stale shared memory. Identity makes such a slot a no-op.
+  // (max identity = -FLT_MAX, denom identity = 0.)
+  if (lidx < subgroupCount) {
+    partialM[lidx] = -FLT_MAX;
+    partialD[lidx] = 0.0;
+  }
+  workgroupBarrier();
+  if (subgroupElect()) {
+    let slot = lidx / sgSize;
+    partialM[slot] = sgM;
+    partialD[slot] = sgD;
+  }
+  workgroupBarrier();
+  if (lidx == 0u) {
+    var accM = -FLT_MAX;
+    var accD = 0.0;
+    for (var i = 0u; i < subgroupCount; i = i + 1u) {
+      let mNew = max(accM, partialM[i]);
+      accD = accD * exp_shift(accM, mNew) + partialD[i] * exp_shift(partialM[i], mNew);
+      accM = mNew;
+    }
+    combinedMD = vec2<f32>(accM, accD);
+  }
+  workgroupBarrier();
+  return combinedMD;
+}
+{% else %}
+{% set mdStreamed = mdStreams is defined %}
+{% set mdStreams = mdStreams if mdStreams is defined else 1 %}
+{% set mdExtent = "WG" if mdStreams == 1 else "WG * " ~ mdStreams ~ "u" %}
+var<workgroup> partialM: array<f32, {{ mdExtent }}>;
+var<workgroup> partialD: array<f32, {{ mdExtent }}>;
+{% if mdStreamed %}
+// In-place fold of {{ mdStreams }} streams. The caller stores its per-thread
+// partials into partialM/partialD first and reads the merged pair of stream s
+// from slot s * WG afterwards.
+fn combine_partials_streams(lidx: u32) {
+  workgroupBarrier();
+  var stride = WG / 2u;
+  loop {
+    if (stride == 0u) {
+      break;
+    }
+    if (lidx < stride) {
+{% for s in range(mdStreams) %}
+      {
+        let slot = {{ s }}u * WG + lidx;
+        let m1 = partialM[slot];
+        let d1 = partialD[slot];
+        let m2 = partialM[slot + stride];
+        let d2 = partialD[slot + stride];
+        let mNew = max(m1, m2);
+        partialD[slot] = d1 * exp_shift(m1, mNew) + d2 * exp_shift(m2, mNew);
+        partialM[slot] = mNew;
+      }
+{% endfor %}
+    }
+    workgroupBarrier();
+    stride = stride / 2u;
+  }
+}
+{% else %}
+fn combine_partials(m: f32, d: f32, lidx: u32) -> vec2<f32> {
+  partialM[lidx] = m;
+  partialD[lidx] = d;
+  workgroupBarrier();
+  var stride = WG / 2u;
+  loop {
+    if (stride == 0u) {
+      break;
+    }
+    if (lidx < stride) {
+      let m1 = partialM[lidx];
+      let d1 = partialD[lidx];
+      let m2 = partialM[lidx + stride];
+      let d2 = partialD[lidx + stride];
+      let mNew = max(m1, m2);
+      partialD[lidx] = d1 * exp_shift(m1, mNew) + d2 * exp_shift(m2, mNew);
+      partialM[lidx] = mNew;
+    }
+    workgroupBarrier();
+    stride = stride / 2u;
+  }
+  let merged = vec2<f32>(partialM[0], partialD[0]);
+  // Trailing barrier so back-to-back calls cannot race a next call's partial
+  // stores against this call's reads of slot 0.
+  workgroupBarrier();
+  return merged;
+}
+{% endif %}
+{% endif %}
+@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(workgroup_id) wid: vec3<u32>,
+        @builtin(num_workgroups) nwg: vec3<u32>,
+        @builtin(local_invocation_id) lid: vec3<u32>{% if combineSubgroups %},
+        @builtin(subgroup_size) sgSize: u32{% endif %}) {
+  // 2D-folded row index: wid.y carries the high bits past the
+  // maxComputeWorkgroupsPerDimension dispatch limit (rows > maxComputeWorkgroupsPerDimension).
+  let row = wid.x + wid.y * nwg.x;
+  if (row >= params.rows) {
+    return;
+  }
+{% if source.useVec4 %}
+  let rowBase = row * params.vecCols;
+{% else %}
+  let rowBase = row * params.cols;
+{% endif %}
+  // Thread-local online (m, d) over a strided chunk; f32 accumulation.
+  var m = -FLT_MAX;
+  var d = 0.0;
+{% if source.useVec4 %}
+  for (var c = lid.x; c < params.vecCols; c = c + WG) {
+    let v = vec4<f32>(x[rowBase + c]);
+    let mNew = max(m, max(max(v.x, v.y), max(v.z, v.w)));
+    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);
+    m = mNew;
+  }
+{% else %}
+  for (var c = lid.x; c < params.cols; c = c + WG) {
+    let v = f32(x[rowBase + c]);
+    let mNew = max(m, v);
+    d = d * exp_shift(m, mNew) + exp_shift(v, mNew);
+    m = mNew;
+  }
+{% endif %}
+  let rowMD = combine_partials(m, d, lid.x{% if combineSubgroups %}, sgSize{% endif %});
+  let finalM = rowMD.x;
+  let finalD = rowMD.y;
+  // Runtime-derived quiet NaN + bit-exact NaN test, as in the 1-pass kernel:
+  // `m - m` (Inf - Inf) and `d != d` are both legal fast-math folds (to 0 and
+  // false), and a constant NaN bit pattern is rejected at WGSL
+  // const-evaluation.
+  let dBits = bitcast<u32>(finalD);
+  let rowIsNan = finalM > FLT_MAX || is_nan_f32(finalD) || finalD == 0.0;
+  let nanOut = bitcast<f32>(dBits | 0x7fc00000u);
+  // (x - m) - log(d). x == m gives exactly 0 - log(d) (IEEE subtraction of
+  // equal finite values is +0), matching the serial kernel's shift trick.
+  let logD = log(finalD);
+{% if source.useVec4 %}
+  for (var c = lid.x; c < params.vecCols; c = c + WG) {
+    let v = vec4<f32>(x[rowBase + c]);
+    var o = vec4<f32>(
+      shifted_value(v.x, finalM),
+      shifted_value(v.y, finalM),
+      shifted_value(v.z, finalM),
+      shifted_value(v.w, finalM)) - logD;
+    if (rowIsNan) {
+      o = vec4<f32>(nanOut);
+    }
+    y[rowBase + c] = vec4<{{ scalar }}>(o);
+  }
+{% else %}
+  for (var c = lid.x; c < params.cols; c = c + WG) {
+    let v = f32(x[rowBase + c]);
+    var o = shifted_value(v, finalM) - logD;
+    if (rowIsNan) {
+      o = nanOut;
+    }
+    y[rowBase + c] = {{ scalar }}(o);
+  }
+{% endif %}
+}

build/webgpu/softmax-row-stage-strided-vec4.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,138 @@

+// Vec4-over-inner specialization for passes 1 and 2 of strided three-pass
+// softmax. A strided row reduces REDUCE_SIZE elements at INNER_SIZE stride, so
+// adjacent scalar threads would read non-coalesced locations. When INNER_SIZE
+// is divisible by four, four adjacent inner positions remain contiguous at
+// every axis coordinate. One workgroup therefore handles four rows packed as a
+// vec4: one coalesced 128-bit access per axis step through one reduction tree.
+// rowMax and rowSum use the same bytes the scalar normalize pass reads.
+{{ env.wgsl.resourceDeclarations }}
+const WGS: u32 = {{ tunables.WORKGROUP_SIZE }}u;
+{% set stridedAxisUses = ["REDUCE_SIZE", "INNER_SIZE", "STRIDED_ROWS"] %}
+// Compiled indexing for row-wise reductions over a non-last axis. A row is one
+// slice along the reduction axis, indexed as outer * INNER_SIZE + inner.
+{% set reduce_size = source.xShape[source.axis] %}
+{% set inner_size = namespace(value=1) %}
+{% for i in range(source.axis + 1, source.xShape | length) %}
+{% set inner_size.value = inner_size.value * source.xShape[i] %}
+{% endfor %}
+{% set outer_size = namespace(value=1) %}
+{% for i in range(source.axis) %}
+{% set outer_size.value = outer_size.value * source.xShape[i] %}
+{% endfor %}
+{% set usesAll = stridedAxisUses is not defined %}
+{% set usesElementOffset = usesAll or "element_offset" in stridedAxisUses %}
+{% set usesRowForOffset = usesAll or "row_for_offset" in stridedAxisUses %}
+{% set usesReduceSize = usesAll or usesElementOffset or "REDUCE_SIZE" in stridedAxisUses %}
+{% set usesInnerSize = usesAll or usesElementOffset or "INNER_SIZE" in stridedAxisUses %}
+{% set usesStridedRows = usesAll or "STRIDED_ROWS" in stridedAxisUses %}
+{% set usesInnerSizeSafe = usesAll or usesElementOffset or usesRowForOffset or "INNER_SIZE_SAFE" in stridedAxisUses %}
+{% set usesAxisStrideSafe = usesAll or usesRowForOffset or "AXIS_STRIDE_SAFE" in stridedAxisUses %}
+{% if usesReduceSize %}
+const REDUCE_SIZE: u32 = {{ reduce_size }}u;
+{% endif %}
+{% if usesInnerSize %}
+const INNER_SIZE: u32 = {{ inner_size.value }}u;
+{% endif %}
+{% if usesStridedRows %}
+// Total one-per-row workgroup count (product of every non-axis dim). Lets the
+// row-wise reduce kernels 2D-fold their dispatch past the maxComputeWorkgroupsPerDimension limit and guard
+// the over-dispatched tail without a runtime `params` uniform.
+const STRIDED_ROWS: u32 = {{ outer_size.value * inner_size.value }}u;
+{% endif %}
+{% if usesInnerSizeSafe %}
+const INNER_SIZE_SAFE: u32 = {{ inner_size.value if inner_size.value > 0 else 1 }}u;
+{% endif %}
+{% if usesAxisStrideSafe %}
+const AXIS_STRIDE_SAFE: u32 = {{ (reduce_size * inner_size.value) if (reduce_size * inner_size.value) > 0 else 1 }}u;
+{% endif %}
+{% if usesElementOffset %}
+// Flat offset of element k (along the reduce axis) of `row`.
+fn element_offset(row: u32, k: u32) -> u32 {
+  let outer = row / INNER_SIZE_SAFE;
+  let inner = row % INNER_SIZE_SAFE;
+  return outer * REDUCE_SIZE * INNER_SIZE + k * INNER_SIZE + inner;
+}
+{% endif %}
+{% if usesRowForOffset %}
+// Row that owns the flat element `offset` (inverse of element_offset over
+// the non-axis coordinates).
+fn row_for_offset(offset: u32) -> u32 {
+  let outer = offset / AXIS_STRIDE_SAFE;
+  let inner = offset % INNER_SIZE_SAFE;
+  return outer * INNER_SIZE_SAFE + inner;
+}
+{%- endif -%}
+{% if source.stage == "max" %}
+// max(1u,..) so an empty strided tensor (INNER_SIZE 0) still compiles — the kernel
+// returns before any divide (VEC4_ROWS 0 -> guard), this only avoids a const `/0`.
+const INNER4: u32 = max(1u, INNER_SIZE / 4u);  // vec4 groups of inner positions per outer
+const STRIDE4: u32 = INNER_SIZE / 4u;          // vec4 stride along the reduce axis
+const VEC4_ROWS: u32 = STRIDED_ROWS / 4u;      // = outer_size * INNER4
+{% else %}
+// max(1u,..) so an empty strided tensor (INNER_SIZE 0) still compiles (kernel
+// returns before any divide; only avoids a const `/0`).
+const INNER4: u32 = max(1u, INNER_SIZE / 4u);
+const STRIDE4: u32 = INNER_SIZE / 4u;
+const VEC4_ROWS: u32 = STRIDED_ROWS / 4u;
+{% endif %}
+var<workgroup> partial: array<vec4<f32>, WGS>;
+{% set reduce_fn = "reduce_max_vec4" if source.stage == "max" else "reduce_add_vec4" %}
+fn {{ reduce_fn }}(value: vec4<f32>, tid: u32) -> vec4<f32> {
+  partial[tid] = value;
+  workgroupBarrier();
+  var stride: u32 = WGS / 2u;
+  loop {
+    if (stride == 0u) { break; }
+{% if source.stage == "max" %}
+    if (tid < stride) { partial[tid] = max(partial[tid], partial[tid + stride]); }
+{% else %}
+    if (tid < stride) { partial[tid] = partial[tid] + partial[tid + stride]; }
+{% endif %}
+    stride = stride / 2u;
+    workgroupBarrier();
+  }
+  return partial[0];
+}
+@compute @workgroup_size(WGS)
+fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>,
+        @builtin(num_workgroups) nwg: vec3<u32>) {
+  let rg = wg.x + wg.y * nwg.x;  // vec4-row-group index
+  if (rg >= VEC4_ROWS) { return; }
+  let outer = rg / INNER4;
+  let inner_group = rg % INNER4;
+{% if source.stage == "max" %}
+  // vec4 index of element k=0 for this group (offset is a multiple of 4 since
+  // INNER_SIZE % 4 == 0 and inner_group*4 is too).
+{% endif %}
+  let base4 = (outer * REDUCE_SIZE * INNER_SIZE + inner_group * 4u) / 4u;
+{% if source.stage == "max" %}
+  var acc = vec4<f32>(-3.4028234663852886e38);
+  for (var k = lid.x; k < REDUCE_SIZE; k = k + WGS) {
+    acc = max(acc, x[base4 + k * STRIDE4]);
+  }
+  let rowM = reduce_max_vec4(acc, lid.x);
+  if (lid.x == 0u) { rowMax[rg] = rowM; }
+{% else %}
+  let m = rowMax[rg];
+  // +Inf row max -> NaN row (ONNX). Derive the quiet NaN per-component from the
+  // runtime bits (a constant NaN pattern is rejected at const-eval); m - m folds
+  // to 0 for finite m so the exp arg stays correct on the non-Inf lanes.
+  let nanv = bitcast<vec4<f32>>(bitcast<vec4<u32>>(m) | vec4<u32>(0x7fc00000u));
+  let isInf = m > vec4<f32>(3.4028234663852886e38);
+  var acc = vec4<f32>(0.0);
+  for (var k = lid.x; k < REDUCE_SIZE; k = k + WGS) {
+    let idx = base4 + k * STRIDE4;
+    let e = select(exp(x[idx] - m), nanv, isInf);
+    y[idx] = e;
+    acc = acc + e;
+  }
+  let rowSumValue = reduce_add_vec4(acc, lid.x);
+  if (lid.x == 0u) { rowSum[rg] = rowSumValue; }
+{% endif %}
+}

build/webgpu/softmax-row-stage.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,165 @@

+// First and second passes of the three-pass softmax, with one workgroup per
+// row. The max stage reduces the row maximum. The exp_sum stage writes
+// exp(x - rowMax) and reduces the row denominator.
+{% if T == "f16" %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+const WGS: u32 = {{ tunables.WORKGROUP_SIZE }}u;
+{% if source.strided %}
+{% set stridedAxisUses = ["STRIDED_ROWS", "element_offset"] %}
+// Compiled indexing for row-wise reductions over a non-last axis. A row is one
+// slice along the reduction axis, indexed as outer * INNER_SIZE + inner.
+{% set reduce_size = source.xShape[source.axis] %}
+{% set inner_size = namespace(value=1) %}
+{% for i in range(source.axis + 1, source.xShape | length) %}
+{% set inner_size.value = inner_size.value * source.xShape[i] %}
+{% endfor %}
+{% set outer_size = namespace(value=1) %}
+{% for i in range(source.axis) %}
+{% set outer_size.value = outer_size.value * source.xShape[i] %}
+{% endfor %}
+{% set usesAll = stridedAxisUses is not defined %}
+{% set usesElementOffset = usesAll or "element_offset" in stridedAxisUses %}
+{% set usesRowForOffset = usesAll or "row_for_offset" in stridedAxisUses %}
+{% set usesReduceSize = usesAll or usesElementOffset or "REDUCE_SIZE" in stridedAxisUses %}
+{% set usesInnerSize = usesAll or usesElementOffset or "INNER_SIZE" in stridedAxisUses %}
+{% set usesStridedRows = usesAll or "STRIDED_ROWS" in stridedAxisUses %}
+{% set usesInnerSizeSafe = usesAll or usesElementOffset or usesRowForOffset or "INNER_SIZE_SAFE" in stridedAxisUses %}
+{% set usesAxisStrideSafe = usesAll or usesRowForOffset or "AXIS_STRIDE_SAFE" in stridedAxisUses %}
+{% if usesReduceSize %}
+const REDUCE_SIZE: u32 = {{ reduce_size }}u;
+{% endif %}
+{% if usesInnerSize %}
+const INNER_SIZE: u32 = {{ inner_size.value }}u;
+{% endif %}
+{% if usesStridedRows %}
+// Total one-per-row workgroup count (product of every non-axis dim). Lets the
+// row-wise reduce kernels 2D-fold their dispatch past the maxComputeWorkgroupsPerDimension limit and guard
+// the over-dispatched tail without a runtime `params` uniform.
+const STRIDED_ROWS: u32 = {{ outer_size.value * inner_size.value }}u;
+{% endif %}
+{% if usesInnerSizeSafe %}
+const INNER_SIZE_SAFE: u32 = {{ inner_size.value if inner_size.value > 0 else 1 }}u;
+{% endif %}
+{% if usesAxisStrideSafe %}
+const AXIS_STRIDE_SAFE: u32 = {{ (reduce_size * inner_size.value) if (reduce_size * inner_size.value) > 0 else 1 }}u;
+{% endif %}
+{% if usesElementOffset %}
+// Flat offset of element k (along the reduce axis) of `row`.
+fn element_offset(row: u32, k: u32) -> u32 {
+  let outer = row / INNER_SIZE_SAFE;
+  let inner = row % INNER_SIZE_SAFE;
+  return outer * REDUCE_SIZE * INNER_SIZE + k * INNER_SIZE + inner;
+}
+{% endif %}
+{% if usesRowForOffset %}
+// Row that owns the flat element `offset` (inverse of element_offset over
+// the non-axis coordinates).
+fn row_for_offset(offset: u32) -> u32 {
+  let outer = offset / AXIS_STRIDE_SAFE;
+  let inner = offset % INNER_SIZE_SAFE;
+  return outer * INNER_SIZE_SAFE + inner;
+}
+{%- endif -%}
+{% endif %}
+{% macro wgsl_tree_reduce_f32(name, mode, buffer="partial", wg="WG", trailingBarrier=true) %}
+fn {{ name }}(value: f32, tid: u32) -> f32 {
+  {{ buffer }}[tid] = value;
+  workgroupBarrier();
+  // Ceil-halving keeps every lane when the workgroup size is not a power of
+  // two. For even n this matches the power-of-two tree order; for odd n, lanes
+  // [0, n-half) fold the upper tail while the middle lane carries forward.
+  var n: u32 = {{ wg }};
+  loop {
+    let half = (n + 1u) / 2u;
+    if (tid < n - half) {
+{% if mode == "max" %}
+      {{ buffer }}[tid] = max({{ buffer }}[tid], {{ buffer }}[tid + half]);
+{% else %}
+      {{ buffer }}[tid] = {{ buffer }}[tid] + {{ buffer }}[tid + half];
+{% endif %}
+    }
+    workgroupBarrier();
+    n = half;
+    if (n == 1u) {
+      break;
+    }
+  }
+  // The default trailing barrier makes this helper safe for back-to-back calls: every lane reads
+  // slot 0 here, so the next call's first store must not run until all lanes have read it.
+  // `trailingBarrier=false` is safe only when the buffer is never written again before kernel exit.
+  let reduced = {{ buffer }}[0];
+{% if trailingBarrier %}
+  workgroupBarrier();
+{% endif %}
+  return reduced;
+}
+{% endmacro %}
+var<workgroup> partial: array<f32, WGS>;
+{% if source.stage == "max" %}
+{{ wgsl_tree_reduce_f32("reduce_max_f32", "max", "partial", "WGS") }}
+{% else %}
+{{ wgsl_tree_reduce_f32("reduce_add_f32", "add", "partial", "WGS") }}
+{% endif %}
+{% set count = "REDUCE_SIZE" if source.strided else "params.cols" %}
+{% set offset_of = "element_offset(row, k)" if source.strided else "row * params.cols + k" %}
+@compute @workgroup_size(WGS)
+fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>,
+        @builtin(num_workgroups) nwg: vec3<u32>) {
+{% if source.strided %}
+  // Strided mode has no `params` uniform (offsets come from compile-time
+  // constants), so it folds against the compile-time STRIDED_ROWS instead of
+  // params.rows. wg.y carries the high bits of the row past the maxComputeWorkgroupsPerDimension limit;
+{% if source.stage == "max" %}
+  // the guard drops the over-dispatched tail before any rowMax[row] store.
+{% else %}
+  // the guard drops the over-dispatched tail before any rowSum[row]/y[offset]
+  // store.
+{% endif %}
+  let row = wg.x + wg.y * nwg.x;
+  if (row >= STRIDED_ROWS) { return; }
+{% else %}
+  // 2D-folded row index: wg.y carries the high bits past the maxComputeWorkgroupsPerDimension
+{% if source.stage == "max" %}
+  // workgroup-per-dimension dispatch limit (rows > maxComputeWorkgroupsPerDimension). The guard returns
+  // the whole over-dispatched workgroup before any read/write — without it the
+  // clamped OOB rowMax[row] store would corrupt the last real row. wg.x/wg.y/
+{% else %}
+  // workgroup-per-dimension dispatch limit (rows > maxComputeWorkgroupsPerDimension). Return the whole
+  // over-dispatched workgroup before any read/write — without it the clamped
+  // OOB rowSum[row]/y[offset] stores would corrupt the last real row. wg.x/wg.y/
+{% endif %}
+  // nwg.x are workgroup-uniform, so the return cannot diverge across a barrier.
+  let row = wg.x + wg.y * nwg.x;
+  if (row >= params.rows) { return; }
+{% endif %}
+{% if source.stage == "max" %}
+  var acc = -3.4028234663852886e38;
+  for (var k = lid.x; k < {{ count }}; k = k + WGS) {
+    acc = max(acc, f32(x[{{ offset_of }}]));
+  }
+  let rowM = reduce_max_f32(acc, lid.x);
+  if (lid.x == 0u) { rowMax[row] = rowM; }
+{% else %}
+  let m = rowMax[row];
+  var acc = 0.0;
+  for (var k = lid.x; k < {{ count }}; k = k + WGS) {
+    let offset = {{ offset_of }};
+    // +Inf row max must yield NaN (ONNX semantics). `m - m` (Inf - Inf) is
+    // folded to 0 by fast-math backends, and a constant NaN bit pattern is
+    // rejected at WGSL const-evaluation — so derive the quiet NaN from the
+    // runtime value (any bits | 0x7fc00000 is a NaN).
+    let e = select(exp(f32(x[offset]) - m), bitcast<f32>(bitcast<u32>(m) | 0x7fc00000u), m > 3.4028234663852886e38);
+    y[offset] = {{ T }}(e);
+    acc = acc + e;
+  }
+  let rowSumValue = reduce_add_f32(acc, lid.x);
+  if (lid.x == 0u) { rowSum[row] = rowSumValue; }
+{% endif %}
+}

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

	@@ -0,0 +1,116 @@

+// Scratch-free online softmax/log-softmax for a non-last reduction axis. One
+// invocation owns one (outer, inner) row, so neighboring lanes make coalesced
+// reads at every reduction step. This is both the large-inner fast path and the
+// capacity-safe fallback when row-stat scratch cannot be allocated.
+{% set operation = source.op if source.op is defined else "softmax" %}
+{% if T == "f16" %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+{% set stridedAxisUses = ["STRIDED_ROWS", "element_offset"] %}
+// Compiled indexing for row-wise reductions over a non-last axis. A row is one
+// slice along the reduction axis, indexed as outer * INNER_SIZE + inner.
+{% set reduce_size = source.xShape[source.axis] %}
+{% set inner_size = namespace(value=1) %}
+{% for i in range(source.axis + 1, source.xShape | length) %}
+{% set inner_size.value = inner_size.value * source.xShape[i] %}
+{% endfor %}
+{% set outer_size = namespace(value=1) %}
+{% for i in range(source.axis) %}
+{% set outer_size.value = outer_size.value * source.xShape[i] %}
+{% endfor %}
+{% set usesAll = stridedAxisUses is not defined %}
+{% set usesElementOffset = usesAll or "element_offset" in stridedAxisUses %}
+{% set usesRowForOffset = usesAll or "row_for_offset" in stridedAxisUses %}
+{% set usesReduceSize = usesAll or usesElementOffset or "REDUCE_SIZE" in stridedAxisUses %}
+{% set usesInnerSize = usesAll or usesElementOffset or "INNER_SIZE" in stridedAxisUses %}
+{% set usesStridedRows = usesAll or "STRIDED_ROWS" in stridedAxisUses %}
+{% set usesInnerSizeSafe = usesAll or usesElementOffset or usesRowForOffset or "INNER_SIZE_SAFE" in stridedAxisUses %}
+{% set usesAxisStrideSafe = usesAll or usesRowForOffset or "AXIS_STRIDE_SAFE" in stridedAxisUses %}
+{% if usesReduceSize %}
+const REDUCE_SIZE: u32 = {{ reduce_size }}u;
+{% endif %}
+{% if usesInnerSize %}
+const INNER_SIZE: u32 = {{ inner_size.value }}u;
+{% endif %}
+{% if usesStridedRows %}
+// Total one-per-row workgroup count (product of every non-axis dim). Lets the
+// row-wise reduce kernels 2D-fold their dispatch past the maxComputeWorkgroupsPerDimension limit and guard
+// the over-dispatched tail without a runtime `params` uniform.
+const STRIDED_ROWS: u32 = {{ outer_size.value * inner_size.value }}u;
+{% endif %}
+{% if usesInnerSizeSafe %}
+const INNER_SIZE_SAFE: u32 = {{ inner_size.value if inner_size.value > 0 else 1 }}u;
+{% endif %}
+{% if usesAxisStrideSafe %}
+const AXIS_STRIDE_SAFE: u32 = {{ (reduce_size * inner_size.value) if (reduce_size * inner_size.value) > 0 else 1 }}u;
+{% endif %}
+{% if usesElementOffset %}
+// Flat offset of element k (along the reduce axis) of `row`.
+fn element_offset(row: u32, k: u32) -> u32 {
+  let outer = row / INNER_SIZE_SAFE;
+  let inner = row % INNER_SIZE_SAFE;
+  return outer * REDUCE_SIZE * INNER_SIZE + k * INNER_SIZE + inner;
+}
+{% endif %}
+{% if usesRowForOffset %}
+// Row that owns the flat element `offset` (inverse of element_offset over
+// the non-axis coordinates).
+fn row_for_offset(offset: u32) -> u32 {
+  let outer = offset / AXIS_STRIDE_SAFE;
+  let inner = offset % INNER_SIZE_SAFE;
+  return outer * INNER_SIZE_SAFE + inner;
+}
+{%- endif -%}
+// FLT_MAX, not -inf, as the online (m, d) accumulator init: merges must keep
+// `m - m` finite so an empty lane / all--inf row contributes the exact
+// accumulator identity (m, d) = (-FLT_MAX, 0). Operator epilogues interpret
+// a zero final denominator according to their public semantics. Using -inf
+// here changes +inf-row behavior.
+const FLT_MAX: f32 = 3.4028234663852886e38;
+fn is_finite_f32(value: f32) -> bool {
+  return select(false, value <= FLT_MAX, value >= -FLT_MAX);
+}
+// x - m that is exactly 0 when x equals a finite m, so exp(shifted) == 1
+// exactly at the row max. `x - x` on an infinite max is a legal fast-math
+// fold to 0, which would silently turn +inf rows finite — the explicit
+// equality test keeps the NaN propagation of the serial kernels.
+fn shifted_value(value: f32, maxValue: f32) -> f32 {
+  let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
+  return select(value - maxValue, 0.0, equalFiniteMax);
+}
+const WGS: u32 = {{ tunables.WORKGROUP_SIZE }}u;
+@compute @workgroup_size(WGS)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>,
+        @builtin(num_workgroups) nwg: vec3<u32>) {
+  let row = gid.x + gid.y * nwg.x * WGS;
+  if (row >= STRIDED_ROWS) { return; }
+  var row_max = -{{ "FLT_MAX" if operation == "logsoftmax" else "F32_MAX" }};
+  var row_sum = 0.0;
+  for (var k = 0u; k < REDUCE_SIZE; k += 1u) {
+    let value = f32(x[element_offset(row, k)]);
+    if (value > row_max) {
+      row_sum = row_sum * exp(row_max - value) + 1.0;
+      row_max = value;
+    } else {
+      row_sum += exp(value - row_max);
+    }
+  }
+  let sum_bits = bitcast<u32>(row_sum);
+  let sum_is_nan = (sum_bits & 0x7f800000u) == 0x7f800000u && (sum_bits & 0x007fffffu) != 0u;
+  for (var k = 0u; k < REDUCE_SIZE; k += 1u) {
+    let offset = element_offset(row, k);
+    var out = shifted_value(f32(x[offset]), row_max) - log(row_sum);
+    if (row_max > FLT_MAX || sum_is_nan) {
+      out = bitcast<f32>(bitcast<u32>(row_sum) | 0x7fc00000u);
+    }
+    y[offset] = {{ T }}(out);
+  }
+}

build/webgpu/softmax-strided-packed4-tail.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,176 @@

+// Online LogSoftmax for a non-last axis whose contiguous inner dimension has a
+// one-to-three-element tail. Scalar bindings keep every outer slice legal even
+// when its starting address is not vec4-aligned. Four adjacent inner rows share
+// one workgroup and vec4 reduction tree, and uniform guards suppress invalid
+// components of the final packed group.
+{{ env.wgsl.resourceDeclarations }}
+const WGS: u32 = {{ tunables.WORKGROUP_SIZE }}u;
+{% set stridedAxisUses = ["REDUCE_SIZE", "INNER_SIZE", "STRIDED_ROWS", "INNER_SIZE_SAFE"] %}
+// Compiled indexing for row-wise reductions over a non-last axis. A row is one
+// slice along the reduction axis, indexed as outer * INNER_SIZE + inner.
+{% set reduce_size = source.xShape[source.axis] %}
+{% set inner_size = namespace(value=1) %}
+{% for i in range(source.axis + 1, source.xShape | length) %}
+{% set inner_size.value = inner_size.value * source.xShape[i] %}
+{% endfor %}
+{% set outer_size = namespace(value=1) %}
+{% for i in range(source.axis) %}
+{% set outer_size.value = outer_size.value * source.xShape[i] %}
+{% endfor %}
+{% set usesAll = stridedAxisUses is not defined %}
+{% set usesElementOffset = usesAll or "element_offset" in stridedAxisUses %}
+{% set usesRowForOffset = usesAll or "row_for_offset" in stridedAxisUses %}
+{% set usesReduceSize = usesAll or usesElementOffset or "REDUCE_SIZE" in stridedAxisUses %}
+{% set usesInnerSize = usesAll or usesElementOffset or "INNER_SIZE" in stridedAxisUses %}
+{% set usesStridedRows = usesAll or "STRIDED_ROWS" in stridedAxisUses %}
+{% set usesInnerSizeSafe = usesAll or usesElementOffset or usesRowForOffset or "INNER_SIZE_SAFE" in stridedAxisUses %}
+{% set usesAxisStrideSafe = usesAll or usesRowForOffset or "AXIS_STRIDE_SAFE" in stridedAxisUses %}
+{% if usesReduceSize %}
+const REDUCE_SIZE: u32 = {{ reduce_size }}u;
+{% endif %}
+{% if usesInnerSize %}
+const INNER_SIZE: u32 = {{ inner_size.value }}u;
+{% endif %}
+{% if usesStridedRows %}
+// Total one-per-row workgroup count (product of every non-axis dim). Lets the
+// row-wise reduce kernels 2D-fold their dispatch past the maxComputeWorkgroupsPerDimension limit and guard
+// the over-dispatched tail without a runtime `params` uniform.
+const STRIDED_ROWS: u32 = {{ outer_size.value * inner_size.value }}u;
+{% endif %}
+{% if usesInnerSizeSafe %}
+const INNER_SIZE_SAFE: u32 = {{ inner_size.value if inner_size.value > 0 else 1 }}u;
+{% endif %}
+{% if usesAxisStrideSafe %}
+const AXIS_STRIDE_SAFE: u32 = {{ (reduce_size * inner_size.value) if (reduce_size * inner_size.value) > 0 else 1 }}u;
+{% endif %}
+{% if usesElementOffset %}
+// Flat offset of element k (along the reduce axis) of `row`.
+fn element_offset(row: u32, k: u32) -> u32 {
+  let outer = row / INNER_SIZE_SAFE;
+  let inner = row % INNER_SIZE_SAFE;
+  return outer * REDUCE_SIZE * INNER_SIZE + k * INNER_SIZE + inner;
+}
+{% endif %}
+{% if usesRowForOffset %}
+// Row that owns the flat element `offset` (inverse of element_offset over
+// the non-axis coordinates).
+fn row_for_offset(offset: u32) -> u32 {
+  let outer = offset / AXIS_STRIDE_SAFE;
+  let inner = offset % INNER_SIZE_SAFE;
+  return outer * INNER_SIZE_SAFE + inner;
+}
+{%- endif -%}
+// FLT_MAX, not -inf, as the online (m, d) accumulator init: merges must keep
+// `m - m` finite so an empty lane / all--inf row contributes the exact
+// accumulator identity (m, d) = (-FLT_MAX, 0). Operator epilogues interpret
+// a zero final denominator according to their public semantics. Using -inf
+// here changes +inf-row behavior.
+const FLT_MAX: f32 = 3.4028234663852886e38;
+fn is_finite_f32(value: f32) -> bool {
+  return select(false, value <= FLT_MAX, value >= -FLT_MAX);
+}
+// x - m that is exactly 0 when x equals a finite m, so exp(shifted) == 1
+// exactly at the row max. `x - x` on an infinite max is a legal fast-math
+// fold to 0, which would silently turn +inf rows finite — the explicit
+// equality test keeps the NaN propagation of the serial kernels.
+fn shifted_value(value: f32, maxValue: f32) -> f32 {
+  let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
+  return select(value - maxValue, 0.0, equalFiniteMax);
+}
+fn exp_shift(value: f32, maxValue: f32) -> f32 {
+  return exp(shifted_value(value, maxValue));
+}
+// Bit-exact NaN test: `v != v` can be folded to false by fast-math backends,
+// and a constant NaN bit pattern is rejected at WGSL const-evaluation.
+fn is_nan_f32(value: f32) -> bool {
+  let bits = bitcast<u32>(value);
+  return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
+}
+const PACKED_INNER: u32 = (INNER_SIZE + 3u) / 4u;
+const OUTER_SIZE: u32 = STRIDED_ROWS / INNER_SIZE_SAFE;
+const PACKED_ROWS: u32 = OUTER_SIZE * PACKED_INNER;
+var<workgroup> partial: array<vec4<f32>, WGS>;
+fn reduce_max4(value: vec4<f32>, tid: u32) -> vec4<f32> {
+  partial[tid] = value;
+  workgroupBarrier();
+  for (var stride = WGS >> 1u; stride > 0u; stride = stride >> 1u) {
+    if (tid < stride) {
+      partial[tid] = max(partial[tid], partial[tid + stride]);
+    }
+    workgroupBarrier();
+  }
+  return partial[0];
+}
+fn reduce_add4(value: vec4<f32>, tid: u32) -> vec4<f32> {
+  partial[tid] = value;
+  workgroupBarrier();
+  for (var stride = WGS >> 1u; stride > 0u; stride = stride >> 1u) {
+    if (tid < stride) {
+      partial[tid] = partial[tid] + partial[tid + stride];
+    }
+    workgroupBarrier();
+  }
+  return partial[0];
+}
+fn stable_exp(value: f32, maximum: f32) -> f32 {
+  let nanv = bitcast<f32>(bitcast<u32>(maximum) | 0x7fc00000u);
+  return select(exp_shift(value, maximum), nanv, maximum > FLT_MAX);
+}
+fn normalized(value: f32, maximum: f32, sum: f32) -> f32 {
+  let rowIsNan = maximum > FLT_MAX || is_nan_f32(sum);
+  let nanOut = bitcast<f32>(bitcast<u32>(sum) | 0x7fc00000u);
+  return select(shifted_value(value, maximum) - log(sum), nanOut, rowIsNan);
+}
+@compute @workgroup_size(WGS)
+fn main(
+  @builtin(workgroup_id) wg: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>,
+  @builtin(num_workgroups) nwg: vec3<u32>
+) {
+  let group = wg.x + wg.y * nwg.x;
+  if (group >= PACKED_ROWS) {
+    return;
+  }
+  let outer = group / PACKED_INNER;
+  let inner0 = (group % PACKED_INNER) * 4u;
+  var localMax = vec4<f32>(-FLT_MAX);
+  for (var k = lid.x; k < REDUCE_SIZE; k = k + WGS) {
+    let base = outer * REDUCE_SIZE * INNER_SIZE + k * INNER_SIZE + inner0;
+    if (inner0 < INNER_SIZE) { localMax.x = max(localMax.x, x[base]); }
+    if (inner0 + 1u < INNER_SIZE) { localMax.y = max(localMax.y, x[base + 1u]); }
+    if (inner0 + 2u < INNER_SIZE) { localMax.z = max(localMax.z, x[base + 2u]); }
+    if (inner0 + 3u < INNER_SIZE) { localMax.w = max(localMax.w, x[base + 3u]); }
+  }
+  let maximum = reduce_max4(localMax, lid.x);
+  var localSum = vec4<f32>(0.0);
+  for (var k = lid.x; k < REDUCE_SIZE; k = k + WGS) {
+    let base = outer * REDUCE_SIZE * INNER_SIZE + k * INNER_SIZE + inner0;
+    if (inner0 < INNER_SIZE) { localSum.x = localSum.x + stable_exp(x[base], maximum.x); }
+    if (inner0 + 1u < INNER_SIZE) { localSum.y = localSum.y + stable_exp(x[base + 1u], maximum.y); }
+    if (inner0 + 2u < INNER_SIZE) { localSum.z = localSum.z + stable_exp(x[base + 2u], maximum.z); }
+    if (inner0 + 3u < INNER_SIZE) { localSum.w = localSum.w + stable_exp(x[base + 3u], maximum.w); }
+  }
+  let sum = reduce_add4(localSum, lid.x);
+  for (var k = lid.x; k < REDUCE_SIZE; k = k + WGS) {
+    let base = outer * REDUCE_SIZE * INNER_SIZE + k * INNER_SIZE + inner0;
+    if (inner0 < INNER_SIZE) { y[base] = normalized(x[base], maximum.x, sum.x); }
+    if (inner0 + 1u < INNER_SIZE) { y[base + 1u] = normalized(x[base + 1u], maximum.y, sum.y); }
+    if (inner0 + 2u < INNER_SIZE) { y[base + 2u] = normalized(x[base + 2u], maximum.z, sum.z); }
+    if (inner0 + 3u < INNER_SIZE) { y[base + 3u] = normalized(x[base + 3u], maximum.w, sum.w); }
+  }
+}

build/webgpu/test.json ADDED Viewed

	@@ -0,0 +1,720 @@

+{
+  "op": "ai.onnx.LogSoftmax",
+  "fixtureArrays": {
+    "ort_seed_123_input_x": [1.0856307, 0.99734545, 0.2829785, 1.5062947, 0.5786002, 1.6514366, 2.4266791, 0.42891264, 1.2659363, 0.8667404, 0.6788862, 0.09470897, 1.4913896, 0.638902, 0.44398195, 0.43435127, 2.20593, 2.1867862, 1.004054, 0.3861864, 0.7373686, 1.4907321, 0.9358339, 1.175829, 1.2538806, 0.6377515, 0.9071052, 1.4286807, 0.14006872, 0.8617549, 0.25561938, 2.798589, 1.7715331, 0.69987726, 0.92746246, 0.17363568, 0.002845916, 0.6882227, 0.87953633, 0.28362733, 0.8053665, 1.7276695, 0.3908998, 0.57380587, 0.33858904, 0.011830495, 2.3923652, 0.41291216, 0.978736, 2.2381434, 1.2940853, 1.0387882, 1.7437122, 0.79806274, 0.02968323, 1.0693159, 0.8907064, 1.7548862, 1.4956441, 1.0693927],
+    "onnx_backend_logsoftmax_input_x": [1.764052391052246, 0.40015721321105957, 0.978738009929657, 2.2408931255340576, 1.8675580024719238, 0.9772778749465942, 0.9500884413719177, 0.15135720372200012, 0.10321885347366333, 0.4105985164642334, 0.14404356479644775, 1.4542734622955322, 0.7610377073287964, 0.12167501449584961, 0.44386324286460876, 0.3336743414402008, 1.4940791130065918, 0.2051582634449005, 0.3130677044391632, 0.8540957570075989, 2.5529897212982178, 0.653618574142456, 0.8644362092018127, 0.7421650290489197, 2.269754648208618, 1.4543657302856445, 0.04575851559638977, 0.18718385696411133, 1.5327792167663574, 1.4693588018417358, 0.154947429895401, 0.37816253304481506, 0.8877857327461243, 1.980796456336975, 0.34791216254234314, 0.15634897351264954, 1.2302906513214111, 1.202379822731018, 0.38732680678367615, 0.302302747964859, 1.0485529899597168, 1.420017957687378, 1.7062702178955078, 1.950775384902954, 0.5096521973609924, 0.4380742907524109, 1.2527953386306763, 0.7774903774261475, 1.6138978004455566, 0.21274028718471527, 0.8954665660858154, 0.38690251111984253, 0.5108051300048828, 1.18063223361969, 0.02818222902715206, 0.4283318817615509, 0.06651721894741058, 0.30247190594673157, 0.6343221068382263, 0.3627411723136902]
+  },
+  "cases": [
+    {
+      "name": "strided_packed4_tail_rank3_axis1_inner6",
+      "provenance": {
+        "notes": "Exercises the scalar-aligned packed-4 strided reducer with two valid lanes in the final inner group."
+      },
+      "attrs": { "axis": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 17, 6],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.07, "scale": 2.0 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 17, 6], "tolerance": 0.00001, "relTolerance": 0.00001 } }
+    },
+    {
+      "name": "strided_capacity_large_axis_tail_boundary",
+      "provenance": {
+        "notes": "Compact correctness sibling for the large-inner selector crossover. The case-local threshold selects the capacity path while axis length 2048 and inner=6 retain the production gate and non-multiple-of-four tail."
+      },
+      "tunables": { "STRIDED_CAPACITY_MIN_ROWS": 1 },
+      "attrs": { "axis": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 2048, 6],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.007, "scale": 2.0 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 2048, 6], "tolerance": 0.00001, "relTolerance": 0.00001 } }
+    },
+    {
+      "name": "dispatch_cliff_online_2dfold_65540x40",
+      "attrs": { "axis": -1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [65540, 40],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.007, "scale": 2.0 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [65540, 40], "tolerance": 0.0001, "relTolerance": 0.0001 } }
+    },
+    {
+      "name": "dispatch_cliff_stable3pass_513x32768",
+      "attrs": { "axis": -1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [513, 32768],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.031, "scale": 3.0 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [513, 32768], "tolerance": 0.0001, "relTolerance": 0.00001 } }
+    },
+    {
+      "name": "rank2_last_axis",
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 4],
+          "data": { "kind": "values", "values": [1.0, 2.0, 3.0, 4.0, -2.0, 0.0, 2.0, 4.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 4] } }
+    },
+    {
+      "name": "packed_rows_tail_axis1_33x32",
+      "provenance": {
+        "notes": "One row beyond a full 32-row packed workgroup locks the shared portable multirow kernel's partial-workgroup barrier and output guards."
+      },
+      "attrs": { "axis": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [33, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.037, "cosStep": 0.019, "scale": 2.0 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [33, 32], "tolerance": 0.00001, "relTolerance": 0.00001 } }
+    },
+    {
+      "name": "many_short_rows_axis1_64x32",
+      "provenance": {
+        "notes": "Compact correctness sibling for many-short-row LogSoftmax benchmarking; keeps the reduced axis short without making the correctness suite benchmark-sized."
+      },
+      "attrs": { "axis": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.037, "cosStep": 0.019, "scale": 2.0 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [64, 32], "tolerance": 0.00001, "relTolerance": 0.00001 } }
+    },
+    {
+      "name": "ort_simple_axis1",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/math/logsoftmax_test.cc",
+        "test": "LogSoftmaxOperator.Simple"
+      },
+      "attrs": { "axis": 1 },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [1, 3], "data": { "kind": "values", "values": [-1.0, 0.0, 1.0] } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 3], "tolerance": 0.000001 } }
+    },
+    {
+      "name": "ort_large_number_axis1",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/math/logsoftmax_test.cc",
+        "test": "LogSoftmaxOperator.LargeNumber"
+      },
+      "attrs": { "axis": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 4],
+          "data": { "kind": "values", "values": [0.0, 1.0, 2.0, 3.0, 10000.0, 10001.0, 10002.0, 10003.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 4], "tolerance": 0.0005 } }
+    },
+    {
+      "name": "rank3_last_axis_f16",
+      "attrs": { "axis": -1 },
+      "inputs": { "x": { "dtype": "float16", "shape": [2, 2, 3] } },
+      "outputs": { "y": { "dtype": "float16", "shape": [2, 2, 3] } },
+      "tolerance": 0.002
+    },
+    {
+      "name": "float32_min_uniform_regression",
+      "attrs": { "axis": -1 },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [1, 64], "data": { "kind": "constant", "value": -3.4028234663852886e+38 } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 64], "tolerance": 0.000001 } }
+    },
+    {
+      "name": "longrow_split_axis1_1x65536",
+      "attrs": { "axis": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 65536],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.031, "scale": 3.0 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 65536], "tolerance": 0.0001, "relTolerance": 0.00001 } }
+    },
+    {
+      "name": "longrow_split_float32_min_uniform",
+      "attrs": { "axis": -1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 65536],
+          "data": { "kind": "constant", "value": -3.4028234663852886e+38 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 65536], "tolerance": 0.00001, "relTolerance": 0.00001 } }
+    },
+    {
+      "name": "positive_infinity_rows_return_nan",
+      "attrs": { "axis": -1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 4],
+          "data": { "kind": "values", "values": ["Infinity", 1.0, 2.0, -3.0, "Infinity", "Infinity", 5.0, 5.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 4], "tolerance": 0.000001, "allowNaN": true } }
+    },
+    {
+      "name": "singleton_axis_mixed_finite_nonfinite_f16",
+      "attrs": { "axis": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [1, 4],
+          "data": { "kind": "values", "values": [0.25, "Infinity", "-Infinity", "NaN"] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float16", "shape": [1, 4], "tolerance": 0, "allowNaN": true } }
+    },
+    {
+      "name": "empty_last_axis_dim_zero",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/math/softmax_test.cc",
+        "test": "SoftmaxOperator.DimWithZero",
+        "notes": "LogSoftmax analogue of ORT's zero-dimension Softmax coverage; the output is empty, so the framework plans the op as a no-op (zero dispatch) per ONNX zero-dimension semantics."
+      },
+      "attrs": { "axis": -1 },
+      "inputs": { "x": { "dtype": "float32", "shape": [1, 0], "data": { "kind": "values", "values": [] } } },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 0], "tolerance": 0 } }
+    },
+    {
+      "name": "empty_axis0_dim_zero",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/math/softmax_test.cc",
+        "test": "SoftmaxOperator.DimWithZero",
+        "notes": "LogSoftmax leading-axis analogue of ORT's zero-dimension Softmax coverage; the reduced axis itself has size zero and the output is empty."
+      },
+      "attrs": { "axis": 0 },
+      "inputs": { "x": { "dtype": "float32", "shape": [0, 1], "data": { "kind": "values", "values": [] } } },
+      "outputs": { "y": { "dtype": "float32", "shape": [0, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "axis0_non_last_rank2",
+      "attrs": { "axis": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 4],
+          "data": { "kind": "values", "values": [1.0, 2.0, 3.0, 4.0, 2.0, 1.0, 4.0, 3.0, 3.0, 4.0, 1.0, 2.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [3, 4], "tolerance": 0.000001 } }
+    },
+    {
+      "name": "ort_opset13_rank3_axis0_large_offsets",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/math/logsoftmax_test.cc",
+        "test": "LogSoftmaxOperator.ThreeDimsAxis0",
+        "notes": "Opset-13 axis=0 semantics reduce only the leading axis; large per-column offsets exercise stable strided-axis normalization."
+      },
+      "attrs": { "axis": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 2, 2],
+          "data": {
+            "kind": "values",
+            "values": [10000.0, -10000.0, 0.0, 5.0, 10001.0, -9999.0, -1.0, 6.0, 9999.0, -10001.0, 2.0, 4.0]
+          }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [3, 2, 2], "tolerance": 0.000001 } }
+    },
+    {
+      "name": "axis0_positive_infinity_rows_return_nan",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/math/logsoftmax_test.cc",
+        "test": "LogSoftmaxOperator.ThreeDimsAxis0",
+        "notes": "Positive infinity must poison the whole reduction group even when the reduced axis is strided instead of contiguous."
+      },
+      "attrs": { "axis": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 2, 2],
+          "data": { "kind": "values", "values": ["Infinity", 1.0, 2.0, "Infinity", 3.0, "Infinity", "Infinity", 4.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 2, 2], "tolerance": 0, "allowNaN": true } }
+    },
+    {
+      "name": "rank6_last_axis",
+      "attrs": { "axis": -1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 2, 1, 2, 1, 3],
+          "data": { "kind": "values", "values": [1.0, 2.0, 3.0, -1.0, 0.0, 1.0, 4.0, 4.0, 5.0, -2.0, -3.0, -4.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 2, 1, 2, 1, 3], "tolerance": 0.000001 } }
+    },
+    {
+      "name": "rank3_axis1_middle_large_values",
+      "attrs": { "axis": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 3, 2],
+          "data": {
+            "kind": "values",
+            "values": [1000.0, -1000.0, 1001.0, -1001.0, 999.0, -999.0, -50.0, 50.0, -51.0, 49.0, -52.0, 48.0]
+          }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 3, 2], "tolerance": 0.000001 } }
+    },
+    {
+      "name": "ort_opset13_rank3_axis1_seed123",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/math/logsoftmax_test.cc",
+        "test": "LogSoftmaxOperator.ThreeDimsAxis1_opset13"
+      },
+      "attrs": { "axis": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 4, 5],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/ort_seed_123_input_x" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [3, 4, 5], "tolerance": 0.000001 } }
+    },
+    {
+      "name": "ort_opset13_rank3_axis2_seed123",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/math/logsoftmax_test.cc",
+        "test": "LogSoftmaxOperator.ThreeDimsAxis2_opset13"
+      },
+      "attrs": { "axis": 2 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 4, 5],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/ort_seed_123_input_x" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [3, 4, 5], "tolerance": 0.000001 } }
+    },
+    {
+      "name": "rank3_axis_minus2_middle_strided",
+      "attrs": { "axis": -2 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 3, 3],
+          "data": { "kind": "values", "values": [1.0, 2.0, 3.0, 2.0, 3.0, 4.0, 3.0, 4.0, 5.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 3, 3], "tolerance": 0.000001 } }
+    },
+    {
+      "name": "ort_opset13_rank3_default_axis_seed123",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/math/logsoftmax_test.cc",
+        "test": "LogSoftmaxOperator.ThreeDimsDefaultAxis_opset13"
+      },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 4, 5],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/ort_seed_123_input_x" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [3, 4, 5], "tolerance": 0.000001 } }
+    },
+    {
+      "name": "ort_negative_axis_rank3_last_seed123",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/math/logsoftmax_test.cc",
+        "test": "LogSoftmaxOperator.ThreeDimsNegativeAxis"
+      },
+      "attrs": { "axis": -1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 4, 5],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/ort_seed_123_input_x" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [3, 4, 5], "tolerance": 0.000001 } }
+    },
+    {
+      "name": "ort_axis1_large_dim_1025",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/math/logsoftmax_test.cc",
+        "test": "LogSoftmaxOperator.2DInputReduceOnAxis1WithLargeDim"
+      },
+      "attrs": { "axis": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 1025],
+          "data": {
+            "kind": "values",
+            "values": [0.01, 0.02, 0.03, 0.04, 0.05, 0.06, 0.07, 0.08, 0.09, 0.1, 0.11, 0.12, 0.13, 0.14, 0.15, 0.16, 0.17, 0.18, 0.19, 0.2, 0.21, 0.22, 0.23, 0.24, 0.25, 0.26, 0.27, 0.28, 0.29, 0.3, 0.31, 0.32, 0.33, 0.34, 0.35, 0.36, 0.37, 0.38, 0.39, 0.4, 0.41, 0.42, 0.43, 0.44, 0.45, 0.46, 0.47, 0.48, 0.49, 0.5, 0.51, 0.52, 0.53, 0.54, 0.55, 0.56, 0.57, 0.58, 0.59, 0.6, 0.61, 0.62, 0.63, 0.64, 0.65, 0.66, 0.67, 0.68, 0.69, 0.7, 0.71, 0.72, 0.73, 0.74, 0.75, 0.76, 0.77, 0.78, 0.79, 0.8, 0.81, 0.82, 0.83, 0.84, 0.85, 0.86, 0.87, 0.88, 0.89, 0.9, 0.91, 0.92, 0.93, 0.94, 0.95, 0.96, 0.97, 0.98, 0.99, 1.0, 1.01, 1.02, 1.03, 1.04, 1.05, 1.06, 1.07, 1.08, 1.09, 1.1, 1.11, 1.12, 1.13, 1.14, 1.15, 1.16, 1.17, 1.18, 1.19, 1.2, 1.21, 1.22, 1.23, 1.24, 1.25, 1.26, 1.27, 1.28, 1.29, 1.3, 1.31, 1.32, 1.33, 1.34, 1.35, 1.36, 1.37, 1.38, 1.39, 1.4, 1.41, 1.42, 1.43, 1.44, 1.45, 1.46, 1.47, 1.48, 1.49, 1.5, 1.51, 1.52, 1.53, 1.54, 1.55, 1.56, 1.57, 1.58, 1.59, 1.6, 1.61, 1.62, 1.63, 1.64, 1.65, 1.66, 1.67, 1.68, 1.69, 1.7, 1.71, 1.72, 1.73, 1.74, 1.75, 1.76, 1.77, 1.78, 1.79, 1.8, 1.81, 1.82, 1.83, 1.84, 1.85, 1.86, 1.87, 1.88, 1.89, 1.9, 1.91, 1.92, 1.93, 1.94, 1.95, 1.96, 1.97, 1.98, 1.99, 2.0, 2.01, 2.02, 2.03, 2.04, 2.05, 2.06, 2.07, 2.08, 2.09, 2.1, 2.11, 2.12, 2.13, 2.14, 2.15, 2.16, 2.17, 2.18, 2.19, 2.2, 2.21, 2.22, 2.23, 2.24, 2.25, 2.26, 2.27, 2.28, 2.29, 2.3, 2.31, 2.32, 2.33, 2.34, 2.35, 2.36, 2.37, 2.38, 2.39, 2.4, 2.41, 2.42, 2.43, 2.44, 2.45, 2.46, 2.47, 2.48, 2.49, 2.5, 2.51, 2.52, 2.53, 2.54, 2.55, 2.56, 2.57, 2.58, 2.59, 2.6, 2.61, 2.62, 2.63, 2.64, 2.65, 2.66, 2.67, 2.68, 2.69, 2.7, 2.71, 2.72, 2.73, 2.74, 2.75, 2.76, 2.77, 2.78, 2.79, 2.8, 2.81, 2.82, 2.83, 2.84, 2.85, 2.86, 2.87, 2.88, 2.89, 2.9, 2.91, 2.92, 2.93, 2.94, 2.95, 2.96, 2.97, 2.98, 2.99, 3.0, 3.01, 3.02, 3.03, 3.04, 3.05, 3.06, 3.07, 3.08, 3.09, 3.1, 3.11, 3.12, 3.13, 3.14, 3.15, 3.16, 3.17, 3.18, 3.19, 3.2, 3.21, 3.22, 3.23, 3.24, 3.25, 3.26, 3.27, 3.28, 3.29, 3.3, 3.31, 3.32, 3.33, 3.34, 3.35, 3.36, 3.37, 3.38, 3.39, 3.4, 3.41, 3.42, 3.43, 3.44, 3.45, 3.46, 3.47, 3.48, 3.49, 3.5, 3.51, 3.52, 3.53, 3.54, 3.55, 3.56, 3.57, 3.58, 3.59, 3.6, 3.61, 3.62, 3.63, 3.64, 3.65, 3.66, 3.67, 3.68, 3.69, 3.7, 3.71, 3.72, 3.73, 3.74, 3.75, 3.76, 3.77, 3.78, 3.79, 3.8, 3.81, 3.82, 3.83, 3.84, 3.85, 3.86, 3.87, 3.88, 3.89, 3.9, 3.91, 3.92, 3.93, 3.94, 3.95, 3.96, 3.97, 3.98, 3.99, 4.0, 4.01, 4.02, 4.03, 4.04, 4.05, 4.06, 4.07, 4.08, 4.09, 4.1, 4.11, 4.12, 4.13, 4.14, 4.15, 4.16, 4.17, 4.18, 4.19, 4.2, 4.21, 4.22, 4.23, 4.24, 4.25, 4.26, 4.27, 4.28, 4.29, 4.3, 4.31, 4.32, 4.33, 4.34, 4.35, 4.36, 4.37, 4.38, 4.39, 4.4, 4.41, 4.42, 4.43, 4.44, 4.45, 4.46, 4.47, 4.48, 4.49, 4.5, 4.51, 4.52, 4.53, 4.54, 4.55, 4.56, 4.57, 4.58, 4.59, 4.6, 4.61, 4.62, 4.63, 4.64, 4.65, 4.66, 4.67, 4.68, 4.69, 4.7, 4.71, 4.72, 4.73, 4.74, 4.75, 4.76, 4.77, 4.78, 4.79, 4.8, 4.81, 4.82, 4.83, 4.84, 4.85, 4.86, 4.87, 4.88, 4.89, 4.9, 4.91, 4.92, 4.93, 4.94, 4.95, 4.96, 4.97, 4.98, 4.99, 5.0, 5.01, 5.02, 5.03, 5.04, 5.05, 5.06, 5.07, 5.08, 5.09, 5.1, 5.11, 5.12, 5.13, 5.14, 5.15, 5.16, 5.17, 5.18, 5.19, 5.2, 5.21, 5.22, 5.23, 5.24, 5.25, 5.26, 5.27, 5.28, 5.29, 5.3, 5.31, 5.32, 5.33, 5.34, 5.35, 5.36, 5.37, 5.38, 5.39, 5.4, 5.41, 5.42, 5.43, 5.44, 5.45, 5.46, 5.47, 5.48, 5.49, 5.5, 5.51, 5.52, 5.53, 5.54, 5.55, 5.56, 5.57, 5.58, 5.59, 5.6, 5.61, 5.62, 5.63, 5.64, 5.65, 5.66, 5.67, 5.68, 5.69, 5.7, 5.71, 5.72, 5.73, 5.74, 5.75, 5.76, 5.77, 5.78, 5.79, 5.8, 5.81, 5.82, 5.83, 5.84, 5.85, 5.86, 5.87, 5.88, 5.89, 5.9, 5.91, 5.92, 5.93, 5.94, 5.95, 5.96, 5.97, 5.98, 5.99, 6.0, 6.01, 6.02, 6.03, 6.04, 6.05, 6.06, 6.07, 6.08, 6.09, 6.1, 6.11, 6.12, 6.13, 6.14, 6.15, 6.16, 6.17, 6.18, 6.19, 6.2, 6.21, 6.22, 6.23, 6.24, 6.25, 6.26, 6.27, 6.28, 6.29, 6.3, 6.31, 6.32, 6.33, 6.34, 6.35, 6.36, 6.37, 6.38, 6.39, 6.4, 6.41, 6.42, 6.43, 6.44, 6.45, 6.46, 6.47, 6.48, 6.49, 6.5, 6.51, 6.52, 6.53, 6.54, 6.55, 6.56, 6.57, 6.58, 6.59, 6.6, 6.61, 6.62, 6.63, 6.64, 6.65, 6.66, 6.67, 6.68, 6.69, 6.7, 6.71, 6.72, 6.73, 6.74, 6.75, 6.76, 6.77, 6.78, 6.79, 6.8, 6.81, 6.82, 6.83, 6.84, 6.85, 6.86, 6.87, 6.88, 6.89, 6.9, 6.91, 6.92, 6.93, 6.94, 6.95, 6.96, 6.97, 6.98, 6.99, 7.0, 7.01, 7.02, 7.03, 7.04, 7.05, 7.06, 7.07, 7.08, 7.09, 7.1, 7.11, 7.12, 7.13, 7.14, 7.15, 7.16, 7.17, 7.18, 7.19, 7.2, 7.21, 7.22, 7.23, 7.24, 7.25, 7.26, 7.27, 7.28, 7.29, 7.3, 7.31, 7.32, 7.33, 7.34, 7.35, 7.36, 7.37, 7.38, 7.39, 7.4, 7.41, 7.42, 7.43, 7.44, 7.45, 7.46, 7.47, 7.48, 7.49, 7.5, 7.51, 7.52, 7.53, 7.54, 7.55, 7.56, 7.57, 7.58, 7.59, 7.6, 7.61, 7.62, 7.63, 7.64, 7.65, 7.66, 7.67, 7.68, 7.69, 7.7, 7.71, 7.72, 7.73, 7.74, 7.75, 7.76, 7.77, 7.78, 7.79, 7.8, 7.81, 7.82, 7.83, 7.84, 7.85, 7.86, 7.87, 7.88, 7.89, 7.9, 7.91, 7.92, 7.93, 7.94, 7.95, 7.96, 7.97, 7.98, 7.99, 8.0, 8.01, 8.02, 8.03, 8.04, 8.05, 8.06, 8.07, 8.08, 8.09, 8.1, 8.11, 8.12, 8.13, 8.14, 8.15, 8.16, 8.17, 8.18, 8.19, 8.2, 8.21, 8.22, 8.23, 8.24, 8.25, 8.26, 8.27, 8.28, 8.29, 8.3, 8.31, 8.32, 8.33, 8.34, 8.35, 8.36, 8.37, 8.38, 8.39, 8.4, 8.41, 8.42, 8.43, 8.44, 8.45, 8.46, 8.47, 8.48, 8.49, 8.5, 8.51, 8.52, 8.53, 8.54, 8.55, 8.56, 8.57, 8.58, 8.59, 8.6, 8.61, 8.62, 8.63, 8.64, 8.65, 8.66, 8.67, 8.68, 8.69, 8.7, 8.71, 8.72, 8.73, 8.74, 8.75, 8.76, 8.77, 8.78, 8.79, 8.8, 8.81, 8.82, 8.83, 8.84, 8.85, 8.86, 8.87, 8.88, 8.89, 8.9, 8.91, 8.92, 8.93, 8.94, 8.95, 8.96, 8.97, 8.98, 8.99, 9.0, 9.01, 9.02, 9.03, 9.04, 9.05, 9.06, 9.07, 9.08, 9.09, 9.1, 9.11, 9.12, 9.13, 9.14, 9.15, 9.16, 9.17, 9.18, 9.19, 9.2, 9.21, 9.22, 9.23, 9.24, 9.25, 9.26, 9.27, 9.28, 9.29, 9.3, 9.31, 9.32, 9.33, 9.34, 9.35, 9.36, 9.37, 9.38, 9.39, 9.4, 9.41, 9.42, 9.43, 9.44, 9.45, 9.46, 9.47, 9.48, 9.49, 9.5, 9.51, 9.52, 9.53, 9.54, 9.55, 9.56, 9.57, 9.58, 9.59, 9.6, 9.61, 9.62, 9.63, 9.64, 9.65, 9.66, 9.67, 9.68, 9.69, 9.7, 9.71, 9.72, 9.73, 9.74, 9.75, 9.76, 9.77, 9.78, 9.79, 9.8, 9.81, 9.82, 9.83, 9.84, 9.85, 9.86, 9.87, 9.88, 9.89, 9.9, 9.91, 9.92, 9.93, 9.94, 9.95, 9.96, 9.97, 9.98, 9.99, 10.0, 10.01, 10.02, 10.03, 10.04, 10.05, 10.06, 10.07, 10.08, 10.09, 10.1, 10.11, 10.12, 10.13, 10.14, 10.15, 10.16, 10.17, 10.18, 10.19, 10.2, 10.21, 10.22, 10.23, 10.24, 10.25]
+          }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 1025], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "onnx_backend_logsoftmax_axis_0",
+      "provenance": { "source": "cmake/external/onnx/onnx/backend/test/data/node/test_logsoftmax_axis_0" },
+      "attrs": { "axis": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 4, 5],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/onnx_backend_logsoftmax_input_x" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [3, 4, 5], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "onnx_backend_logsoftmax_axis_1",
+      "provenance": { "source": "cmake/external/onnx/onnx/backend/test/data/node/test_logsoftmax_axis_1" },
+      "attrs": { "axis": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 4, 5],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/onnx_backend_logsoftmax_input_x" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [3, 4, 5], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "onnx_backend_logsoftmax_axis_2",
+      "provenance": { "source": "cmake/external/onnx/onnx/backend/test/data/node/test_logsoftmax_axis_2" },
+      "attrs": { "axis": 2 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 4, 5],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/onnx_backend_logsoftmax_input_x" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [3, 4, 5], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "onnx_backend_logsoftmax_default_axis",
+      "provenance": { "source": "cmake/external/onnx/onnx/backend/test/data/node/test_logsoftmax_default_axis" },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 4, 5],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/onnx_backend_logsoftmax_input_x" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [3, 4, 5], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "onnx_backend_logsoftmax_example_1",
+      "provenance": { "source": "cmake/external/onnx/onnx/backend/test/data/node/test_logsoftmax_example_1" },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [1, 3], "data": { "kind": "values", "values": [-1.0, 0.0, 1.0] } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 3], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "onnx_backend_logsoftmax_large_number",
+      "provenance": { "source": "cmake/external/onnx/onnx/backend/test/data/node/test_logsoftmax_large_number" },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 4],
+          "data": { "kind": "values", "values": [0.0, 1.0, 2.0, 3.0, 10000.0, 10001.0, 10002.0, 10003.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 4], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "onnx_backend_logsoftmax_negative_axis",
+      "provenance": { "source": "cmake/external/onnx/onnx/backend/test/data/node/test_logsoftmax_negative_axis" },
+      "attrs": { "axis": -1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 4, 5],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/onnx_backend_logsoftmax_input_x" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [3, 4, 5], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "online_wg_vec4_2x1024",
+      "attrs": { "axis": -1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.07, "scale": 2.0 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 1024], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "online_wg_scalar_odd_2x511",
+      "attrs": { "axis": -1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 511],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.11, "scale": 2.0 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 511], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "online_wg_unaligned_even_2x510",
+      "attrs": { "axis": -1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 510],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.09, "scale": 2.0 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 510], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "online_wg_batch_2048x64",
+      "attrs": { "axis": -1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2048, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 2.0 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2048, 64], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "online_wg_all_negative_infinity_2x64",
+      "attrs": { "axis": -1 },
+      "inputs": { "x": { "dtype": "float32", "shape": [2, 64], "data": { "kind": "negativeInfinity" } } },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 64], "tolerance": 0.000001, "allowNaN": true } }
+    },
+    {
+      "name": "online_wg_neg_inf_padding_mixed_2x40",
+      "attrs": { "axis": -1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 40],
+          "data": { "kind": "cycle", "values": ["-Infinity", 0.5, 1.25, "-Infinity", -2.5] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 40], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "online_wg_f16_2x128",
+      "attrs": { "axis": -1 },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [2, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.21, "cosStep": 0.09, "scale": 2.0 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float16", "shape": [2, 128], "tolerance": 0.005 } }
+    },
+    {
+      "name": "dispatch_cliff_strided_rows_140000",
+      "provenance": {
+        "notes": "Non-last (strided) axis log-softmax over 140000 outer rows (>65535) forces the folded one-workgroup-per-row dispatch in the strided_3pass reduce_max/exp_sum passes (shared with Softmax). Validates the wg.x + wg.y*nwg.x reconstruction guarded by the compile-time STRIDED_ROWS const. Before the 2D fold the plan-time dispatch-limit guard throws."
+      },
+      "attrs": { "axis": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [70000, 4, 2],
+          "data": { "kind": "cycle", "values": [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [70000, 4, 2], "tolerance": 0.0001 } }
+    },
+    {
+      "name": "rank1_axis_minus1_f32",
+      "attrs": { "axis": -1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.041, "cosStep": 0.017, "scale": 3.0 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [128], "tolerance": 0.000001 } }
+    },
+    {
+      "name": "rank1_axis0_f32",
+      "attrs": { "axis": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.041, "cosStep": 0.017, "scale": 3.0 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [128], "tolerance": 0.000001 } }
+    },
+    {
+      "name": "f16_strided_axis1_rank3",
+      "attrs": { "axis": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [4, 8, 5],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.07, "scale": 2.0 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float16", "shape": [4, 8, 5], "tolerance": 0.005 } }
+    },
+    {
+      "name": "f16_longrow_split_1x65536",
+      "attrs": { "axis": -1 },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [1, 65536],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.031, "scale": 2.0 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float16", "shape": [1, 65536], "tolerance": 0.005 } }
+    },
+    {
+      "name": "nan_in_input_propagates_rank2",
+      "attrs": { "axis": -1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 4],
+          "data": { "kind": "values", "values": ["NaN", 1.0, 2.0, 3.0, 1.0, 2.0, 3.0, 4.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 4], "tolerance": 0.000001, "allowNaN": true } }
+    },
+    {
+      "name": "longrow_split_positive_infinity_poisons_row_2x65536",
+      "attrs": { "axis": -1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 65536],
+          "data": { "kind": "cycle", "values": ["Infinity", 0.5, 1.25, -2.5, 0.75, -1.0, 2.0] }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "float32", "shape": [2, 65536], "tolerance": 0.0001, "relTolerance": 0.0001, "allowNaN": true }
+      }
+    },
+    {
+      "name": "longrow_split_all_negative_infinity_row_nan_1x65536",
+      "attrs": { "axis": -1 },
+      "inputs": { "x": { "dtype": "float32", "shape": [1, 65536], "data": { "kind": "negativeInfinity" } } },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 65536], "tolerance": 0.000001, "allowNaN": true } }
+    },
+    {
+      "name": "vec4_combine_positive_infinity_poisons_row_2x68",
+      "attrs": { "axis": -1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 68],
+          "data": { "kind": "cycle", "values": ["Infinity", 0.25, -1.5, 2.0, 0.5, -0.75, 1.75, -2.25] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 68], "tolerance": 0.000001, "allowNaN": true } }
+    },
+    {
+      "name": "longrow_split_multirow_isolation_3x65536",
+      "attrs": { "axis": -1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 65536],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.023, "scale": 4.0, "offset": 100.0 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [3, 65536], "tolerance": 0.0001, "relTolerance": 0.00001 } }
+    },
+    {
+      "name": "f16_longrow_split_positive_infinity_poisons_row_1x65536",
+      "attrs": { "axis": -1 },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [1, 65536],
+          "data": { "kind": "cycle", "values": ["Infinity", 0.5, 1.0, -1.5, 2.0, -0.5] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float16", "shape": [1, 65536], "tolerance": 0.005, "allowNaN": true } }
+    },
+    {
+      "name": "rank7_last_axis",
+      "attrs": { "axis": -1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 1, 2, 1, 2, 1, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 2.0 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 1, 2, 1, 2, 1, 3], "tolerance": 0.000001 } }
+    },
+    {
+      "name": "strided_rowmax_scratch_over_128mib_f16_capacity_fallback",
+      "provenance": {
+        "notes": "Row-scratch capacity cliff: 33554433 strided rows need a 128 MiB+1 f32 rowMax buffer, so strided_3pass fails axisRowScratchFits and only the single-pass online fallback remains. f16 keeps the 64 MiB input inside the same spec-minimum storage-binding limit the f32 scratch overflows (and rejects the two f32-only strided routes), which is exactly the device situation the fallback exists for."
+      },
+      "attrs": { "axis": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [1, 33554433],
+          "data": { "kind": "cycle", "values": [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float16", "shape": [1, 33554433], "tolerance": 0.002 } }
+    },
+    {
+      "name": "rank8_last_axis",
+      "attrs": { "axis": -1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 1, 2, 1, 2, 1, 2, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.21 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 1, 2, 1, 2, 1, 2, 3], "tolerance": 0.000001 } }
+    }
+  ]
+}