sync 91d990483a17

Browse files

Files changed (12) hide show

README.md +33 -9
build/webgpu/bench.json +388 -2
build/webgpu/manifest.json +598 -821
build/webgpu/metadata.json +34 -13
build/webgpu/test.json +615 -33
build/webgpu/topk-axis.wgsl.jinja +44 -27
build/webgpu/topk-large-block.wgsl.jinja +68 -45
build/webgpu/topk-portable-rows-smallk.wgsl.jinja +27 -17
build/webgpu/topk-small-rows-batched.wgsl.jinja +24 -13
build/webgpu/topk-strided-smallk.wgsl.jinja +24 -13
build/webgpu/topk-subgroup-rows.wgsl.jinja +24 -13
build/webgpu/topk-top1-last-axis.wgsl.jinja +30 -19

README.md CHANGED Viewed

@@ -18,20 +18,20 @@ See the [ONNX `TopK` spec](https://onnx.ai/onnx/operators/onnx__TopK.html) for t
 ## Inputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `X` | `x` | `T` | — | — | Values from which the top `k` entries are selected along `axis`. | required |
 ## Outputs
-| Name | Bind key | Logical dtype | WebGPU storage | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- | --- |
-| `Values` | `values` | `T` | runtime-selected; narrow integers and bool use 32-bit slots | same as `X` | derived; see description | Selected values; the reduced axis has length `k`. | required |
-| `Indices` | `indices` | `I` | `uint32` | same as `X` | derived; see description | Logical int64 indices of the selected values along the reduced axis; WebGPU stores these bounded indices as uint32. | required |
 ## Runtime arguments
-| Name | Kind | Semantic | Description | Presence |
 | --- | --- | --- | --- | --- |
 | `k` | `u32` | `kernel.k` | Number of values to select along the configured axis. | required |
@@ -52,13 +52,33 @@ Default values (overridable per request):
 | `T` | `float32`, `float16`, `int8`, `int16`, `int32`, `uint8`, `uint32` |
 | `I` | `int64` |
 ## Device requirements
 Some implementation variants require `subgroups`. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.
 ## Files
-- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, 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
@@ -73,10 +93,14 @@ Some implementation variants require `subgroups`. These are route-specific capab
 ## Use with `@huggingface/kernels`
-The loader derives every required output's shape and logical dtype from the manifest contract and this call.
-It then allocates the result tensors automatically.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

 ## Inputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `x` | `X` | `T` | — | — | Values from which the top `k` entries are selected along `axis`. | required |
 ## Outputs
+| Name | Upstream name | Logical dtype | WebGPU storage | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- | --- |
+| `values` | `Values` | `T` | runtime-selected; narrow integers and bool use 32-bit slots | same as `x` | derived | Selected values; the reduced axis has length `k`. | required |
+| `indices` | `Indices` | `I` | `uint32` | same as `x` | derived | Logical int64 indices of the selected values along the reduced axis; WebGPU stores these bounded indices as uint32. | required |
 ## Runtime arguments
+| Name | Kind | Upstream attribute | Description | Presence |
 | --- | --- | --- | --- | --- |
 | `k` | `u32` | `kernel.k` | Number of values to select along the configured axis. | required |
 | `T` | `float32`, `float16`, `int8`, `int16`, `int32`, `uint8`, `uint32` |
 | `I` | `int64` |
+## Implementation variants
+One implementation is selected per call from the device capabilities, the request shapes and the dtypes; these notes say what each one covers.
+- `axis_smallk_tournament` — Scans a strided non-last axis with one workgroup per output position and retains only a short candidate list instead of sorting the whole axis. It is favored when many independent outputs amortize the strided scan and remains the bounded-storage route when the axis does not fit shared memory.
+- `last_axis_large_top1` — Finds one winner on rows too wide for the direct shared-memory route by reducing blocks to scratch candidates. A second pass selects the final candidate.
+- `axis_bitonic` — Shared bitonic selection for arbitrary axes; floating inputs encode order keys once and gather original values after sorting, preserving stable ties and value bits.
+- `subgroup_rows_smallk` — Retains only K candidates per row and selects a row-parallel reduction layout according to row width and available subgroup geometry. It avoids shared storage proportional to the full axis.
+- `subgroup_min_rows_smallk` — Retains only K candidates per row and selects a row-parallel reduction layout according to row width and available subgroup geometry. It avoids shared storage proportional to the full axis.
+- `portable_rows_smallk` — Retains only K candidates per row and selects a row-parallel reduction layout according to row width and available subgroup geometry. It avoids shared storage proportional to the full axis.
+- `small_rows_batched` — Retains only K candidates per row and selects a row-parallel reduction layout according to row width and available subgroup geometry. It avoids shared storage proportional to the full axis.
+- `last_axis_large_one_merge` — Staged stable TopK with bounded candidate buffers. Floating inputs encode order keys on load and gather selected source values after the final merge; integer inputs retain exact integer keys.
+- `last_axis_large_two_merge` — Staged stable TopK with bounded candidate buffers. Floating inputs encode order keys on load and gather selected source values after the final merge; integer inputs retain exact integer keys.
+- `last_axis_large_one_merge_int` — Staged stable TopK with bounded candidate buffers. Floating inputs encode order keys on load and gather selected source values after the final merge; integer inputs retain exact integer keys.
+- `last_axis_large_two_merge_int` — Staged stable TopK with bounded candidate buffers. Floating inputs encode order keys on load and gather selected source values after the final merge; integer inputs retain exact integer keys.
+- `axis_large_one_merge` — Staged stable TopK with bounded candidate buffers. Floating inputs encode order keys on load and gather selected source values after the final merge; integer inputs retain exact integer keys.
+- `axis_large_two_merge` — Staged stable TopK with bounded candidate buffers. Floating inputs encode order keys on load and gather selected source values after the final merge; integer inputs retain exact integer keys.
+- `last_axis_large_three_merge` — Staged stable TopK with bounded candidate buffers. Floating inputs encode order keys on load and gather selected source values after the final merge; integer inputs retain exact integer keys.
+- `last_axis_large_three_merge_int` — Staged stable TopK with bounded candidate buffers. Floating inputs encode order keys on load and gather selected source values after the final merge; integer inputs retain exact integer keys.
 ## Device requirements
 Some implementation variants require `subgroups`. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.
 ## Files
+- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
 - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 ## Use with `@huggingface/kernels`
+```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.2
+```
+Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

build/webgpu/bench.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "ai.onnx.TopK",
   "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] },
   "cases": [
     {
@@ -274,7 +273,7 @@
       }
     },
     {
-      "name": "topk-axis1-f32-2048x256-k4-occupied-sibling",
       "preset": "smoke",
       "vars": { "rows": 2048, "cols": 256, "k": 4 },
       "args": { "k": 4 },
@@ -468,6 +467,393 @@
           { "name": "io", "type": "bandwidth", "value": "args.axis * args.inner * 4 + args.k * args.inner * 8" }
         ]
       }
     }
   ]
 }

 {
   "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] },
   "cases": [
     {
       }
     },
     {
+      "name": "topk-axis1-f32-2048x256-k4-many-rows",
       "preset": "smoke",
       "vars": { "rows": 2048, "cols": 256, "k": 4 },
       "args": { "k": 4 },
           { "name": "io", "type": "bandwidth", "value": "args.axis * args.inner * 4 + args.k * args.inner * 8" }
         ]
       }
+    },
+    {
+      "name": "topk-router-f32-131072x8-k2-unsorted-contract",
+      "preset": "stress",
+      "provenance": {
+        "notes": "This router shape selects two values from each of 131,072 eight-element rows with sorted=0. The unspecified output order permits the batched small-row selection path without a full row sort."
+      },
+      "vars": { "rows": 131072, "cols": 8, "k": 2 },
+      "args": { "k": 2 },
+      "attrs": { "axis": -1, "largest": 1, "sorted": 0 },
+      "inputs": { "x": { "shape": [131072, 8], "dtype": "float32", "dist": "normal", "seed": 213, "scale": 2 } },
+      "outputs": {
+        "values": { "shape": [131072, 2], "dtype": "float32", "dist": "empty" },
+        "indices": { "shape": [131072, 2], "dtype": "uint32", "dist": "empty" }
+      },
+      "bench": {
+        "metrics": [
+          { "name": "io", "type": "bandwidth", "value": "args.rows * args.cols * 4 + args.rows * args.k * 8" }
+        ]
+      }
+    },
+    {
+      "name": "topk-vocab-f32-64x50257-k50-two-merge-pathology",
+      "preset": "stress",
+      "provenance": {
+        "notes": "64 GPT-2 vocabulary rows with k=50: 99 blocks of 512 produce 4950 candidates per row, above the 2048 merge chunk, so every row takes the two-merge large-axis route. Control: topk-vocab-f32-64x50257-k1-large-top1-control (same input bytes, the k=1 large-top1 route)."
+      },
+      "vars": { "rows": 64, "cols": 50257, "k": 50 },
+      "args": { "k": 50 },
+      "attrs": { "axis": -1, "largest": 1, "sorted": 1 },
+      "inputs": { "x": { "shape": [64, 50257], "dtype": "float32", "dist": "normal", "seed": 218, "scale": 4 } },
+      "outputs": {
+        "values": { "shape": [64, 50], "dtype": "float32", "dist": "empty" },
+        "indices": { "shape": [64, 50], "dtype": "uint32", "dist": "empty" }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "name": "io",
+            "type": "bandwidth",
+            "value": "args.rows * args.cols * 4 + args.rows * ceilDiv(args.cols, (512 if args.k <= 128 else 2048)) * args.k * 16 + args.rows * ceilDiv(ceilDiv(args.cols, (512 if args.k <= 128 else 2048)) * args.k, 2048) * args.k * 16 + args.rows * args.k * 8"
+          }
+        ]
+      }
+    },
+    {
+      "name": "topk-vocab-f32-64x50257-k1-large-top1-control",
+      "preset": "stress",
+      "provenance": {
+        "notes": "Sixty-four GPT-2 vocabulary rows with k=1 exercise a single block reduction followed by one combine pass."
+      },
+      "vars": { "rows": 64, "cols": 50257, "k": 1 },
+      "args": { "k": 1 },
+      "attrs": { "axis": -1, "largest": 1, "sorted": 1 },
+      "inputs": { "x": { "shape": [64, 50257], "dtype": "float32", "dist": "normal", "seed": 219, "scale": 4 } },
+      "outputs": {
+        "values": { "shape": [64, 1], "dtype": "float32", "dist": "empty" },
+        "indices": { "shape": [64, 1], "dtype": "uint32", "dist": "empty" }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "name": "io",
+            "type": "bandwidth",
+            "value": "args.rows * args.cols * 4 + args.rows * ceilDiv(args.cols, 512) * 8 + args.rows * args.k * 8"
+          }
+        ]
+      }
+    },
+    {
+      "name": "topk-router-f32-rank3-8x512x32-k2-batched-rows",
+      "provenance": {
+        "notes": "Shape [8,512,32] represents 4,096 independent 32-element rows. Flattening the leading dimensions allows the batched small-row path to select k=2 from each row."
+      },
+      "preset": "smoke",
+      "vars": { "rows": 4096, "cols": 32, "k": 2 },
+      "args": { "k": 2 },
+      "attrs": { "axis": -1, "largest": 1, "sorted": 1 },
+      "inputs": { "x": { "shape": [8, 512, 32], "dtype": "float32", "dist": "normal", "seed": 119, "scale": 2 } },
+      "outputs": {
+        "values": { "shape": [8, 512, 2], "dtype": "float32" },
+        "indices": { "shape": [8, 512, 2], "dtype": "uint32" }
+      }
+    },
+    {
+      "name": "bitonic_prefix_boundary_float32_axis256_k17_largest0",
+      "preset": "stress",
+      "attrs": { "axis": -1, "largest": 0 },
+      "args": { "k": 17 },
+      "inputs": { "x": { "dtype": "float32", "shape": [64, 256], "dist": "normal", "seed": 1739 } },
+      "outputs": {
+        "values": { "dtype": "float32", "shape": [64, 17] },
+        "indices": { "dtype": "uint32", "shape": [64, 17] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "numel(shapes.x) * dtypeBytes(args.dtype) + numel(shapes.values) * (dtypeBytes(args.dtype) + 4)"
+          }
+        ]
+      },
+      "vars": { "dtype": "float32" }
+    },
+    {
+      "name": "bitonic_prefix_boundary_float32_axis256_k17_largest1",
+      "preset": "stress",
+      "attrs": { "axis": -1, "largest": 1 },
+      "args": { "k": 17 },
+      "inputs": { "x": { "dtype": "float32", "shape": [64, 256], "dist": "normal", "seed": 1739 } },
+      "outputs": {
+        "values": { "dtype": "float32", "shape": [64, 17] },
+        "indices": { "dtype": "uint32", "shape": [64, 17] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "numel(shapes.x) * dtypeBytes(args.dtype) + numel(shapes.values) * (dtypeBytes(args.dtype) + 4)"
+          }
+        ]
+      },
+      "vars": { "dtype": "float32" }
+    },
+    {
+      "name": "bitonic_prefix_boundary_float32_axis257_k17_largest0",
+      "preset": "stress",
+      "attrs": { "axis": -1, "largest": 0 },
+      "args": { "k": 17 },
+      "inputs": { "x": { "dtype": "float32", "shape": [64, 257], "dist": "normal", "seed": 1739 } },
+      "outputs": {
+        "values": { "dtype": "float32", "shape": [64, 17] },
+        "indices": { "dtype": "uint32", "shape": [64, 17] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "numel(shapes.x) * dtypeBytes(args.dtype) + numel(shapes.values) * (dtypeBytes(args.dtype) + 4)"
+          }
+        ]
+      },
+      "vars": { "dtype": "float32" }
+    },
+    {
+      "name": "bitonic_prefix_boundary_float32_axis257_k17_largest1",
+      "preset": "stress",
+      "attrs": { "axis": -1, "largest": 1 },
+      "args": { "k": 17 },
+      "inputs": { "x": { "dtype": "float32", "shape": [64, 257], "dist": "normal", "seed": 1739 } },
+      "outputs": {
+        "values": { "dtype": "float32", "shape": [64, 17] },
+        "indices": { "dtype": "uint32", "shape": [64, 17] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "numel(shapes.x) * dtypeBytes(args.dtype) + numel(shapes.values) * (dtypeBytes(args.dtype) + 4)"
+          }
+        ]
+      },
+      "vars": { "dtype": "float32" }
+    },
+    {
+      "name": "bitonic_prefix_boundary_float16_axis256_k17_largest0",
+      "preset": "stress",
+      "attrs": { "axis": -1, "largest": 0 },
+      "args": { "k": 17 },
+      "inputs": { "x": { "dtype": "float16", "shape": [64, 256], "dist": "normal", "seed": 1739 } },
+      "outputs": {
+        "values": { "dtype": "float16", "shape": [64, 17] },
+        "indices": { "dtype": "uint32", "shape": [64, 17] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "numel(shapes.x) * dtypeBytes(args.dtype) + numel(shapes.values) * (dtypeBytes(args.dtype) + 4)"
+          }
+        ]
+      },
+      "vars": { "dtype": "float16" }
+    },
+    {
+      "name": "bitonic_prefix_boundary_float16_axis256_k17_largest1",
+      "preset": "stress",
+      "attrs": { "axis": -1, "largest": 1 },
+      "args": { "k": 17 },
+      "inputs": { "x": { "dtype": "float16", "shape": [64, 256], "dist": "normal", "seed": 1739 } },
+      "outputs": {
+        "values": { "dtype": "float16", "shape": [64, 17] },
+        "indices": { "dtype": "uint32", "shape": [64, 17] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "numel(shapes.x) * dtypeBytes(args.dtype) + numel(shapes.values) * (dtypeBytes(args.dtype) + 4)"
+          }
+        ]
+      },
+      "vars": { "dtype": "float16" }
+    },
+    {
+      "name": "bitonic_prefix_boundary_float16_axis257_k17_largest0",
+      "preset": "stress",
+      "attrs": { "axis": -1, "largest": 0 },
+      "args": { "k": 17 },
+      "inputs": { "x": { "dtype": "float16", "shape": [64, 257], "dist": "normal", "seed": 1739 } },
+      "outputs": {
+        "values": { "dtype": "float16", "shape": [64, 17] },
+        "indices": { "dtype": "uint32", "shape": [64, 17] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "numel(shapes.x) * dtypeBytes(args.dtype) + numel(shapes.values) * (dtypeBytes(args.dtype) + 4)"
+          }
+        ]
+      },
+      "vars": { "dtype": "float16" }
+    },
+    {
+      "name": "bitonic_prefix_boundary_float16_axis257_k17_largest1",
+      "preset": "stress",
+      "attrs": { "axis": -1, "largest": 1 },
+      "args": { "k": 17 },
+      "inputs": { "x": { "dtype": "float16", "shape": [64, 257], "dist": "normal", "seed": 1739 } },
+      "outputs": {
+        "values": { "dtype": "float16", "shape": [64, 17] },
+        "indices": { "dtype": "uint32", "shape": [64, 17] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "numel(shapes.x) * dtypeBytes(args.dtype) + numel(shapes.values) * (dtypeBytes(args.dtype) + 4)"
+          }
+        ]
+      },
+      "vars": { "dtype": "float16" }
+    },
+    {
+      "name": "order-keys-float32-n257-axis0-k17",
+      "preset": "all",
+      "vars": { "rows": 7, "cols": 257, "k": 17 },
+      "args": { "k": 17 },
+      "attrs": { "axis": 0, "largest": 1, "sorted": 1 },
+      "inputs": { "x": { "dtype": "float32", "shape": [257, 7], "dist": "normal", "seed": 613, "scale": 2 } },
+      "outputs": {
+        "values": { "dtype": "float32", "shape": [17, 7] },
+        "indices": { "dtype": "uint32", "shape": [17, 7] }
+      },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          { "name": "io", "type": "bandwidth", "value": "args.rows * args.cols * 4 + args.rows * args.k * 8" }
+        ]
+      }
+    },
+    {
+      "name": "order-keys-float32-n257-axis1-k17",
+      "preset": "all",
+      "vars": { "rows": 7, "cols": 257, "k": 17 },
+      "args": { "k": 17 },
+      "attrs": { "axis": 1, "largest": 1, "sorted": 1 },
+      "inputs": { "x": { "dtype": "float32", "shape": [7, 257], "dist": "normal", "seed": 613, "scale": 2 } },
+      "outputs": {
+        "values": { "dtype": "float32", "shape": [7, 17] },
+        "indices": { "dtype": "uint32", "shape": [7, 17] }
+      },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          { "name": "io", "type": "bandwidth", "value": "args.rows * args.cols * 4 + args.rows * args.k * 8" }
+        ]
+      }
+    },
+    {
+      "name": "order-keys-float32-n4099-axis0-k17",
+      "preset": "all",
+      "vars": { "rows": 7, "cols": 4099, "k": 17 },
+      "args": { "k": 17 },
+      "attrs": { "axis": 0, "largest": 1, "sorted": 1 },
+      "inputs": { "x": { "dtype": "float32", "shape": [4099, 7], "dist": "normal", "seed": 613, "scale": 2 } },
+      "outputs": {
+        "values": { "dtype": "float32", "shape": [17, 7] },
+        "indices": { "dtype": "uint32", "shape": [17, 7] }
+      },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          { "name": "io", "type": "bandwidth", "value": "args.rows * args.cols * 4 + args.rows * args.k * 8" }
+        ]
+      }
+    },
+    {
+      "name": "order-keys-float32-n4099-axis1-k17",
+      "preset": "all",
+      "vars": { "rows": 7, "cols": 4099, "k": 17 },
+      "args": { "k": 17 },
+      "attrs": { "axis": 1, "largest": 1, "sorted": 1 },
+      "inputs": { "x": { "dtype": "float32", "shape": [7, 4099], "dist": "normal", "seed": 613, "scale": 2 } },
+      "outputs": {
+        "values": { "dtype": "float32", "shape": [7, 17] },
+        "indices": { "dtype": "uint32", "shape": [7, 17] }
+      },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          { "name": "io", "type": "bandwidth", "value": "args.rows * args.cols * 4 + args.rows * args.k * 8" }
+        ]
+      }
+    },
+    {
+      "name": "order-keys-float16-n257-axis0-k17",
+      "preset": "all",
+      "vars": { "rows": 7, "cols": 257, "k": 17 },
+      "args": { "k": 17 },
+      "attrs": { "axis": 0, "largest": 1, "sorted": 1 },
+      "inputs": { "x": { "dtype": "float16", "shape": [257, 7], "dist": "normal", "seed": 613, "scale": 2 } },
+      "outputs": {
+        "values": { "dtype": "float16", "shape": [17, 7] },
+        "indices": { "dtype": "uint32", "shape": [17, 7] }
+      },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          { "name": "io", "type": "bandwidth", "value": "args.rows * args.cols * 2 + args.rows * args.k * 6" }
+        ]
+      }
+    },
+    {
+      "name": "order-keys-float16-n257-axis1-k17",
+      "preset": "all",
+      "vars": { "rows": 7, "cols": 257, "k": 17 },
+      "args": { "k": 17 },
+      "attrs": { "axis": 1, "largest": 1, "sorted": 1 },
+      "inputs": { "x": { "dtype": "float16", "shape": [7, 257], "dist": "normal", "seed": 613, "scale": 2 } },
+      "outputs": {
+        "values": { "dtype": "float16", "shape": [7, 17] },
+        "indices": { "dtype": "uint32", "shape": [7, 17] }
+      },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          { "name": "io", "type": "bandwidth", "value": "args.rows * args.cols * 2 + args.rows * args.k * 6" }
+        ]
+      }
+    },
+    {
+      "name": "order-keys-float16-n4099-axis0-k17",
+      "preset": "all",
+      "vars": { "rows": 7, "cols": 4099, "k": 17 },
+      "args": { "k": 17 },
+      "attrs": { "axis": 0, "largest": 1, "sorted": 1 },
+      "inputs": { "x": { "dtype": "float16", "shape": [4099, 7], "dist": "normal", "seed": 613, "scale": 2 } },
+      "outputs": {
+        "values": { "dtype": "float16", "shape": [17, 7] },
+        "indices": { "dtype": "uint32", "shape": [17, 7] }
+      },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          { "name": "io", "type": "bandwidth", "value": "args.rows * args.cols * 2 + args.rows * args.k * 6" }
+        ]
+      }
+    },
+    {
+      "name": "order-keys-float16-n4099-axis1-k17",
+      "preset": "all",
+      "vars": { "rows": 7, "cols": 4099, "k": 17 },
+      "args": { "k": 17 },
+      "attrs": { "axis": 1, "largest": 1, "sorted": 1 },
+      "inputs": { "x": { "dtype": "float16", "shape": [7, 4099], "dist": "normal", "seed": 613, "scale": 2 } },
+      "outputs": {
+        "values": { "dtype": "float16", "shape": [7, 17] },
+        "indices": { "dtype": "uint32", "shape": [7, 17] }
+      },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          { "name": "io", "type": "bandwidth", "value": "args.rows * args.cols * 2 + args.rows * args.k * 6" }
+        ]
+      }
     }
   ]
 }

build/webgpu/manifest.json CHANGED Viewed

@@ -2,90 +2,65 @@
   "domain": "ai.onnx",
   "name": "TopK",
   "sinceVersion": 11,
-  "description": "Retrieves the top-`k` largest or smallest elements along the selected axis, returning values and stable lower-index tie-breaking indices.",
-  "inputs": [
-    { "role": "X", "dtype": "T", "description": "Values from which the top `k` entries are selected along `axis`." }
-  ],
-  "outputs": [
-    {
-      "role": "Values",
       "dtype": "T",
-      "rank": "ranks.X",
-      "description": "Selected values; the reduced axis has length `k`.",
-      "shape": "prefix(shapes.X, normalizedAxis) + [args.k] + suffix(shapes.X, normalizedAxis + 1)"
     },
-    {
-      "role": "Indices",
       "dtype": "I",
-      "rank": "ranks.X",
-      "description": "Logical int64 indices of the selected values along the reduced axis; WebGPU stores these bounded indices as uint32.",
-      "shape": "prefix(shapes.X, normalizedAxis) + [args.k] + suffix(shapes.X, normalizedAxis + 1)"
     }
-  ],
-  "attributes": { "axis": -1, "largest": 1, "sorted": 1 },
-  "attributeDescriptions": {
-    "axis": "Axis to reduce; negative values count from the back.",
-    "largest": "Select largest values when 1, smallest values when 0.",
-    "sorted": "Sort selected values when 1. A sorted result is also valid when output order is unspecified (`sorted=0`)."
   },
   "attributeConstraints": { "largest": { "values": [0, 1] }, "sorted": { "values": [0, 1] } },
   "typeConstraints": { "T": ["float32", "float16", "int8", "int16", "int32", "uint8", "uint32"], "I": ["int64"] },
-  "args": {
-    "x": { "kind": "tensor", "semantic": "X", "role": "input" },
-    "values": { "kind": "tensor", "semantic": "Values", "role": "output" },
-    "indices": { "kind": "tensor", "semantic": "Indices", "role": "output", "dtype": "uint32" },
-    "k": {
-      "kind": "u32",
-      "semantic": "kernel.k",
-      "role": "attribute",
-      "description": "Number of values to select along the configured axis."
-    }
-  },
   "tunables": {
-    "WORKGROUP_SIZE": 256,
-    "AXIS_SHARED_MAX_VALUES": 2048,
-    "TOP1_BLOCK_SIZE": 512,
-    "SMALL_K_BLOCK_SIZE": 512,
-    "LARGE_K_BLOCK_SIZE": 2048,
-    "SMALL_K_LIMIT": 128,
-    "MERGE_CHUNK_SIZE": 2048,
-    "SMALL_ROWS_MIN_ROWS": 4096,
-    "SMALL_ROWS_MAX_AXIS": 64,
-    "SMALL_ROWS_MAX_K": 8,
-    "PORTABLE_ROW_WORKGROUP_SIZE": 32,
-    "PORTABLE_ROWS_MIN": 1024,
-    "PORTABLE_ROWS_MAX_AXIS": 256,
-    "TOURNAMENT_MAX_AXIS": 4096,
-    "TOURNAMENT_MAX_K": 16,
-    "TOURNAMENT_MIN_OUTPUTS": 256
-  },
-  "tunableDescriptions": {
-    "AXIS_SHARED_MAX_VALUES": "Caps the padded axis size admitted to the single-workgroup shared bitonic route; larger axes use staged block selection.",
-    "MERGE_CHUNK_SIZE": "Caps the candidate count consumed by one large-axis merge workgroup; exceeding it adds another merge stage.",
-    "SMALL_ROWS_MIN_ROWS": "Sets the minimum row count for the one-lane-per-row small-K route; `SMALL_ROWS_MAX_AXIS` and `SMALL_ROWS_MAX_K` bound each lane's serial work.",
-    "PORTABLE_ROWS_MIN": "Sets the minimum row count for the portable small-K route, ensuring enough independent rows for cooperative per-row merging.",
-    "TOURNAMENT_MIN_OUTPUTS": "Sets the minimum independent output count that retains the strided small-K tournament when shared bitonic sorting also fits."
   },
   "derive": {
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
-    "foldedDispatchCapacity": "device.limits.maxComputeWorkgroupsPerDimension * device.limits.maxComputeWorkgroupsPerDimension",
     "wave32Adapter": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 32 and device.adapterInfo.subgroupMaxSize == 32",
     "subgroupsWave32": "device.features.has(\"subgroups\") and wave32Adapter",
-    "normalizedAxis": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.X",
-    "axisInRange": "ranks.X >= 1 and normalizedAxis >= 0 and normalizedAxis < ranks.X",
-    "axisOuter": "outer(shapes.X, normalizedAxis)",
-    "axisDim": "dim(shapes.X, normalizedAxis)",
-    "axisInner": "inner(shapes.X, normalizedAxis)",
     "outputPositions": "axisOuter * axisInner",
-    "outputShapeOk": "ranks.Values == ranks.X and ranks.Indices == ranks.X and outer(shapes.Values, normalizedAxis) == axisOuter and outer(shapes.Indices, normalizedAxis) == axisOuter and dim(shapes.Values, normalizedAxis) == args.k and dim(shapes.Indices, normalizedAxis) == args.k and inner(shapes.Values, normalizedAxis) == axisInner and inner(shapes.Indices, normalizedAxis) == axisInner",
     "dtypeOk": "f16Ok(dtypes.T)",
     "shapeContract": "axisInRange and outputShapeOk and args.k <= axisDim",
     "baseContract": "shapeContract and dtypeOk",
-    "lastAxis": "normalizedAxis == ranks.X - 1",
     "floatInput": "dtypes.T == \"f32\" or dtypes.T == \"f16\"",
-    "workgroupSize": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
-    "workgroupStorageFits": "workgroupSize * 8 <= device.limits.maxComputeWorkgroupStorageSize",
     "outputDispatchFits": "outputPositions <= foldedDispatchCapacity",
     "axisSharedSize": "pow2ceil(max(1, axisDim))",
     "axisSharedFits": "axisSharedSize <= tunables.AXIS_SHARED_MAX_VALUES and axisSharedSize * 8 <= device.limits.maxComputeWorkgroupStorageSize",
     "top1Blocks": "ceilDiv(axisDim, tunables.TOP1_BLOCK_SIZE)",
@@ -110,631 +85,216 @@
     "secondMergeDispatchFits": "outputPositions * secondGroups <= foldedDispatchCapacity",
     "largeSharedFits": "largeBlockSize * 8 <= device.limits.maxComputeWorkgroupStorageSize and mergeChunkSize * 8 <= device.limits.maxComputeWorkgroupStorageSize"
   },
-  "bindingSets": {
-    "smallRows": [
-      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
-      { "name": "values", "arg": "values", "semantic": "Values", "buffer": { "type": "storage" }, "elementType": "$T" },
-      {
-        "name": "indices",
-        "arg": "indices",
-        "semantic": "Indices",
-        "buffer": { "type": "storage" },
-        "elementType": "$I"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "outputPositions" },
-            { "name": "cols", "type": "u32", "value": "axisDim" },
-            { "name": "k", "type": "u32", "value": "args.k" }
-          ]
-        }
-      }
-    ],
-    "axisGeometry": [
-      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
-      { "name": "values", "arg": "values", "semantic": "Values", "buffer": { "type": "storage" }, "elementType": "$T" },
-      {
-        "name": "indices",
-        "arg": "indices",
-        "semantic": "Indices",
-        "buffer": { "type": "storage" },
-        "elementType": "$I"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "outputPositions", "type": "u32", "value": "outputPositions" },
-            { "name": "axis", "type": "u32", "value": "axisDim" },
-            { "name": "inner", "type": "u32", "value": "axisInner" },
-            { "name": "k", "type": "u32", "value": "args.k" }
-          ]
-        }
-      }
-    ],
-    "directTop1": [
-      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
-      { "name": "values", "arg": "values", "semantic": "Values", "buffer": { "type": "storage" }, "elementType": "$T" },
-      {
-        "name": "indices",
-        "arg": "indices",
-        "semantic": "Indices",
-        "buffer": { "type": "storage" },
-        "elementType": "$I"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "outputPositions" },
-            { "name": "cols", "type": "u32", "value": "axisDim" }
-          ]
-        }
-      }
-    ],
-    "top1Block": [
-      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
-      { "name": "candidateVals", "semantic": "candidateVals", "buffer": { "type": "storage" }, "elementType": "f32" },
-      { "name": "candidateIdxs", "semantic": "candidateIdxs", "buffer": { "type": "storage" }, "elementType": "u32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "outputPositions" },
-            { "name": "cols", "type": "u32", "value": "axisDim" },
-            { "name": "blocks", "type": "u32", "value": "top1Blocks" }
-          ]
-        }
-      }
-    ],
-    "top1Output": [
-      {
-        "name": "candidateVals",
-        "semantic": "candidateVals",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      {
-        "name": "candidateIdxs",
-        "semantic": "candidateIdxs",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      { "name": "values", "arg": "values", "semantic": "Values", "buffer": { "type": "storage" }, "elementType": "$T" },
-      {
-        "name": "indices",
-        "arg": "indices",
-        "semantic": "Indices",
-        "buffer": { "type": "storage" },
-        "elementType": "$I"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "outputPositions" },
-            { "name": "blocks", "type": "u32", "value": "top1Blocks" }
-          ]
-        }
-      }
-    ],
-    "largeBlock": [
-      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
-      { "name": "candidateVals", "semantic": "candidateVals", "buffer": { "type": "storage" }, "elementType": "f32" },
-      { "name": "candidateIdxs", "semantic": "candidateIdxs", "buffer": { "type": "storage" }, "elementType": "u32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "outputPositions" },
-            { "name": "cols", "type": "u32", "value": "axisDim" },
-            { "name": "k", "type": "u32", "value": "args.k" },
-            { "name": "blocks", "type": "u32", "value": "largeBlockCount" }
-          ]
-        }
-      }
-    ],
-    "mergeOutput": [
-      {
-        "name": "candidateVals",
-        "semantic": "candidateVals",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      {
-        "name": "candidateIdxs",
-        "semantic": "candidateIdxs",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      { "name": "values", "arg": "values", "semantic": "Values", "buffer": { "type": "storage" }, "elementType": "$T" },
-      {
-        "name": "indices",
-        "arg": "indices",
-        "semantic": "Indices",
-        "buffer": { "type": "storage" },
-        "elementType": "$I"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "outputPositions" },
-            { "name": "inCandidates", "type": "u32", "value": "firstCandidates" },
-            { "name": "k", "type": "u32", "value": "args.k" }
-          ]
-        }
-      }
-    ],
-    "mergeScratch": [
-      {
-        "name": "candidateVals",
-        "semantic": "candidateVals",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      {
-        "name": "candidateIdxs",
-        "semantic": "candidateIdxs",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      { "name": "candidateVals2", "semantic": "candidateVals2", "buffer": { "type": "storage" }, "elementType": "f32" },
-      { "name": "candidateIdxs2", "semantic": "candidateIdxs2", "buffer": { "type": "storage" }, "elementType": "u32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "outputPositions" },
-            { "name": "inCandidates", "type": "u32", "value": "firstCandidates" },
-            { "name": "outGroups", "type": "u32", "value": "secondGroups" },
-            { "name": "k", "type": "u32", "value": "args.k" }
-          ]
-        }
-      }
-    ],
-    "mergeOutput2": [
-      {
-        "name": "candidateVals",
-        "semantic": "candidateVals2",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      {
-        "name": "candidateIdxs",
-        "semantic": "candidateIdxs2",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      { "name": "values", "arg": "values", "semantic": "Values", "buffer": { "type": "storage" }, "elementType": "$T" },
-      {
-        "name": "indices",
-        "arg": "indices",
-        "semantic": "Indices",
-        "buffer": { "type": "storage" },
-        "elementType": "$I"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "outputPositions" },
-            { "name": "inCandidates", "type": "u32", "value": "secondCandidates" },
-            { "name": "k", "type": "u32", "value": "args.k" }
-          ]
-        }
-      }
-    ],
-    "largeBlockIntKeys": [
-      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
-      { "name": "candidateVals", "semantic": "candidateVals", "buffer": { "type": "storage" }, "elementType": "u32" },
-      { "name": "candidateIdxs", "semantic": "candidateIdxs", "buffer": { "type": "storage" }, "elementType": "u32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "outputPositions" },
-            { "name": "cols", "type": "u32", "value": "axisDim" },
-            { "name": "k", "type": "u32", "value": "args.k" },
-            { "name": "blocks", "type": "u32", "value": "largeBlockCount" }
-          ]
-        }
-      }
-    ],
-    "mergeOutputIntKeys": [
-      {
-        "name": "candidateVals",
-        "semantic": "candidateVals",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      {
-        "name": "candidateIdxs",
-        "semantic": "candidateIdxs",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      { "name": "values", "arg": "values", "semantic": "Values", "buffer": { "type": "storage" }, "elementType": "$T" },
-      {
-        "name": "indices",
-        "arg": "indices",
-        "semantic": "Indices",
-        "buffer": { "type": "storage" },
-        "elementType": "$I"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "outputPositions" },
-            { "name": "inCandidates", "type": "u32", "value": "firstCandidates" },
-            { "name": "k", "type": "u32", "value": "args.k" }
-          ]
-        }
-      }
-    ],
-    "mergeScratchIntKeys": [
-      {
-        "name": "candidateVals",
-        "semantic": "candidateVals",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      {
-        "name": "candidateIdxs",
-        "semantic": "candidateIdxs",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      { "name": "candidateVals2", "semantic": "candidateVals2", "buffer": { "type": "storage" }, "elementType": "u32" },
-      { "name": "candidateIdxs2", "semantic": "candidateIdxs2", "buffer": { "type": "storage" }, "elementType": "u32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "outputPositions" },
-            { "name": "inCandidates", "type": "u32", "value": "firstCandidates" },
-            { "name": "outGroups", "type": "u32", "value": "secondGroups" },
-            { "name": "k", "type": "u32", "value": "args.k" }
-          ]
-        }
-      }
-    ],
-    "mergeOutput2IntKeys": [
-      {
-        "name": "candidateVals",
-        "semantic": "candidateVals2",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      {
-        "name": "candidateIdxs",
-        "semantic": "candidateIdxs2",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      { "name": "values", "arg": "values", "semantic": "Values", "buffer": { "type": "storage" }, "elementType": "$T" },
-      {
-        "name": "indices",
-        "arg": "indices",
-        "semantic": "Indices",
-        "buffer": { "type": "storage" },
-        "elementType": "$I"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "outputPositions" },
-            { "name": "inCandidates", "type": "u32", "value": "secondCandidates" },
-            { "name": "k", "type": "u32", "value": "args.k" }
-          ]
-        }
-      }
-    ],
-    "mergeScratch2": [
-      {
-        "name": "candidateVals",
-        "semantic": "candidateVals2",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      {
-        "name": "candidateIdxs",
-        "semantic": "candidateIdxs2",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      { "name": "candidateVals2", "semantic": "candidateVals3", "buffer": { "type": "storage" }, "elementType": "f32" },
-      { "name": "candidateIdxs2", "semantic": "candidateIdxs3", "buffer": { "type": "storage" }, "elementType": "u32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "outputPositions" },
-            { "name": "inCandidates", "type": "u32", "value": "secondCandidates" },
-            { "name": "outGroups", "type": "u32", "value": "thirdGroups" },
-            { "name": "k", "type": "u32", "value": "args.k" }
-          ]
-        }
-      }
-    ],
-    "mergeOutput3": [
-      {
-        "name": "candidateVals",
-        "semantic": "candidateVals3",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      {
-        "name": "candidateIdxs",
-        "semantic": "candidateIdxs3",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      { "name": "values", "arg": "values", "semantic": "Values", "buffer": { "type": "storage" }, "elementType": "$T" },
-      {
-        "name": "indices",
-        "arg": "indices",
-        "semantic": "Indices",
-        "buffer": { "type": "storage" },
-        "elementType": "$I"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "outputPositions" },
-            { "name": "inCandidates", "type": "u32", "value": "thirdCandidates" },
-            { "name": "k", "type": "u32", "value": "args.k" }
-          ]
-        }
-      }
-    ],
-    "mergeScratch2IntKeys": [
-      {
-        "name": "candidateVals",
-        "semantic": "candidateVals2",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      {
-        "name": "candidateIdxs",
-        "semantic": "candidateIdxs2",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      { "name": "candidateVals2", "semantic": "candidateVals3", "buffer": { "type": "storage" }, "elementType": "u32" },
-      { "name": "candidateIdxs2", "semantic": "candidateIdxs3", "buffer": { "type": "storage" }, "elementType": "u32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "outputPositions" },
-            { "name": "inCandidates", "type": "u32", "value": "secondCandidates" },
-            { "name": "outGroups", "type": "u32", "value": "thirdGroups" },
-            { "name": "k", "type": "u32", "value": "args.k" }
-          ]
-        }
-      }
-    ],
-    "mergeOutput3IntKeys": [
-      {
-        "name": "candidateVals",
-        "semantic": "candidateVals3",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      {
-        "name": "candidateIdxs",
-        "semantic": "candidateIdxs3",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      { "name": "values", "arg": "values", "semantic": "Values", "buffer": { "type": "storage" }, "elementType": "$T" },
-      {
-        "name": "indices",
-        "arg": "indices",
-        "semantic": "Indices",
-        "buffer": { "type": "storage" },
-        "elementType": "$I"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "outputPositions" },
-            { "name": "inCandidates", "type": "u32", "value": "thirdCandidates" },
-            { "name": "k", "type": "u32", "value": "args.k" }
-          ]
-        }
-      }
-    ],
-    "largeBlockStrided": [
-      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
-      { "name": "candidateVals", "semantic": "candidateVals", "buffer": { "type": "storage" }, "elementType": "f32" },
-      { "name": "candidateIdxs", "semantic": "candidateIdxs", "buffer": { "type": "storage" }, "elementType": "u32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "outputPositions" },
-            { "name": "cols", "type": "u32", "value": "axisDim" },
-            { "name": "k", "type": "u32", "value": "args.k" },
-            { "name": "inner", "type": "u32", "value": "axisInner" },
-            { "name": "blocks", "type": "u32", "value": "largeBlockCount" }
-          ]
-        }
-      }
-    ],
-    "mergeOutputStrided": [
-      {
-        "name": "candidateVals",
-        "semantic": "candidateVals",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      {
-        "name": "candidateIdxs",
-        "semantic": "candidateIdxs",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      { "name": "values", "arg": "values", "semantic": "Values", "buffer": { "type": "storage" }, "elementType": "$T" },
-      {
-        "name": "indices",
-        "arg": "indices",
-        "semantic": "Indices",
-        "buffer": { "type": "storage" },
-        "elementType": "$I"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "outputPositions" },
-            { "name": "inCandidates", "type": "u32", "value": "firstCandidates" },
-            { "name": "k", "type": "u32", "value": "args.k" },
-            { "name": "inner", "type": "u32", "value": "axisInner" }
-          ]
-        }
-      }
-    ],
-    "mergeOutput2Strided": [
-      {
-        "name": "candidateVals",
-        "semantic": "candidateVals2",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      {
-        "name": "candidateIdxs",
-        "semantic": "candidateIdxs2",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      { "name": "values", "arg": "values", "semantic": "Values", "buffer": { "type": "storage" }, "elementType": "$T" },
-      {
-        "name": "indices",
-        "arg": "indices",
-        "semantic": "Indices",
-        "buffer": { "type": "storage" },
-        "elementType": "$I"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "rows", "type": "u32", "value": "outputPositions" },
-            { "name": "inCandidates", "type": "u32", "value": "secondCandidates" },
-            { "name": "k", "type": "u32", "value": "args.k" },
-            { "name": "inner", "type": "u32", "value": "axisInner" }
-          ]
-        }
-      }
-    ],
-    "noopParams": [
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "dummy", "type": "u32", "value": 0 }] }
-      }
-    ]
   },
   "variants": [
     {
       "id": "axis_smallk_tournament",
-      "description": "Scans a strided non-last axis with one workgroup per output position and retains only a short candidate list instead of sorting the whole axis. It is favored when many independent outputs amortize the strided scan and remains the bounded-storage route when the axis does not fit shared memory.",
       "priority": 25,
       "demoteWhen": ["axisSharedFits and not (axisInner > 1 and outputPositions >= tunables.TOURNAMENT_MIN_OUTPUTS)"],
-      "when": ["baseContract", "attrs.sorted == 1", "not lastAxis", "args.k >= 1", "args.k <= tunables.TOURNAMENT_MAX_K", "axisDim <= tunables.TOURNAMENT_MAX_AXIS", "workgroupStorageFits", "outputDispatchFits"],
-      "constants": {
         "scalar": "dtypes.T",
         "usesF16": "dtypes.T == \"f16\"",
-        "localItems": "ceilDiv(axisDim, workgroupSize)"
       },
       "passes": [
         {
           "id": "main",
           "name": "TopK.AxisSmallKTournament",
-          "source": { "shader": "topk-strided-smallk.wgsl.jinja", "inputs": { "nativeValues": true } },
-          "bindings": "axisGeometry",
-          "dispatch": { "workgroups": "outputPositions" }
         }
       ]
     },
     {
       "id": "last_axis_large_top1",
-      "description": "Finds one winner on rows too wide for the direct shared-memory route by reducing blocks to scratch candidates. A second pass selects the final candidate.",
       "priority": 20,
       "when": ["baseContract", "floatInput", "lastAxis", "args.k == 1", "axisDim > tunables.AXIS_SHARED_MAX_VALUES", "workgroupStorageFits", "top1ScratchFits", "outputDispatchFits"],
-      "constants": {
         "scalar": "dtypes.T",
         "usesF16": "dtypes.T == \"f16\"",
-        "top1BlockSize": "tunables.TOP1_BLOCK_SIZE"
       },
       "intermediates": [
         { "id": "candidateVals", "dtype": "float32", "shape": "[top1ScratchElements]" },
@@ -744,41 +304,74 @@
         {
           "id": "block",
           "name": "TopK.LargeTop1Block",
-          "source": {
-            "shader": "topk-top1-last-axis.wgsl.jinja",
-            "inputs": { "stage": "\"block\"", "nativeValues": false }
-          },
-          "bindings": "top1Block",
-          "dispatch": { "workgroups": "outputPositions * top1Blocks" }
         },
         {
           "id": "output",
           "name": "TopK.LargeTop1Output",
-          "source": {
-            "shader": "topk-top1-last-axis.wgsl.jinja",
-            "inputs": { "stage": "\"output\"", "nativeValues": false }
-          },
-          "bindings": "top1Output",
-          "dispatch": { "workgroups": "outputPositions" }
         }
       ]
     },
     {
       "id": "last_axis_top1",
       "priority": 15,
-      "supersededBy": ["last_axis_large_top1"],
       "when": ["baseContract", "lastAxis", "args.k == 1", "workgroupStorageFits", "outputDispatchFits"],
-      "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" },
       "passes": [
         {
           "id": "main",
           "name": "TopK.LastAxisTop1",
-          "source": {
-            "shader": "topk-top1-last-axis.wgsl.jinja",
-            "inputs": { "stage": "\"direct\"", "nativeValues": true }
-          },
-          "bindings": "directTop1",
-          "dispatch": { "workgroups": "outputPositions" }
         }
       ]
     },
@@ -786,14 +379,24 @@
       "id": "axis_bitonic",
       "priority": 0,
       "when": ["baseContract", "axisSharedFits", "outputDispatchFits"],
-      "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"", "sharedSize": "axisSharedSize" },
       "passes": [
         {
           "id": "main",
           "name": "TopK.AxisBitonic",
           "shader": "topk-axis.wgsl.jinja",
-          "bindings": "axisGeometry",
-          "dispatch": { "workgroups": "0 if args.k == 0 else outputPositions" }
         }
       ]
     },
@@ -806,40 +409,56 @@
           "id": "noop",
           "name": "TopK.Top0",
           "shader": "topk-noop.wgsl.jinja",
-          "bindings": "noopParams",
           "dispatch": { "x": 0 }
         }
       ]
     },
     {
       "id": "subgroup_rows_smallk",
-      "description": "Retains only K candidates per row and selects a row-parallel reduction layout according to row width and available subgroup geometry. It avoids shared storage proportional to the full axis.",
       "priority": 35,
-      "when": ["baseContract", "floatInput", "subgroupsWave32", "device.wgslLanguageFeatures.has(\"subgroup_id\")", "ranks.X == 2", "lastAxis", "attrs.sorted == 1", "args.k >= 1", "args.k <= tunables.SMALL_ROWS_MAX_K", "axisDim >= device.adapterInfo.subgroupMinSize * 4", "axisDim <= workgroupSize", "outputPositions >= workgroupSize * 4", "workgroupSize % device.adapterInfo.subgroupMinSize == 0", "outputDispatchFits"],
-      "constants": {
         "scalar": "dtypes.T",
         "usesF16": "dtypes.T == \"f16\"",
         "subgroupWidth": "device.adapterInfo.subgroupMinSize",
-        "subgroupsPerWorkgroup": "workgroupSize / device.adapterInfo.subgroupMinSize",
-        "localItems": "ceilDiv(axisDim, device.adapterInfo.subgroupMinSize)"
       },
       "passes": [
         {
           "id": "main",
           "name": "TopK.SubgroupRowsSmallK",
-          "source": { "shader": "topk-subgroup-rows.wgsl.jinja", "inputs": { "nativeValues": true } },
-          "subgroupCollectivesWidth": 32,
-          "bindings": "smallRows",
-          "dispatch": { "workgroups": "ceilDiv(outputPositions, workgroupSize / device.adapterInfo.subgroupMinSize)" }
         }
       ]
     },
     {
       "id": "subgroup_min_rows_smallk",
-      "description": "Retains only K candidates per row and selects a row-parallel reduction layout according to row width and available subgroup geometry. It avoids shared storage proportional to the full axis.",
       "priority": 35,
-      "when": ["baseContract", "floatInput", "not subgroupsWave32", "device.features.has(\"subgroups\")", "device.wgslLanguageFeatures.has(\"subgroup_id\")", "has(device.adapterInfo, \"subgroupMinSize\")", "has(device.adapterInfo, \"subgroupMaxSize\")", "device.adapterInfo.subgroupMinSize >= 4", "device.adapterInfo.subgroupMinSize <= device.adapterInfo.subgroupMaxSize", "device.adapterInfo.subgroupMinSize <= device.limits.maxComputeInvocationsPerWorkgroup", "device.adapterInfo.subgroupMinSize <= device.limits.maxComputeWorkgroupSizeX", "ranks.X == 2", "lastAxis", "attrs.sorted == 1", "args.k >= 1", "args.k <= tunables.SMALL_ROWS_MAX_K", "axisDim > tunables.SMALL_ROWS_MAX_AXIS", "axisDim <= tunables.PORTABLE_ROWS_MAX_AXIS", "outputPositions >= device.adapterInfo.subgroupMinSize * 4", "outputDispatchFits"],
-      "constants": {
         "scalar": "dtypes.T",
         "usesF16": "dtypes.T == \"f16\"",
         "workgroupSize": "device.adapterInfo.subgroupMinSize",
@@ -851,46 +470,88 @@
         {
           "id": "main",
           "name": "TopK.SubgroupMinRowsSmallK",
-          "source": { "shader": "topk-subgroup-rows.wgsl.jinja", "inputs": { "nativeValues": true } },
-          "bindings": "smallRows",
-          "dispatch": { "workgroups": "outputPositions" }
         }
       ]
     },
     {
       "id": "portable_rows_smallk",
-      "description": "Retains only K candidates per row and selects a row-parallel reduction layout according to row width and available subgroup geometry. It avoids shared storage proportional to the full axis.",
       "priority": 34,
-      "when": ["baseContract", "floatInput", "not subgroupsWave32 or not device.wgslLanguageFeatures.has(\"subgroup_id\")", "ranks.X == 2", "lastAxis", "attrs.sorted == 1", "args.k >= 1", "args.k <= tunables.SMALL_ROWS_MAX_K", "axisDim > tunables.SMALL_ROWS_MAX_AXIS", "axisDim <= tunables.PORTABLE_ROWS_MAX_AXIS", "outputPositions >= tunables.PORTABLE_ROWS_MIN", "tunables.PORTABLE_ROW_WORKGROUP_SIZE <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.PORTABLE_ROW_WORKGROUP_SIZE <= device.limits.maxComputeWorkgroupSizeX", "tunables.PORTABLE_ROW_WORKGROUP_SIZE * 8 + 4 <= device.limits.maxComputeWorkgroupStorageSize", "outputDispatchFits"],
-      "constants": {
         "scalar": "dtypes.T",
         "usesF16": "dtypes.T == \"f16\"",
         "rowWorkgroupSize": "tunables.PORTABLE_ROW_WORKGROUP_SIZE",
-        "localItems": "ceilDiv(axisDim, tunables.PORTABLE_ROW_WORKGROUP_SIZE)"
       },
       "passes": [
         {
           "id": "main",
           "name": "TopK.PortableRowsSmallK",
-          "source": { "shader": "topk-portable-rows-smallk.wgsl.jinja", "inputs": { "nativeValues": true } },
-          "bindings": "smallRows",
-          "dispatch": { "workgroups": "outputPositions" }
         }
       ]
     },
     {
       "id": "small_rows_batched",
-      "description": "Retains only K candidates per row and selects a row-parallel reduction layout according to row width and available subgroup geometry. It avoids shared storage proportional to the full axis.",
       "priority": 30,
-      "when": ["baseContract", "floatInput", "ranks.X == 2", "lastAxis", "attrs.sorted == 1", "args.k >= 1", "args.k <= tunables.SMALL_ROWS_MAX_K", "axisOuter >= tunables.SMALL_ROWS_MIN_ROWS", "axisDim <= tunables.SMALL_ROWS_MAX_AXIS", "outputDispatchFits"],
-      "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" },
       "passes": [
         {
           "id": "main",
           "name": "TopK.SmallRowsBatched",
-          "source": { "shader": "topk-small-rows-batched.wgsl.jinja", "inputs": { "nativeValues": true } },
-          "bindings": "smallRows",
-          "dispatch": { "threads": "outputPositions", "workgroupSize": "workgroupSize" }
         }
       ]
     },
@@ -898,26 +559,37 @@
       "id": "last_axis_large_one_merge",
       "priority": 10,
       "when": ["baseContract", "lastAxis", "args.k > 1", "axisDim > tunables.AXIS_SHARED_MAX_VALUES", "firstCandidates <= mergeChunkSize", "firstScratchFits", "largeSharedFits", "largeBlockDispatchFits", "outputDispatchFits", "floatInput"],
-      "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"", "blockSize": "largeBlockSize" },
       "intermediates": [
-        { "id": "candidateVals", "dtype": "float32", "shape": "[firstScratchElements]" },
         { "id": "candidateIdxs", "dtype": "uint32", "shape": "[firstScratchElements]" }
       ],
       "passes": [
         {
           "id": "block",
           "name": "TopK.LargeBlock",
-          "source": { "shader": "topk-large-block.wgsl.jinja", "inputs": { "stage": "\"block\"" } },
-          "bindings": "largeBlock",
-          "dispatch": { "workgroups": "outputPositions * largeBlockCount" }
         },
         {
           "id": "merge",
           "name": "TopK.LargeMergeOutput",
-          "source": { "shader": "topk-large-block.wgsl.jinja", "inputs": { "stage": "\"output\"" } },
-          "constants": { "sharedSize": "pow2ceil(firstCandidates)" },
-          "bindings": "mergeOutput",
-          "dispatch": { "workgroups": "outputPositions" }
         }
       ]
     },
@@ -925,36 +597,51 @@
       "id": "last_axis_large_two_merge",
       "priority": 10,
       "when": ["baseContract", "lastAxis", "args.k > 1", "axisDim > tunables.AXIS_SHARED_MAX_VALUES", "firstCandidates > mergeChunkSize", "secondCandidates <= mergeChunkSize", "firstScratchFits", "secondScratchFits", "largeSharedFits", "largeBlockDispatchFits", "secondMergeDispatchFits", "outputDispatchFits", "floatInput"],
-      "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"", "blockSize": "largeBlockSize" },
       "intermediates": [
-        { "id": "candidateVals", "dtype": "float32", "shape": "[firstScratchElements]" },
         { "id": "candidateIdxs", "dtype": "uint32", "shape": "[firstScratchElements]" },
-        { "id": "candidateVals2", "dtype": "float32", "shape": "[secondScratchElements]" },
         { "id": "candidateIdxs2", "dtype": "uint32", "shape": "[secondScratchElements]" }
       ],
       "passes": [
         {
           "id": "block",
           "name": "TopK.LargeBlock",
-          "source": { "shader": "topk-large-block.wgsl.jinja", "inputs": { "stage": "\"block\"" } },
-          "bindings": "largeBlock",
-          "dispatch": { "workgroups": "outputPositions * largeBlockCount" }
         },
         {
           "id": "merge_scratch",
           "name": "TopK.LargeMergeScratch",
-          "source": { "shader": "topk-large-block.wgsl.jinja", "inputs": { "stage": "\"scratch\"" } },
-          "constants": { "sharedSize": "mergeChunkSize" },
-          "bindings": "mergeScratch",
-          "dispatch": { "workgroups": "outputPositions * secondGroups" }
         },
         {
           "id": "merge",
           "name": "TopK.LargeMergeOutput",
-          "source": { "shader": "topk-large-block.wgsl.jinja", "inputs": { "stage": "\"output\"" } },
-          "constants": { "sharedSize": "pow2ceil(secondCandidates)" },
-          "bindings": "mergeOutput2",
-          "dispatch": { "workgroups": "outputPositions" }
         }
       ]
     },
@@ -962,7 +649,13 @@
       "id": "last_axis_large_one_merge_int",
       "priority": 10,
       "when": ["baseContract", "lastAxis", "args.k > 1", "axisDim > tunables.AXIS_SHARED_MAX_VALUES", "firstCandidates <= mergeChunkSize", "firstScratchFits", "largeSharedFits", "largeBlockDispatchFits", "outputDispatchFits", "not floatInput"],
-      "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"", "blockSize": "largeBlockSize" },
       "intermediates": [
         { "id": "candidateVals", "dtype": "uint32", "shape": "[firstScratchElements]" },
         { "id": "candidateIdxs", "dtype": "uint32", "shape": "[firstScratchElements]" }
@@ -971,17 +664,22 @@
         {
           "id": "block",
           "name": "TopK.LargeBlock",
-          "source": { "shader": "topk-large-block.wgsl.jinja", "inputs": { "stage": "\"block\"" } },
-          "bindings": "largeBlockIntKeys",
-          "dispatch": { "workgroups": "outputPositions * largeBlockCount" }
         },
         {
           "id": "merge",
           "name": "TopK.LargeMergeOutput",
-          "source": { "shader": "topk-large-block.wgsl.jinja", "inputs": { "stage": "\"output\"" } },
-          "constants": { "sharedSize": "pow2ceil(firstCandidates)" },
-          "bindings": "mergeOutputIntKeys",
-          "dispatch": { "workgroups": "outputPositions" }
         }
       ]
     },
@@ -989,7 +687,13 @@
       "id": "last_axis_large_two_merge_int",
       "priority": 10,
       "when": ["baseContract", "lastAxis", "args.k > 1", "axisDim > tunables.AXIS_SHARED_MAX_VALUES", "firstCandidates > mergeChunkSize", "secondCandidates <= mergeChunkSize", "firstScratchFits", "secondScratchFits", "largeSharedFits", "largeBlockDispatchFits", "secondMergeDispatchFits", "outputDispatchFits", "not floatInput"],
-      "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"", "blockSize": "largeBlockSize" },
       "intermediates": [
         { "id": "candidateVals", "dtype": "uint32", "shape": "[firstScratchElements]" },
         { "id": "candidateIdxs", "dtype": "uint32", "shape": "[firstScratchElements]" },
@@ -1000,25 +704,34 @@
         {
           "id": "block",
           "name": "TopK.LargeBlock",
-          "source": { "shader": "topk-large-block.wgsl.jinja", "inputs": { "stage": "\"block\"" } },
-          "bindings": "largeBlockIntKeys",
-          "dispatch": { "workgroups": "outputPositions * largeBlockCount" }
         },
         {
           "id": "merge_scratch",
           "name": "TopK.LargeMergeScratch",
-          "source": { "shader": "topk-large-block.wgsl.jinja", "inputs": { "stage": "\"scratch\"" } },
-          "constants": { "sharedSize": "mergeChunkSize" },
-          "bindings": "mergeScratchIntKeys",
-          "dispatch": { "workgroups": "outputPositions * secondGroups" }
         },
         {
           "id": "merge",
           "name": "TopK.LargeMergeOutput",
-          "source": { "shader": "topk-large-block.wgsl.jinja", "inputs": { "stage": "\"output\"" } },
-          "constants": { "sharedSize": "pow2ceil(secondCandidates)" },
-          "bindings": "mergeOutput2IntKeys",
-          "dispatch": { "workgroups": "outputPositions" }
         }
       ]
     },
@@ -1026,26 +739,37 @@
       "id": "axis_large_one_merge",
       "priority": 10,
       "when": ["baseContract", "not lastAxis", "args.k >= 1", "axisDim > tunables.AXIS_SHARED_MAX_VALUES", "firstCandidates <= mergeChunkSize", "firstScratchFits", "largeSharedFits", "largeBlockDispatchFits", "outputDispatchFits", "floatInput"],
-      "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"", "blockSize": "largeBlockSize" },
       "intermediates": [
-        { "id": "candidateVals", "dtype": "float32", "shape": "[firstScratchElements]" },
         { "id": "candidateIdxs", "dtype": "uint32", "shape": "[firstScratchElements]" }
       ],
       "passes": [
         {
           "id": "block",
           "name": "TopK.LargeBlock",
-          "source": { "shader": "topk-large-block.wgsl.jinja", "inputs": { "stage": "\"block\"", "strided": true } },
-          "bindings": "largeBlockStrided",
-          "dispatch": { "workgroups": "outputPositions * largeBlockCount" }
         },
         {
           "id": "merge",
           "name": "TopK.LargeMergeOutput",
-          "source": { "shader": "topk-large-block.wgsl.jinja", "inputs": { "stage": "\"output\"", "strided": true } },
-          "constants": { "sharedSize": "pow2ceil(firstCandidates)" },
-          "bindings": "mergeOutputStrided",
-          "dispatch": { "workgroups": "outputPositions" }
         }
       ]
     },
@@ -1053,36 +777,51 @@
       "id": "axis_large_two_merge",
       "priority": 10,
       "when": ["baseContract", "not lastAxis", "args.k >= 1", "axisDim > tunables.AXIS_SHARED_MAX_VALUES", "firstCandidates > mergeChunkSize", "secondCandidates <= mergeChunkSize", "firstScratchFits", "secondScratchFits", "largeSharedFits", "largeBlockDispatchFits", "secondMergeDispatchFits", "outputDispatchFits", "floatInput"],
-      "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"", "blockSize": "largeBlockSize" },
       "intermediates": [
-        { "id": "candidateVals", "dtype": "float32", "shape": "[firstScratchElements]" },
         { "id": "candidateIdxs", "dtype": "uint32", "shape": "[firstScratchElements]" },
-        { "id": "candidateVals2", "dtype": "float32", "shape": "[secondScratchElements]" },
         { "id": "candidateIdxs2", "dtype": "uint32", "shape": "[secondScratchElements]" }
       ],
       "passes": [
         {
           "id": "block",
           "name": "TopK.LargeBlock",
-          "source": { "shader": "topk-large-block.wgsl.jinja", "inputs": { "stage": "\"block\"", "strided": true } },
-          "bindings": "largeBlockStrided",
-          "dispatch": { "workgroups": "outputPositions * largeBlockCount" }
         },
         {
           "id": "merge_scratch",
           "name": "TopK.LargeMergeScratch",
-          "source": { "shader": "topk-large-block.wgsl.jinja", "inputs": { "stage": "\"scratch\"" } },
-          "constants": { "sharedSize": "mergeChunkSize" },
-          "bindings": "mergeScratch",
-          "dispatch": { "workgroups": "outputPositions * secondGroups" }
         },
         {
           "id": "merge",
           "name": "TopK.LargeMergeOutput",
-          "source": { "shader": "topk-large-block.wgsl.jinja", "inputs": { "stage": "\"output\"", "strided": true } },
-          "constants": { "sharedSize": "pow2ceil(secondCandidates)" },
-          "bindings": "mergeOutput2Strided",
-          "dispatch": { "workgroups": "outputPositions" }
         }
       ]
     },
@@ -1090,46 +829,65 @@
       "id": "last_axis_large_three_merge",
       "priority": 10,
       "when": ["baseContract", "lastAxis", "args.k > 1", "axisDim > tunables.AXIS_SHARED_MAX_VALUES", "firstCandidates > mergeChunkSize", "secondCandidates > mergeChunkSize", "thirdCandidates <= mergeChunkSize", "firstScratchFits", "secondScratchFits", "thirdScratchFits", "largeSharedFits", "largeBlockDispatchFits", "secondMergeDispatchFits", "thirdMergeDispatchFits", "outputDispatchFits", "floatInput"],
-      "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"", "blockSize": "largeBlockSize" },
       "intermediates": [
-        { "id": "candidateVals", "dtype": "float32", "shape": "[firstScratchElements]" },
         { "id": "candidateIdxs", "dtype": "uint32", "shape": "[firstScratchElements]" },
-        { "id": "candidateVals2", "dtype": "float32", "shape": "[secondScratchElements]" },
         { "id": "candidateIdxs2", "dtype": "uint32", "shape": "[secondScratchElements]" },
-        { "id": "candidateVals3", "dtype": "float32", "shape": "[thirdScratchElements]" },
         { "id": "candidateIdxs3", "dtype": "uint32", "shape": "[thirdScratchElements]" }
       ],
       "passes": [
         {
           "id": "block",
           "name": "TopK.LargeBlock",
-          "source": { "shader": "topk-large-block.wgsl.jinja", "inputs": { "stage": "\"block\"" } },
-          "bindings": "largeBlock",
-          "dispatch": { "workgroups": "outputPositions * largeBlockCount" }
         },
         {
           "id": "merge_scratch",
           "name": "TopK.LargeMergeScratch",
-          "source": { "shader": "topk-large-block.wgsl.jinja", "inputs": { "stage": "\"scratch\"" } },
-          "constants": { "sharedSize": "mergeChunkSize" },
-          "bindings": "mergeScratch",
-          "dispatch": { "workgroups": "outputPositions * secondGroups" }
         },
         {
           "id": "merge_scratch2",
           "name": "TopK.LargeMergeScratch2",
-          "source": { "shader": "topk-large-block.wgsl.jinja", "inputs": { "stage": "\"scratch\"" } },
-          "constants": { "sharedSize": "mergeChunkSize" },
-          "bindings": "mergeScratch2",
-          "dispatch": { "workgroups": "outputPositions * thirdGroups" }
         },
         {
           "id": "merge",
           "name": "TopK.LargeMergeOutput",
-          "source": { "shader": "topk-large-block.wgsl.jinja", "inputs": { "stage": "\"output\"" } },
-          "constants": { "sharedSize": "pow2ceil(thirdCandidates)" },
-          "bindings": "mergeOutput3",
-          "dispatch": { "workgroups": "outputPositions" }
         }
       ]
     },
@@ -1137,7 +895,13 @@
       "id": "last_axis_large_three_merge_int",
       "priority": 10,
       "when": ["baseContract", "lastAxis", "args.k > 1", "axisDim > tunables.AXIS_SHARED_MAX_VALUES", "firstCandidates > mergeChunkSize", "secondCandidates > mergeChunkSize", "thirdCandidates <= mergeChunkSize", "firstScratchFits", "secondScratchFits", "thirdScratchFits", "largeSharedFits", "largeBlockDispatchFits", "secondMergeDispatchFits", "thirdMergeDispatchFits", "outputDispatchFits", "not floatInput"],
-      "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"", "blockSize": "largeBlockSize" },
       "intermediates": [
         { "id": "candidateVals", "dtype": "uint32", "shape": "[firstScratchElements]" },
         { "id": "candidateIdxs", "dtype": "uint32", "shape": "[firstScratchElements]" },
@@ -1150,33 +914,46 @@
         {
           "id": "block",
           "name": "TopK.LargeBlock",
-          "source": { "shader": "topk-large-block.wgsl.jinja", "inputs": { "stage": "\"block\"" } },
-          "bindings": "largeBlockIntKeys",
-          "dispatch": { "workgroups": "outputPositions * largeBlockCount" }
         },
         {
           "id": "merge_scratch",
           "name": "TopK.LargeMergeScratch",
-          "source": { "shader": "topk-large-block.wgsl.jinja", "inputs": { "stage": "\"scratch\"" } },
-          "constants": { "sharedSize": "mergeChunkSize" },
-          "bindings": "mergeScratchIntKeys",
-          "dispatch": { "workgroups": "outputPositions * secondGroups" }
         },
         {
           "id": "merge_scratch2",
           "name": "TopK.LargeMergeScratch2",
-          "source": { "shader": "topk-large-block.wgsl.jinja", "inputs": { "stage": "\"scratch\"" } },
-          "constants": { "sharedSize": "mergeChunkSize" },
-          "bindings": "mergeScratch2IntKeys",
-          "dispatch": { "workgroups": "outputPositions * thirdGroups" }
         },
         {
           "id": "merge",
           "name": "TopK.LargeMergeOutput",
-          "source": { "shader": "topk-large-block.wgsl.jinja", "inputs": { "stage": "\"output\"" } },
-          "constants": { "sharedSize": "pow2ceil(thirdCandidates)" },
-          "bindings": "mergeOutput3IntKeys",
-          "dispatch": { "workgroups": "outputPositions" }
         }
       ]
     }

   "domain": "ai.onnx",
   "name": "TopK",
   "sinceVersion": 11,
+  "inputs": { "x": { "onnx": "X", "dtype": "T" } },
+  "outputs": {
+    "values": {
+      "onnx": "Values",
       "dtype": "T",
+      "rank": "ranks.x",
+      "shape": "prefix(shapes.x, normalizedAxis) + [args.k] + suffix(shapes.x, normalizedAxis + 1)"
     },
+    "indices": {
+      "onnx": "Indices",
       "dtype": "I",
+      "rank": "ranks.x",
+      "shape": "prefix(shapes.x, normalizedAxis) + [args.k] + suffix(shapes.x, normalizedAxis + 1)",
+      "storage": "uint32"
     }
   },
+  "args": { "k": { "kind": "u32", "onnx": "kernel.k" } },
+  "attributes": { "axis": { "default": -1 }, "largest": { "default": 1 }, "sorted": { "default": 1 } },
   "attributeConstraints": { "largest": { "values": [0, 1] }, "sorted": { "values": [0, 1] } },
   "typeConstraints": { "T": ["float32", "float16", "int8", "int16", "int32", "uint8", "uint32"], "I": ["int64"] },
   "tunables": {
+    "WORKGROUP_SIZE": { "default": 256 },
+    "AXIS_SHARED_MAX_VALUES": { "default": 2048 },
+    "TOP1_BLOCK_SIZE": { "default": 512 },
+    "SMALL_K_BLOCK_SIZE": { "default": 512 },
+    "LARGE_K_BLOCK_SIZE": { "default": 2048 },
+    "SMALL_K_LIMIT": { "default": 128 },
+    "MERGE_CHUNK_SIZE": { "default": 2048 },
+    "SMALL_ROWS_MIN_ROWS": { "default": 4096 },
+    "SMALL_ROWS_MAX_AXIS": { "default": 64 },
+    "SMALL_ROWS_MAX_K": { "default": 8 },
+    "PORTABLE_ROW_WORKGROUP_SIZE": { "default": 32 },
+    "PORTABLE_ROWS_MIN": { "default": 1024 },
+    "PORTABLE_ROWS_MAX_AXIS": { "default": 256 },
+    "TOURNAMENT_MAX_AXIS": { "default": 4096 },
+    "TOURNAMENT_MAX_K": { "default": 16 },
+    "TOURNAMENT_MIN_OUTPUTS": { "default": 256 }
   },
   "derive": {
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
+    "foldedDispatchCapacity": "min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
     "wave32Adapter": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 32 and device.adapterInfo.subgroupMaxSize == 32",
     "subgroupsWave32": "device.features.has(\"subgroups\") and wave32Adapter",
+    "normalizedAxis": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x",
+    "axisInRange": "ranks.x >= 1 and normalizedAxis >= 0 and normalizedAxis < ranks.x",
+    "axisOuter": "outer(shapes.x, normalizedAxis)",
+    "axisDim": "dim(shapes.x, normalizedAxis)",
+    "axisInner": "inner(shapes.x, normalizedAxis)",
     "outputPositions": "axisOuter * axisInner",
+    "outputShapeOk": "ranks.values == ranks.x and ranks.indices == ranks.x and outer(shapes.values, normalizedAxis) == axisOuter and outer(shapes.indices, normalizedAxis) == axisOuter and dim(shapes.values, normalizedAxis) == args.k and dim(shapes.indices, normalizedAxis) == args.k and inner(shapes.values, normalizedAxis) == axisInner and inner(shapes.indices, normalizedAxis) == axisInner",
     "dtypeOk": "f16Ok(dtypes.T)",
     "shapeContract": "axisInRange and outputShapeOk and args.k <= axisDim",
     "baseContract": "shapeContract and dtypeOk",
+    "lastAxis": "normalizedAxis == ranks.x - 1",
     "floatInput": "dtypes.T == \"f32\" or dtypes.T == \"f16\"",
+    "workgroupSizePlan": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
+    "workgroupStorageFits": "workgroupSizePlan * 8 <= device.limits.maxComputeWorkgroupStorageSize",
     "outputDispatchFits": "outputPositions <= foldedDispatchCapacity",
+    "sortRetainedSize": "pow2ceil(max(1, args.k))",
     "axisSharedSize": "pow2ceil(max(1, axisDim))",
     "axisSharedFits": "axisSharedSize <= tunables.AXIS_SHARED_MAX_VALUES and axisSharedSize * 8 <= device.limits.maxComputeWorkgroupStorageSize",
     "top1Blocks": "ceilDiv(axisDim, tunables.TOP1_BLOCK_SIZE)",
     "secondMergeDispatchFits": "outputPositions * secondGroups <= foldedDispatchCapacity",
     "largeSharedFits": "largeBlockSize * 8 <= device.limits.maxComputeWorkgroupStorageSize and mergeChunkSize * 8 <= device.limits.maxComputeWorkgroupStorageSize"
   },
+  "bindings": {
+    "x": { "buffer": "read-only-storage", "elementType": "$T" },
+    "values": { "buffer": "storage", "elementType": "$T" },
+    "indices": { "buffer": "storage", "elementType": "$I" },
+    "params": {
+      "buffer": "uniform",
+      "struct": [
+        { "name": "outputPositions", "type": "u32", "value": "outputPositions" },
+        { "name": "axis", "type": "u32", "value": "axisDim" },
+        { "name": "inner", "type": "u32", "value": "axisInner" },
+        { "name": "k", "type": "u32" }
+      ]
+    },
+    "candidateVals": { "buffer": "storage", "elementType": "f32" },
+    "candidateIdxs": { "buffer": "storage", "elementType": "u32" },
+    "candidateVals_2": { "name": "candidateVals", "buffer": "read-only-storage", "elementType": "f32" },
+    "candidateIdxs_2": { "name": "candidateIdxs", "buffer": "read-only-storage", "elementType": "u32" },
+    "params_7": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "outputPositions" },
+        { "name": "cols", "type": "u32", "value": "axisDim" },
+        { "name": "k", "type": "u32" },
+        { "name": "blocks", "type": "u32", "value": "largeBlockCount" }
+      ]
+    },
+    "candidateVals_3": { "name": "candidateVals", "buffer": "storage", "elementType": "u32" },
+    "params_8": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "outputPositions" },
+        { "name": "cols", "type": "u32", "value": "axisDim" },
+        { "name": "k", "type": "u32" },
+        { "name": "inner", "type": "u32", "value": "axisInner" },
+        { "name": "blocks", "type": "u32", "value": "largeBlockCount" }
+      ]
+    },
+    "params_9": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "outputPositions" },
+        { "name": "inCandidates", "type": "u32", "value": "firstCandidates" },
+        { "name": "k", "type": "u32" }
+      ]
+    },
+    "candidateVals_4": { "name": "candidateVals", "buffer": "read-only-storage", "elementType": "u32" },
+    "candidateIdxs2": { "buffer": "storage", "elementType": "u32" },
+    "params_11": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "outputPositions" },
+        { "name": "inCandidates", "type": "u32", "value": "firstCandidates" },
+        { "name": "outGroups", "type": "u32", "value": "secondGroups" },
+        { "name": "k", "type": "u32" }
+      ]
+    },
+    "candidateVals2_2": { "name": "candidateVals2", "buffer": "storage", "elementType": "u32" },
+    "candidateIdxs_3": {
+      "scratch": "candidateIdxs2",
+      "name": "candidateIdxs",
+      "buffer": "read-only-storage",
+      "elementType": "u32"
+    },
+    "params_12": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "outputPositions" },
+        { "name": "inCandidates", "type": "u32", "value": "secondCandidates" },
+        { "name": "k", "type": "u32" }
+      ]
+    },
+    "candidateVals_6": {
+      "scratch": "candidateVals2",
+      "name": "candidateVals",
+      "buffer": "read-only-storage",
+      "elementType": "u32"
+    },
+    "candidateIdxs2_2": {
+      "scratch": "candidateIdxs3",
+      "name": "candidateIdxs2",
+      "buffer": "storage",
+      "elementType": "u32"
+    },
+    "params_14": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "outputPositions" },
+        { "name": "inCandidates", "type": "u32", "value": "secondCandidates" },
+        { "name": "outGroups", "type": "u32", "value": "thirdGroups" },
+        { "name": "k", "type": "u32" }
+      ]
+    },
+    "candidateVals2_4": {
+      "scratch": "candidateVals3",
+      "name": "candidateVals2",
+      "buffer": "storage",
+      "elementType": "u32"
+    },
+    "candidateIdxs_4": {
+      "scratch": "candidateIdxs3",
+      "name": "candidateIdxs",
+      "buffer": "read-only-storage",
+      "elementType": "u32"
+    },
+    "params_15": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "outputPositions" },
+        { "name": "inCandidates", "type": "u32", "value": "thirdCandidates" },
+        { "name": "k", "type": "u32" }
+      ]
+    },
+    "candidateVals_8": {
+      "scratch": "candidateVals3",
+      "name": "candidateVals",
+      "buffer": "read-only-storage",
+      "elementType": "u32"
+    },
+    "params_9_source": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "outputPositions" },
+        { "name": "inCandidates", "type": "u32", "value": "firstCandidates" },
+        { "name": "k", "type": "u32" },
+        { "name": "cols", "type": "u32", "value": "axisDim" }
+      ]
+    },
+    "params_12_source": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "outputPositions" },
+        { "name": "inCandidates", "type": "u32", "value": "secondCandidates" },
+        { "name": "k", "type": "u32" },
+        { "name": "cols", "type": "u32", "value": "axisDim" }
+      ]
+    },
+    "params_10_source": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "outputPositions" },
+        { "name": "inCandidates", "type": "u32", "value": "firstCandidates" },
+        { "name": "k", "type": "u32" },
+        { "name": "inner", "type": "u32", "value": "axisInner" },
+        { "name": "cols", "type": "u32", "value": "axisDim" }
+      ]
+    },
+    "params_13_source": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "outputPositions" },
+        { "name": "inCandidates", "type": "u32", "value": "secondCandidates" },
+        { "name": "k", "type": "u32" },
+        { "name": "inner", "type": "u32", "value": "axisInner" },
+        { "name": "cols", "type": "u32", "value": "axisDim" }
+      ]
+    },
+    "params_15_source": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "outputPositions" },
+        { "name": "inCandidates", "type": "u32", "value": "thirdCandidates" },
+        { "name": "k", "type": "u32" },
+        { "name": "cols", "type": "u32", "value": "axisDim" }
+      ]
+    }
   },
   "variants": [
     {
       "id": "axis_smallk_tournament",
       "priority": 25,
+      "when": ["baseContract", "not lastAxis", "args.k >= 1", "args.k <= tunables.TOURNAMENT_MAX_K", "axisDim <= tunables.TOURNAMENT_MAX_AXIS", "workgroupStorageFits", "outputDispatchFits"],
       "demoteWhen": ["axisSharedFits and not (axisInner > 1 and outputPositions >= tunables.TOURNAMENT_MIN_OUTPUTS)"],
+      "derive": {
         "scalar": "dtypes.T",
         "usesF16": "dtypes.T == \"f16\"",
+        "localItems": "ceilDiv(axisDim, workgroupSizePlan)",
+        "workgroupSize": "workgroupSizePlan"
       },
       "passes": [
         {
           "id": "main",
           "name": "TopK.AxisSmallKTournament",
+          "shader": "topk-strided-smallk.wgsl.jinja",
+          "derive": { "nativeValues": true },
+          "bindings": ["x", "values", "indices", "params"],
+          "dispatch": { "x": "min(outputPositions, 65535)", "y": "ceilDiv(outputPositions, 65535)", "z": 1 }
         }
       ]
     },
     {
       "id": "last_axis_large_top1",
       "priority": 20,
       "when": ["baseContract", "floatInput", "lastAxis", "args.k == 1", "axisDim > tunables.AXIS_SHARED_MAX_VALUES", "workgroupStorageFits", "top1ScratchFits", "outputDispatchFits"],
+      "derive": {
         "scalar": "dtypes.T",
         "usesF16": "dtypes.T == \"f16\"",
+        "top1BlockSize": "tunables.TOP1_BLOCK_SIZE",
+        "workgroupSize": "workgroupSizePlan"
       },
       "intermediates": [
         { "id": "candidateVals", "dtype": "float32", "shape": "[top1ScratchElements]" },
         {
           "id": "block",
           "name": "TopK.LargeTop1Block",
+          "shader": "topk-top1-last-axis.wgsl.jinja",
+          "derive": { "stage": "\"block\"", "nativeValues": false },
+          "bindings": [
+            "x",
+            "candidateVals",
+            "candidateIdxs",
+            {
+              "name": "params",
+              "struct": [
+                { "name": "rows", "type": "u32", "value": "outputPositions" },
+                { "name": "cols", "type": "u32", "value": "axisDim" },
+                { "name": "blocks", "type": "u32", "value": "top1Blocks" }
+              ]
+            }
+          ],
+          "dispatch": {
+            "x": "min(outputPositions * top1Blocks, 65535)",
+            "y": "ceilDiv(outputPositions * top1Blocks, 65535)",
+            "z": 1
+          }
         },
         {
           "id": "output",
           "name": "TopK.LargeTop1Output",
+          "shader": "topk-top1-last-axis.wgsl.jinja",
+          "derive": { "stage": "\"output\"", "nativeValues": false },
+          "bindings": [
+            "candidateVals_2",
+            "candidateIdxs_2",
+            "values",
+            "indices",
+            {
+              "name": "params",
+              "struct": [
+                { "name": "rows", "type": "u32", "value": "outputPositions" },
+                { "name": "blocks", "type": "u32", "value": "top1Blocks" }
+              ]
+            }
+          ],
+          "dispatch": { "x": "min(outputPositions, 65535)", "y": "ceilDiv(outputPositions, 65535)", "z": 1 }
         }
       ]
     },
     {
       "id": "last_axis_top1",
       "priority": 15,
       "when": ["baseContract", "lastAxis", "args.k == 1", "workgroupStorageFits", "outputDispatchFits"],
+      "supersededBy": ["last_axis_large_top1"],
+      "derive": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"", "workgroupSize": "workgroupSizePlan" },
       "passes": [
         {
           "id": "main",
           "name": "TopK.LastAxisTop1",
+          "shader": "topk-top1-last-axis.wgsl.jinja",
+          "derive": { "stage": "\"direct\"", "nativeValues": true },
+          "bindings": [
+            "x",
+            "values",
+            "indices",
+            {
+              "name": "params",
+              "struct": [
+                { "name": "rows", "type": "u32", "value": "outputPositions" },
+                { "name": "cols", "type": "u32", "value": "axisDim" }
+              ]
+            }
+          ],
+          "dispatch": { "x": "min(outputPositions, 65535)", "y": "ceilDiv(outputPositions, 65535)", "z": 1 }
         }
       ]
     },
       "id": "axis_bitonic",
       "priority": 0,
       "when": ["baseContract", "axisSharedFits", "outputDispatchFits"],
+      "derive": {
+        "scalar": "dtypes.T",
+        "usesF16": "dtypes.T == \"f16\"",
+        "sharedSize": "axisSharedSize",
+        "workgroupSize": "workgroupSizePlan",
+        "floatOrderKeys": "floatInput"
+      },
       "passes": [
         {
           "id": "main",
           "name": "TopK.AxisBitonic",
           "shader": "topk-axis.wgsl.jinja",
+          "bindings": ["x", "values", "indices", "params"],
+          "dispatch": {
+            "x": "min(0 if args.k == 0 else outputPositions, 65535)",
+            "y": "ceilDiv(0 if args.k == 0 else outputPositions, 65535)",
+            "z": 1
+          }
         }
       ]
     },
           "id": "noop",
           "name": "TopK.Top0",
           "shader": "topk-noop.wgsl.jinja",
+          "bindings": [{ "name": "params", "struct": [{ "name": "dummy", "type": "u32", "value": 0 }] }],
           "dispatch": { "x": 0 }
         }
       ]
     },
     {
       "id": "subgroup_rows_smallk",
       "priority": 35,
+      "when": ["baseContract", "floatInput", "subgroupsWave32", "device.wgslLanguageFeatures.has(\"subgroup_id\")", "lastAxis", "args.k >= 1", "args.k <= tunables.SMALL_ROWS_MAX_K", "axisDim >= device.adapterInfo.subgroupMinSize * 4", "axisDim <= workgroupSizePlan", "outputPositions >= workgroupSizePlan * 4", "workgroupSizePlan % device.adapterInfo.subgroupMinSize == 0", "outputDispatchFits"],
+      "derive": {
         "scalar": "dtypes.T",
         "usesF16": "dtypes.T == \"f16\"",
         "subgroupWidth": "device.adapterInfo.subgroupMinSize",
+        "subgroupsPerWorkgroup": "workgroupSizePlan / device.adapterInfo.subgroupMinSize",
+        "localItems": "ceilDiv(axisDim, device.adapterInfo.subgroupMinSize)",
+        "workgroupSize": "workgroupSizePlan"
       },
       "passes": [
         {
           "id": "main",
           "name": "TopK.SubgroupRowsSmallK",
+          "shader": "topk-subgroup-rows.wgsl.jinja",
+          "derive": { "nativeValues": true },
+          "bindings": [
+            "x",
+            "values",
+            "indices",
+            {
+              "name": "params",
+              "struct": [
+                { "name": "rows", "type": "u32", "value": "outputPositions" },
+                { "name": "cols", "type": "u32", "value": "axisDim" },
+                { "name": "k", "type": "u32" }
+              ]
+            }
+          ],
+          "dispatch": {
+            "x": "min(ceilDiv(outputPositions, workgroupSize / device.adapterInfo.subgroupMinSize), 65535)",
+            "y": "ceilDiv(ceilDiv(outputPositions, workgroupSize / device.adapterInfo.subgroupMinSize), 65535)",
+            "z": 1
+          },
+          "subgroupCollectivesWidth": 32
         }
       ]
     },
     {
       "id": "subgroup_min_rows_smallk",
       "priority": 35,
+      "when": ["baseContract", "floatInput", "not subgroupsWave32", "device.features.has(\"subgroups\")", "device.wgslLanguageFeatures.has(\"subgroup_id\")", "has(device.adapterInfo, \"subgroupMinSize\")", "has(device.adapterInfo, \"subgroupMaxSize\")", "device.adapterInfo.subgroupMinSize >= 4", "device.adapterInfo.subgroupMinSize <= device.adapterInfo.subgroupMaxSize", "device.adapterInfo.subgroupMinSize <= device.limits.maxComputeInvocationsPerWorkgroup", "device.adapterInfo.subgroupMinSize <= device.limits.maxComputeWorkgroupSizeX", "lastAxis", "args.k >= 1", "args.k <= tunables.SMALL_ROWS_MAX_K", "axisDim > tunables.SMALL_ROWS_MAX_AXIS", "axisDim <= tunables.PORTABLE_ROWS_MAX_AXIS", "outputPositions >= device.adapterInfo.subgroupMinSize * 4", "outputDispatchFits"],
+      "derive": {
         "scalar": "dtypes.T",
         "usesF16": "dtypes.T == \"f16\"",
         "workgroupSize": "device.adapterInfo.subgroupMinSize",
         {
           "id": "main",
           "name": "TopK.SubgroupMinRowsSmallK",
+          "shader": "topk-subgroup-rows.wgsl.jinja",
+          "derive": { "nativeValues": true },
+          "bindings": [
+            "x",
+            "values",
+            "indices",
+            {
+              "name": "params",
+              "struct": [
+                { "name": "rows", "type": "u32", "value": "outputPositions" },
+                { "name": "cols", "type": "u32", "value": "axisDim" },
+                { "name": "k", "type": "u32" }
+              ]
+            }
+          ],
+          "dispatch": { "x": "min(outputPositions, 65535)", "y": "ceilDiv(outputPositions, 65535)", "z": 1 }
         }
       ]
     },
     {
       "id": "portable_rows_smallk",
       "priority": 34,
+      "when": ["baseContract", "floatInput", "not subgroupsWave32 or not device.wgslLanguageFeatures.has(\"subgroup_id\")", "lastAxis", "args.k >= 1", "args.k <= tunables.SMALL_ROWS_MAX_K", "axisDim > tunables.SMALL_ROWS_MAX_AXIS", "axisDim <= tunables.PORTABLE_ROWS_MAX_AXIS", "outputPositions >= tunables.PORTABLE_ROWS_MIN", "tunables.PORTABLE_ROW_WORKGROUP_SIZE <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.PORTABLE_ROW_WORKGROUP_SIZE <= device.limits.maxComputeWorkgroupSizeX", "tunables.PORTABLE_ROW_WORKGROUP_SIZE * 8 + 4 <= device.limits.maxComputeWorkgroupStorageSize", "outputDispatchFits"],
+      "derive": {
         "scalar": "dtypes.T",
         "usesF16": "dtypes.T == \"f16\"",
         "rowWorkgroupSize": "tunables.PORTABLE_ROW_WORKGROUP_SIZE",
+        "localItems": "ceilDiv(axisDim, tunables.PORTABLE_ROW_WORKGROUP_SIZE)",
+        "workgroupSize": "workgroupSizePlan"
       },
       "passes": [
         {
           "id": "main",
           "name": "TopK.PortableRowsSmallK",
+          "shader": "topk-portable-rows-smallk.wgsl.jinja",
+          "derive": { "nativeValues": true },
+          "bindings": [
+            "x",
+            "values",
+            "indices",
+            {
+              "name": "params",
+              "struct": [
+                { "name": "rows", "type": "u32", "value": "outputPositions" },
+                { "name": "cols", "type": "u32", "value": "axisDim" },
+                { "name": "k", "type": "u32" }
+              ]
+            }
+          ],
+          "dispatch": { "x": "min(outputPositions, 65535)", "y": "ceilDiv(outputPositions, 65535)", "z": 1 }
         }
       ]
     },
     {
       "id": "small_rows_batched",
       "priority": 30,
+      "when": ["baseContract", "floatInput", "lastAxis", "args.k >= 1", "args.k <= tunables.SMALL_ROWS_MAX_K", "axisOuter >= tunables.SMALL_ROWS_MIN_ROWS", "axisDim <= tunables.SMALL_ROWS_MAX_AXIS", "outputDispatchFits"],
+      "derive": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"", "workgroupSize": "workgroupSizePlan" },
       "passes": [
         {
           "id": "main",
           "name": "TopK.SmallRowsBatched",
+          "shader": "topk-small-rows-batched.wgsl.jinja",
+          "derive": { "nativeValues": true },
+          "bindings": [
+            "x",
+            "values",
+            "indices",
+            {
+              "name": "params",
+              "struct": [
+                { "name": "rows", "type": "u32", "value": "outputPositions" },
+                { "name": "cols", "type": "u32", "value": "axisDim" },
+                { "name": "k", "type": "u32" }
+              ]
+            }
+          ],
+          "dispatch": {
+            "x": "min(ceilDiv((outputPositions), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((outputPositions), (workgroupSize)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
       "id": "last_axis_large_one_merge",
       "priority": 10,
       "when": ["baseContract", "lastAxis", "args.k > 1", "axisDim > tunables.AXIS_SHARED_MAX_VALUES", "firstCandidates <= mergeChunkSize", "firstScratchFits", "largeSharedFits", "largeBlockDispatchFits", "outputDispatchFits", "floatInput"],
+      "derive": {
+        "scalar": "dtypes.T",
+        "usesF16": "dtypes.T == \"f16\"",
+        "blockSize": "largeBlockSize",
+        "workgroupSize": "workgroupSizePlan",
+        "floatOrderKeys": true
+      },
       "intermediates": [
+        { "id": "candidateVals", "dtype": "uint32", "shape": "[firstScratchElements]" },
         { "id": "candidateIdxs", "dtype": "uint32", "shape": "[firstScratchElements]" }
       ],
       "passes": [
         {
           "id": "block",
           "name": "TopK.LargeBlock",
+          "shader": "topk-large-block.wgsl.jinja",
+          "derive": { "stage": "\"block\"" },
+          "bindings": ["x", "candidateVals_3", "candidateIdxs", "params_7"],
+          "dispatch": {
+            "x": "min(outputPositions * largeBlockCount, 65535)",
+            "y": "ceilDiv(outputPositions * largeBlockCount, 65535)",
+            "z": 1
+          }
         },
         {
           "id": "merge",
           "name": "TopK.LargeMergeOutput",
+          "shader": "topk-large-block.wgsl.jinja",
+          "derive": { "sharedSize": "pow2ceil(firstCandidates)", "stage": "\"output\"" },
+          "bindings": ["candidateVals_4", "candidateIdxs_2", "values", "indices", "params_9_source", "x"],
+          "dispatch": { "x": "min(outputPositions, 65535)", "y": "ceilDiv(outputPositions, 65535)", "z": 1 }
         }
       ]
     },
       "id": "last_axis_large_two_merge",
       "priority": 10,
       "when": ["baseContract", "lastAxis", "args.k > 1", "axisDim > tunables.AXIS_SHARED_MAX_VALUES", "firstCandidates > mergeChunkSize", "secondCandidates <= mergeChunkSize", "firstScratchFits", "secondScratchFits", "largeSharedFits", "largeBlockDispatchFits", "secondMergeDispatchFits", "outputDispatchFits", "floatInput"],
+      "derive": {
+        "scalar": "dtypes.T",
+        "usesF16": "dtypes.T == \"f16\"",
+        "blockSize": "largeBlockSize",
+        "workgroupSize": "workgroupSizePlan",
+        "floatOrderKeys": true
+      },
       "intermediates": [
+        { "id": "candidateVals", "dtype": "uint32", "shape": "[firstScratchElements]" },
         { "id": "candidateIdxs", "dtype": "uint32", "shape": "[firstScratchElements]" },
+        { "id": "candidateVals2", "dtype": "uint32", "shape": "[secondScratchElements]" },
         { "id": "candidateIdxs2", "dtype": "uint32", "shape": "[secondScratchElements]" }
       ],
       "passes": [
         {
           "id": "block",
           "name": "TopK.LargeBlock",
+          "shader": "topk-large-block.wgsl.jinja",
+          "derive": { "stage": "\"block\"" },
+          "bindings": ["x", "candidateVals_3", "candidateIdxs", "params_7"],
+          "dispatch": {
+            "x": "min(outputPositions * largeBlockCount, 65535)",
+            "y": "ceilDiv(outputPositions * largeBlockCount, 65535)",
+            "z": 1
+          }
         },
         {
           "id": "merge_scratch",
           "name": "TopK.LargeMergeScratch",
+          "shader": "topk-large-block.wgsl.jinja",
+          "derive": { "sharedSize": "mergeChunkSize", "stage": "\"scratch\"" },
+          "bindings": ["candidateVals_4", "candidateIdxs_2", "candidateVals2_2", "candidateIdxs2", "params_11"],
+          "dispatch": {
+            "x": "min(outputPositions * secondGroups, 65535)",
+            "y": "ceilDiv(outputPositions * secondGroups, 65535)",
+            "z": 1
+          }
         },
         {
           "id": "merge",
           "name": "TopK.LargeMergeOutput",
+          "shader": "topk-large-block.wgsl.jinja",
+          "derive": { "sharedSize": "pow2ceil(secondCandidates)", "stage": "\"output\"" },
+          "bindings": ["candidateVals_6", "candidateIdxs_3", "values", "indices", "params_12_source", "x"],
+          "dispatch": { "x": "min(outputPositions, 65535)", "y": "ceilDiv(outputPositions, 65535)", "z": 1 }
         }
       ]
     },
       "id": "last_axis_large_one_merge_int",
       "priority": 10,
       "when": ["baseContract", "lastAxis", "args.k > 1", "axisDim > tunables.AXIS_SHARED_MAX_VALUES", "firstCandidates <= mergeChunkSize", "firstScratchFits", "largeSharedFits", "largeBlockDispatchFits", "outputDispatchFits", "not floatInput"],
+      "derive": {
+        "scalar": "dtypes.T",
+        "usesF16": "dtypes.T == \"f16\"",
+        "blockSize": "largeBlockSize",
+        "workgroupSize": "workgroupSizePlan",
+        "floatOrderKeys": false
+      },
       "intermediates": [
         { "id": "candidateVals", "dtype": "uint32", "shape": "[firstScratchElements]" },
         { "id": "candidateIdxs", "dtype": "uint32", "shape": "[firstScratchElements]" }
         {
           "id": "block",
           "name": "TopK.LargeBlock",
+          "shader": "topk-large-block.wgsl.jinja",
+          "derive": { "stage": "\"block\"" },
+          "bindings": ["x", "candidateVals_3", "candidateIdxs", "params_7"],
+          "dispatch": {
+            "x": "min(outputPositions * largeBlockCount, 65535)",
+            "y": "ceilDiv(outputPositions * largeBlockCount, 65535)",
+            "z": 1
+          }
         },
         {
           "id": "merge",
           "name": "TopK.LargeMergeOutput",
+          "shader": "topk-large-block.wgsl.jinja",
+          "derive": { "sharedSize": "pow2ceil(firstCandidates)", "stage": "\"output\"" },
+          "bindings": ["candidateVals_4", "candidateIdxs_2", "values", "indices", "params_9"],
+          "dispatch": { "x": "min(outputPositions, 65535)", "y": "ceilDiv(outputPositions, 65535)", "z": 1 }
         }
       ]
     },
       "id": "last_axis_large_two_merge_int",
       "priority": 10,
       "when": ["baseContract", "lastAxis", "args.k > 1", "axisDim > tunables.AXIS_SHARED_MAX_VALUES", "firstCandidates > mergeChunkSize", "secondCandidates <= mergeChunkSize", "firstScratchFits", "secondScratchFits", "largeSharedFits", "largeBlockDispatchFits", "secondMergeDispatchFits", "outputDispatchFits", "not floatInput"],
+      "derive": {
+        "scalar": "dtypes.T",
+        "usesF16": "dtypes.T == \"f16\"",
+        "blockSize": "largeBlockSize",
+        "workgroupSize": "workgroupSizePlan",
+        "floatOrderKeys": false
+      },
       "intermediates": [
         { "id": "candidateVals", "dtype": "uint32", "shape": "[firstScratchElements]" },
         { "id": "candidateIdxs", "dtype": "uint32", "shape": "[firstScratchElements]" },
         {
           "id": "block",
           "name": "TopK.LargeBlock",
+          "shader": "topk-large-block.wgsl.jinja",
+          "derive": { "stage": "\"block\"" },
+          "bindings": ["x", "candidateVals_3", "candidateIdxs", "params_7"],
+          "dispatch": {
+            "x": "min(outputPositions * largeBlockCount, 65535)",
+            "y": "ceilDiv(outputPositions * largeBlockCount, 65535)",
+            "z": 1
+          }
         },
         {
           "id": "merge_scratch",
           "name": "TopK.LargeMergeScratch",
+          "shader": "topk-large-block.wgsl.jinja",
+          "derive": { "sharedSize": "mergeChunkSize", "stage": "\"scratch\"" },
+          "bindings": ["candidateVals_4", "candidateIdxs_2", "candidateVals2_2", "candidateIdxs2", "params_11"],
+          "dispatch": {
+            "x": "min(outputPositions * secondGroups, 65535)",
+            "y": "ceilDiv(outputPositions * secondGroups, 65535)",
+            "z": 1
+          }
         },
         {
           "id": "merge",
           "name": "TopK.LargeMergeOutput",
+          "shader": "topk-large-block.wgsl.jinja",
+          "derive": { "sharedSize": "pow2ceil(secondCandidates)", "stage": "\"output\"" },
+          "bindings": ["candidateVals_6", "candidateIdxs_3", "values", "indices", "params_12"],
+          "dispatch": { "x": "min(outputPositions, 65535)", "y": "ceilDiv(outputPositions, 65535)", "z": 1 }
         }
       ]
     },
       "id": "axis_large_one_merge",
       "priority": 10,
       "when": ["baseContract", "not lastAxis", "args.k >= 1", "axisDim > tunables.AXIS_SHARED_MAX_VALUES", "firstCandidates <= mergeChunkSize", "firstScratchFits", "largeSharedFits", "largeBlockDispatchFits", "outputDispatchFits", "floatInput"],
+      "derive": {
+        "scalar": "dtypes.T",
+        "usesF16": "dtypes.T == \"f16\"",
+        "blockSize": "largeBlockSize",
+        "workgroupSize": "workgroupSizePlan",
+        "floatOrderKeys": true
+      },
       "intermediates": [
+        { "id": "candidateVals", "dtype": "uint32", "shape": "[firstScratchElements]" },
         { "id": "candidateIdxs", "dtype": "uint32", "shape": "[firstScratchElements]" }
       ],
       "passes": [
         {
           "id": "block",
           "name": "TopK.LargeBlock",
+          "shader": "topk-large-block.wgsl.jinja",
+          "derive": { "stage": "\"block\"", "strided": true },
+          "bindings": ["x", "candidateVals_3", "candidateIdxs", "params_8"],
+          "dispatch": {
+            "x": "min(outputPositions * largeBlockCount, 65535)",
+            "y": "ceilDiv(outputPositions * largeBlockCount, 65535)",
+            "z": 1
+          }
         },
         {
           "id": "merge",
           "name": "TopK.LargeMergeOutput",
+          "shader": "topk-large-block.wgsl.jinja",
+          "derive": { "sharedSize": "pow2ceil(firstCandidates)", "stage": "\"output\"", "strided": true },
+          "bindings": ["candidateVals_4", "candidateIdxs_2", "values", "indices", "params_10_source", "x"],
+          "dispatch": { "x": "min(outputPositions, 65535)", "y": "ceilDiv(outputPositions, 65535)", "z": 1 }
         }
       ]
     },
       "id": "axis_large_two_merge",
       "priority": 10,
       "when": ["baseContract", "not lastAxis", "args.k >= 1", "axisDim > tunables.AXIS_SHARED_MAX_VALUES", "firstCandidates > mergeChunkSize", "secondCandidates <= mergeChunkSize", "firstScratchFits", "secondScratchFits", "largeSharedFits", "largeBlockDispatchFits", "secondMergeDispatchFits", "outputDispatchFits", "floatInput"],
+      "derive": {
+        "scalar": "dtypes.T",
+        "usesF16": "dtypes.T == \"f16\"",
+        "blockSize": "largeBlockSize",
+        "workgroupSize": "workgroupSizePlan",
+        "floatOrderKeys": true
+      },
       "intermediates": [
+        { "id": "candidateVals", "dtype": "uint32", "shape": "[firstScratchElements]" },
         { "id": "candidateIdxs", "dtype": "uint32", "shape": "[firstScratchElements]" },
+        { "id": "candidateVals2", "dtype": "uint32", "shape": "[secondScratchElements]" },
         { "id": "candidateIdxs2", "dtype": "uint32", "shape": "[secondScratchElements]" }
       ],
       "passes": [
         {
           "id": "block",
           "name": "TopK.LargeBlock",
+          "shader": "topk-large-block.wgsl.jinja",
+          "derive": { "stage": "\"block\"", "strided": true },
+          "bindings": ["x", "candidateVals_3", "candidateIdxs", "params_8"],
+          "dispatch": {
+            "x": "min(outputPositions * largeBlockCount, 65535)",
+            "y": "ceilDiv(outputPositions * largeBlockCount, 65535)",
+            "z": 1
+          }
         },
         {
           "id": "merge_scratch",
           "name": "TopK.LargeMergeScratch",
+          "shader": "topk-large-block.wgsl.jinja",
+          "derive": { "sharedSize": "mergeChunkSize", "stage": "\"scratch\"" },
+          "bindings": ["candidateVals_4", "candidateIdxs_2", "candidateVals2_2", "candidateIdxs2", "params_11"],
+          "dispatch": {
+            "x": "min(outputPositions * secondGroups, 65535)",
+            "y": "ceilDiv(outputPositions * secondGroups, 65535)",
+            "z": 1
+          }
         },
         {
           "id": "merge",
           "name": "TopK.LargeMergeOutput",
+          "shader": "topk-large-block.wgsl.jinja",
+          "derive": { "sharedSize": "pow2ceil(secondCandidates)", "stage": "\"output\"", "strided": true },
+          "bindings": ["candidateVals_6", "candidateIdxs_3", "values", "indices", "params_13_source", "x"],
+          "dispatch": { "x": "min(outputPositions, 65535)", "y": "ceilDiv(outputPositions, 65535)", "z": 1 }
         }
       ]
     },
       "id": "last_axis_large_three_merge",
       "priority": 10,
       "when": ["baseContract", "lastAxis", "args.k > 1", "axisDim > tunables.AXIS_SHARED_MAX_VALUES", "firstCandidates > mergeChunkSize", "secondCandidates > mergeChunkSize", "thirdCandidates <= mergeChunkSize", "firstScratchFits", "secondScratchFits", "thirdScratchFits", "largeSharedFits", "largeBlockDispatchFits", "secondMergeDispatchFits", "thirdMergeDispatchFits", "outputDispatchFits", "floatInput"],
+      "derive": {
+        "scalar": "dtypes.T",
+        "usesF16": "dtypes.T == \"f16\"",
+        "blockSize": "largeBlockSize",
+        "workgroupSize": "workgroupSizePlan",
+        "floatOrderKeys": true
+      },
       "intermediates": [
+        { "id": "candidateVals", "dtype": "uint32", "shape": "[firstScratchElements]" },
         { "id": "candidateIdxs", "dtype": "uint32", "shape": "[firstScratchElements]" },
+        { "id": "candidateVals2", "dtype": "uint32", "shape": "[secondScratchElements]" },
         { "id": "candidateIdxs2", "dtype": "uint32", "shape": "[secondScratchElements]" },
+        { "id": "candidateVals3", "dtype": "uint32", "shape": "[thirdScratchElements]" },
         { "id": "candidateIdxs3", "dtype": "uint32", "shape": "[thirdScratchElements]" }
       ],
       "passes": [
         {
           "id": "block",
           "name": "TopK.LargeBlock",
+          "shader": "topk-large-block.wgsl.jinja",
+          "derive": { "stage": "\"block\"" },
+          "bindings": ["x", "candidateVals_3", "candidateIdxs", "params_7"],
+          "dispatch": {
+            "x": "min(outputPositions * largeBlockCount, 65535)",
+            "y": "ceilDiv(outputPositions * largeBlockCount, 65535)",
+            "z": 1
+          }
         },
         {
           "id": "merge_scratch",
           "name": "TopK.LargeMergeScratch",
+          "shader": "topk-large-block.wgsl.jinja",
+          "derive": { "sharedSize": "mergeChunkSize", "stage": "\"scratch\"" },
+          "bindings": ["candidateVals_4", "candidateIdxs_2", "candidateVals2_2", "candidateIdxs2", "params_11"],
+          "dispatch": {
+            "x": "min(outputPositions * secondGroups, 65535)",
+            "y": "ceilDiv(outputPositions * secondGroups, 65535)",
+            "z": 1
+          }
         },
         {
           "id": "merge_scratch2",
           "name": "TopK.LargeMergeScratch2",
+          "shader": "topk-large-block.wgsl.jinja",
+          "derive": { "sharedSize": "mergeChunkSize", "stage": "\"scratch\"" },
+          "bindings": ["candidateVals_6", "candidateIdxs_3", "candidateVals2_4", "candidateIdxs2_2", "params_14"],
+          "dispatch": {
+            "x": "min(outputPositions * thirdGroups, 65535)",
+            "y": "ceilDiv(outputPositions * thirdGroups, 65535)",
+            "z": 1
+          }
         },
         {
           "id": "merge",
           "name": "TopK.LargeMergeOutput",
+          "shader": "topk-large-block.wgsl.jinja",
+          "derive": { "sharedSize": "pow2ceil(thirdCandidates)", "stage": "\"output\"" },
+          "bindings": ["candidateVals_8", "candidateIdxs_4", "values", "indices", "params_15_source", "x"],
+          "dispatch": { "x": "min(outputPositions, 65535)", "y": "ceilDiv(outputPositions, 65535)", "z": 1 }
         }
       ]
     },
       "id": "last_axis_large_three_merge_int",
       "priority": 10,
       "when": ["baseContract", "lastAxis", "args.k > 1", "axisDim > tunables.AXIS_SHARED_MAX_VALUES", "firstCandidates > mergeChunkSize", "secondCandidates > mergeChunkSize", "thirdCandidates <= mergeChunkSize", "firstScratchFits", "secondScratchFits", "thirdScratchFits", "largeSharedFits", "largeBlockDispatchFits", "secondMergeDispatchFits", "thirdMergeDispatchFits", "outputDispatchFits", "not floatInput"],
+      "derive": {
+        "scalar": "dtypes.T",
+        "usesF16": "dtypes.T == \"f16\"",
+        "blockSize": "largeBlockSize",
+        "workgroupSize": "workgroupSizePlan",
+        "floatOrderKeys": false
+      },
       "intermediates": [
         { "id": "candidateVals", "dtype": "uint32", "shape": "[firstScratchElements]" },
         { "id": "candidateIdxs", "dtype": "uint32", "shape": "[firstScratchElements]" },
         {
           "id": "block",
           "name": "TopK.LargeBlock",
+          "shader": "topk-large-block.wgsl.jinja",
+          "derive": { "stage": "\"block\"" },
+          "bindings": ["x", "candidateVals_3", "candidateIdxs", "params_7"],
+          "dispatch": {
+            "x": "min(outputPositions * largeBlockCount, 65535)",
+            "y": "ceilDiv(outputPositions * largeBlockCount, 65535)",
+            "z": 1
+          }
         },
         {
           "id": "merge_scratch",
           "name": "TopK.LargeMergeScratch",
+          "shader": "topk-large-block.wgsl.jinja",
+          "derive": { "sharedSize": "mergeChunkSize", "stage": "\"scratch\"" },
+          "bindings": ["candidateVals_4", "candidateIdxs_2", "candidateVals2_2", "candidateIdxs2", "params_11"],
+          "dispatch": {
+            "x": "min(outputPositions * secondGroups, 65535)",
+            "y": "ceilDiv(outputPositions * secondGroups, 65535)",
+            "z": 1
+          }
         },
         {
           "id": "merge_scratch2",
           "name": "TopK.LargeMergeScratch2",
+          "shader": "topk-large-block.wgsl.jinja",
+          "derive": { "sharedSize": "mergeChunkSize", "stage": "\"scratch\"" },
+          "bindings": ["candidateVals_6", "candidateIdxs_3", "candidateVals2_4", "candidateIdxs2_2", "params_14"],
+          "dispatch": {
+            "x": "min(outputPositions * thirdGroups, 65535)",
+            "y": "ceilDiv(outputPositions * thirdGroups, 65535)",
+            "z": 1
+          }
         },
         {
           "id": "merge",
           "name": "TopK.LargeMergeOutput",
+          "shader": "topk-large-block.wgsl.jinja",
+          "derive": { "sharedSize": "pow2ceil(thirdCandidates)", "stage": "\"output\"" },
+          "bindings": ["candidateVals_8", "candidateIdxs_4", "values", "indices", "params_15"],
+          "dispatch": { "x": "min(outputPositions, 65535)", "y": "ceilDiv(outputPositions, 65535)", "z": 1 }
         }
       ]
     }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,25 +1,46 @@
 {
   "name": "ai.onnx.TopK",
-  "id": "_ai_onnx_topk_webgpu_148201a",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "mfHtgd+qPu1FdwN2IYJL8z9saLYe1dxUbiQ4utLBkxg=",
-      "manifest.json": "4FlmxzZXEt5u1iom1e5ehHAik/1Nn0aIhyIU/6dai1o=",
-      "test.json": "WrNZJCAriH+P2DxOGsAbrS4beAQMOiOilMoGY7shG5g=",
-      "topk-axis.wgsl.jinja": "lSuHRRplKZnODO49Op7ehqJiOqGmI5F/dahAsauzBoM=",
-      "topk-large-block.wgsl.jinja": "ZF/+hgaZ42J/Gun/JOXOqRzbwn6xbs2ekew8F4sOndk=",
       "topk-noop.wgsl.jinja": "rzwkvcZd8vJ3qLSosuI96CfUW+WUQyFxjc2Hgvq1bdo=",
-      "topk-portable-rows-smallk.wgsl.jinja": "jc8DhH2f7LZMHBFPqON9zet+kLnMb/fNVHBkYdAyzAo=",
-      "topk-small-rows-batched.wgsl.jinja": "7cPf2k8yNMXg1bZrusvl71RIcBSjOLKPBATaj1Y4bfo=",
-      "topk-strided-smallk.wgsl.jinja": "10odKZ0VhoUSwcQZIP3rnTEfN8gfP7OUQOilJEIfsNU=",
-      "topk-subgroup-rows.wgsl.jinja": "zO4KDfA5eM9COJ9chF7wwvYdJbpPKcKkB0qQJLeYMvA=",
-      "topk-top1-last-axis.wgsl.jinja": "C+DVDZ78X/DimUDKyRtlC5HV/nRwiXQil8IjGjXcnB0="
     }
   },
-  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
-  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.TopK" }
 }

 {
   "name": "ai.onnx.TopK",
+  "id": "_ai_onnx_topk_webgpu_52be754",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "PurEWvvW6gbYNkFFlfRDPuV5w56IJxPpwT0HV0cXd0A=",
+      "manifest.json": "YqipHsTNrefVMyk2lSH1Vp9eXwjmXHibQ+FXiPhg3iw=",
+      "test.json": "Fr5n36m9d6NTCaetFemmXK+rj/6dXPaw6INsNyX9jqs=",
+      "topk-axis.wgsl.jinja": "Xl+bV55WupVkNPIof1xc0YfVFd0YlmTvWBoVKXuKH0s=",
+      "topk-large-block.wgsl.jinja": "r0f8d02UPzDntSxjLxFsP7ix3osvWFRmzsyt6YQS5K8=",
       "topk-noop.wgsl.jinja": "rzwkvcZd8vJ3qLSosuI96CfUW+WUQyFxjc2Hgvq1bdo=",
+      "topk-portable-rows-smallk.wgsl.jinja": "WLo2yQv5mq1yMAW6vNAFbCSzRwuQ8fvM3MtuYH/d5PM=",
+      "topk-small-rows-batched.wgsl.jinja": "zaHFjRfth6UmWbFEYmjRbE2hpIdjuqyafituQUBTGvg=",
+      "topk-strided-smallk.wgsl.jinja": "pMdoMVSn+bYSKpP/FHMFLL2hnGeMQyakRPPDTITae4I=",
+      "topk-subgroup-rows.wgsl.jinja": "NPnHo98cYexGK/mRtrzxZvXOjmVF/1soDP/EaBRaApw=",
+      "topk-top1-last-axis.wgsl.jinja": "AYmJPcushgFs3ixDK6sJtDQDNEGbBk/CAAPWRMF3g+Y="
     }
   },
+  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.0",
+    "variants": {
+      "axis_smallk_tournament": ["topk-strided-smallk.wgsl.jinja"],
+      "last_axis_large_top1": ["topk-top1-last-axis.wgsl.jinja"],
+      "last_axis_top1": ["topk-top1-last-axis.wgsl.jinja"],
+      "axis_bitonic": ["topk-axis.wgsl.jinja"],
+      "top0_noop": ["topk-noop.wgsl.jinja"],
+      "subgroup_rows_smallk": ["topk-subgroup-rows.wgsl.jinja"],
+      "subgroup_min_rows_smallk": ["topk-subgroup-rows.wgsl.jinja"],
+      "portable_rows_smallk": ["topk-portable-rows-smallk.wgsl.jinja"],
+      "small_rows_batched": ["topk-small-rows-batched.wgsl.jinja"],
+      "last_axis_large_one_merge": ["topk-large-block.wgsl.jinja"],
+      "last_axis_large_two_merge": ["topk-large-block.wgsl.jinja"],
+      "last_axis_large_one_merge_int": ["topk-large-block.wgsl.jinja"],
+      "last_axis_large_two_merge_int": ["topk-large-block.wgsl.jinja"],
+      "axis_large_one_merge": ["topk-large-block.wgsl.jinja"],
+      "axis_large_two_merge": ["topk-large-block.wgsl.jinja"],
+      "last_axis_large_three_merge": ["topk-large-block.wgsl.jinja"],
+      "last_axis_large_three_merge_int": ["topk-large-block.wgsl.jinja"]
+    }
+  }
 }

build/webgpu/test.json CHANGED Viewed

@@ -1,12 +1,15 @@
 {
-  "op": "ai.onnx.TopK",
   "cases": [
     {
       "name": "onnx_default_last_axis_omitted",
       "provenance": {
         "source": "https://onnx.ai/onnx/operators/onnx__TopK.html",
         "test": "TopK axis default",
-        "notes": "Omitting axis must select the last dimension (ONNX default -1), not axis 1. Rank 3 makes those axes distinct, and both output values and indices are pinned. ONNX int64 indices use the framework's representable uint32 slot."
       },
       "args": { "k": 1 },
       "inputs": {
@@ -53,7 +56,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/math/topk_op_test.cc",
         "test": "TopKOperator.Top1DefaultAxisLargestElements",
-        "notes": "Axis-0 companion for finite subnormal ordering: a positive subnormal is strictly greater than zero; TopK should return its index rather than tie-breaking on a flushed zero."
       },
       "attrs": { "axis": 0, "largest": 1, "sorted": 1 },
       "args": { "k": 1 },
@@ -87,7 +90,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/math/topk_op_test.cc",
         "test": "TopKOperator.Top3ExplicitAxisSmallestElements",
-        "notes": "Axis-0 companion for finite subnormal ordering: a negative subnormal is strictly less than zero; smallest TopK should return its index."
       },
       "attrs": { "axis": 0, "largest": 0, "sorted": 1 },
       "args": { "k": 1 },
@@ -107,7 +110,7 @@
         "x": {
           "dtype": "float32",
           "shape": [2, 5],
-          "data": { "kind": "values", "values": [1.5, 2.0, 2.0, -1.0, 0.5, 0.1, 3.0, 2.5, 3.0, -2.0] }
         }
       },
       "outputs": {
@@ -403,7 +406,7 @@
         "x": {
           "dtype": "float32",
           "shape": [2, 5],
-          "data": { "kind": "values", "values": [1.0, 5.0, 3.0, 4.0, 2.0, 2.0, 8.0, 7.0, 6.0, 1.0] }
         }
       },
       "outputs": {
@@ -1099,7 +1102,7 @@
       },
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_top_k",
-        "notes": "ONNX TopK K input represented as framework scalar arg k. ONNX int64 tensors use framework int32/uint32 slots where representable."
       }
     },
     {
@@ -1119,7 +1122,7 @@
       },
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_top_k_negative_axis",
-        "notes": "ONNX TopK K input represented as framework scalar arg k. ONNX int64 tensors use framework int32/uint32 slots where representable."
       }
     },
     {
@@ -1139,7 +1142,7 @@
       },
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_top_k_smallest",
-        "notes": "ONNX TopK K input represented as framework scalar arg k. ONNX int64 tensors use framework int32/uint32 slots where representable."
       }
     },
     {
@@ -1166,7 +1169,7 @@
       "name": "onnx_backend_top_k_same_values_largest",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_top_k_same_values_largest",
-        "notes": "Official value tensors use int64; this fixture adapts them to supported float32 values and uint32 indices. The test_top_k_same_values vector projects to this same explicit largest=1 request."
       },
       "attrs": { "axis": 0, "largest": 1, "sorted": 1 },
       "args": { "k": 3 },
@@ -1365,7 +1368,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/math/topk_op_test.cc",
         "test": "TopKOperator.Top1ExplicitAxisLargestElements",
-        "notes": "Non-last-axis companion for exact int32 ordering: the axis sorter must compare and emit integer values without routing them through f32."
       },
       "attrs": { "axis": 0, "largest": 1, "sorted": 1 },
       "args": { "k": 1 },
@@ -1396,7 +1399,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/math/topk_op_test.cc",
         "test": "TopKOperator.Top3ExplicitAxisSmallestElements",
-        "notes": "Non-last-axis signed companion for exact int32 ordering below -2^24; f32 round-trip collapses the chosen sentinel values."
       },
       "attrs": { "axis": 0, "largest": 0, "sorted": 1 },
       "args": { "k": 1 },
@@ -1458,7 +1461,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/math/topk_op_test.cc",
         "test": "TopKOperator.Top2ExplicitAxisMultiDInputSmallestElements",
-        "notes": "Sorted k=2 signed-smallest companion for exact int32 values below -2^24."
       },
       "attrs": { "axis": -1, "largest": 0, "sorted": 1 },
       "args": { "k": 2 },
@@ -1520,7 +1523,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/math/topk_op_test.cc",
         "test": "TopKOperator.TopK_Int32_ExplicitAxis",
-        "notes": "Middle-axis signed-smallest companion for exact int32 ordering below -2^24."
       },
       "attrs": { "axis": 1, "largest": 0, "sorted": 1 },
       "args": { "k": 1 },
@@ -1565,7 +1568,7 @@
     {
       "name": "dispatch_cliff_rows_65537",
       "provenance": {
-        "notes": "65537 rows > 65535 forces the folded primary-axis dispatch (x=65535, y=2) in topk-axis.wgsl.jinja. Validates row = workgroup_id.x + workgroup_id.y*nwg.x. The 11-value cycle shifts each row (8 mod 11) so rows differ and a wrong fold would mismatch; 8 distinct values per row avoid tie-break ambiguity. Before the 2D fold the plan-time dispatch-limit guard throws."
       },
       "attrs": { "axis": -1, "largest": 1, "sorted": 1 },
       "args": { "k": 4 },
@@ -1584,7 +1587,7 @@
     {
       "name": "dispatch_cliff_rank3_axis1_cols_65537",
       "provenance": {
-        "notes": "rank3 axis1 top-k with dim2=65537 (>65535) forces the x->y dispatch fold (x=65535, y=2) while z carries the batch axis (dim0). Validates the col = wg.x + wg.y*nwg.x reconstruction in topk-axis.wgsl.jinja (axisLayout 2), guarded by col >= params.dim2. Before the fold the plan-time dispatch-limit guard throws."
       },
       "attrs": { "axis": 1, "largest": 1, "sorted": 1 },
       "args": { "k": 1 },
@@ -1707,7 +1710,7 @@
         "x": {
           "dtype": "float32",
           "shape": [1025, 128],
-          "data": { "kind": "cycle", "values": [0.1, 0.5, 0.9, 0.3, 0.7, 0.2, 0.6, 0.4, 0.8, 0.05, 0.95] }
         }
       },
       "outputs": {
@@ -1718,7 +1721,7 @@
     {
       "name": "small_row_batched_smallest_ties_4096x8_k4",
       "provenance": {
-        "notes": "The consolidated batched-row variant also renders largest=0. Equal values must retain lower input indices in ascending index order."
       },
       "attrs": { "axis": -1, "largest": 0, "sorted": 1 },
       "args": { "k": 4 },
@@ -1769,7 +1772,7 @@
     {
       "name": "rank3_axis1_2049x_k8_tournament_coverage",
       "provenance": {
-        "notes": "Compact coverage lock one element beyond the historical 2048 non-last-axis ceiling. The tournament specializes to nine local candidates per lane and covers the shape without a generic-bitonic fallback."
       },
       "attrs": { "axis": 1, "largest": 1, "sorted": 1 },
       "args": { "k": 8 },
@@ -1826,7 +1829,7 @@
       "provenance": {
         "source": "onnxruntime/core/providers/webgpu/math/top_k.cc",
         "test": "float-typed kernel registration",
-        "notes": "ORT restricts WebGPU TopK to float types; our large-axis merge kernels accept i32/u32, so the candidate scratch must preserve integer exactness (values above 2^24 and u32 above 2^31 are not f32-representable)."
       },
       "attrs": { "axis": -1, "largest": 1, "sorted": 1 },
       "args": { "k": 2 },
@@ -1857,7 +1860,7 @@
       "provenance": {
         "source": "onnxruntime/core/providers/webgpu/math/top_k.cc",
         "test": "float-typed kernel registration",
-        "notes": "ORT restricts WebGPU TopK to float types; our large-axis merge kernels accept i32/u32, so the candidate scratch must preserve integer exactness (values above 2^24 and u32 above 2^31 are not f32-representable)."
       },
       "attrs": { "axis": -1, "largest": 0, "sorted": 1 },
       "args": { "k": 2 },
@@ -1888,7 +1891,7 @@
       "provenance": {
         "source": "onnxruntime/core/providers/webgpu/math/top_k.cc",
         "test": "float-typed kernel registration",
-        "notes": "ORT restricts WebGPU TopK to float types; our large-axis merge kernels accept i32/u32, so the candidate scratch must preserve integer exactness (values above 2^24 and u32 above 2^31 are not f32-representable)."
       },
       "attrs": { "axis": -1, "largest": 1, "sorted": 1 },
       "args": { "k": 3 },
@@ -1922,7 +1925,7 @@
       "provenance": {
         "source": "onnxruntime/core/providers/webgpu/math/top_k.cc",
         "test": "float-typed kernel registration",
-        "notes": "ORT restricts WebGPU TopK to float types; our large-axis merge kernels accept i32/u32, so the candidate scratch must preserve integer exactness (values above 2^24 and u32 above 2^31 are not f32-representable)."
       },
       "attrs": { "axis": -1, "largest": 1, "sorted": 1 },
       "args": { "k": 64 },
@@ -2040,7 +2043,7 @@
     {
       "name": "two_merge_distinct_values_vocab_k16_262144",
       "provenance": {
-        "notes": "Distinct (non-tie) data through the two-merge chain at vocab scale. Guards the scratch-binding semantic redirect: with name-first resolution the final merge silently re-read the stage-1 candidate buffer, which every all-ties fixture masked."
       },
       "attrs": { "axis": -1, "largest": 1, "sorted": 1 },
       "args": { "k": 16 },
@@ -2059,7 +2062,7 @@
     {
       "name": "two_merge_monotonic_ramp_k128_8704",
       "provenance": {
-        "notes": "Monotonic ramp: the true top-128 is the last 128 indices, so any stage reading a stale candidate buffer surfaces immediately (the name-first scratch lookup returned block-1 candidates here)."
       },
       "attrs": { "axis": -1, "largest": 1, "sorted": 1 },
       "args": { "k": 128 },
@@ -2109,7 +2112,7 @@
       "name": "dsa_index_topk_2048_decode_16k",
       "requires": { "limits": { "maxComputeWorkgroupStorageSize": 32768 } },
       "provenance": {
-        "notes": "DeepSeek-V3.2 / AXK2 sparse-attention select step: index_topk defaults to 2048 in both configs. One query row against a 16K indexer key cache is the decode shape that co.huggingface.SparseAttentionScore feeds. The merge chunk is derived from k rather than fixed at MERGE_CHUNK_SIZE: a chunk no larger than k is a fixed point (each stage re-emits ceilDiv(candidates, chunk) * k), so the ladder never converged at k = 2048. Twice k halves the candidates per stage. That chunk needs 32 KiB of workgroup storage, which the case declares."
       },
       "attrs": { "axis": -1, "largest": 1, "sorted": 1 },
       "args": { "k": 2048 },
@@ -2129,7 +2132,7 @@
       "name": "dsa_index_topk_2048_prefill_rows",
       "requires": { "limits": { "maxComputeWorkgroupStorageSize": 32768 } },
       "provenance": {
-        "notes": "Prefill form of the same select step: several query rows, each choosing 2048 of 4096 candidate keys. The merge chunk is derived from k rather than fixed at MERGE_CHUNK_SIZE: a chunk no larger than k is a fixed point (each stage re-emits ceilDiv(candidates, chunk) * k), so the ladder never converged at k = 2048. Twice k halves the candidates per stage. That chunk needs 32 KiB of workgroup storage, which the case declares."
       },
       "attrs": { "axis": -1, "largest": 1, "sorted": 1 },
       "args": { "k": 2048 },
@@ -2148,7 +2151,7 @@
     {
       "name": "small_rows_batched_4096x32_k4_f16",
       "provenance": {
-        "notes": "float16 twin for the batched small-rows rung: at least SMALL_ROWS_MIN_ROWS (4096) rows, an axis within SMALL_ROWS_MAX_AXIS (64) and k within SMALL_ROWS_MAX_K (8). The 33-value cycle is one longer than the axis, so every row is a rotation of the previous one and no row repeats a value, which leaves the top-4 indices unambiguous. Each value is a multiple of 1/64 and so is exact in float16."
       },
       "attrs": { "axis": -1, "largest": 1, "sorted": 1 },
       "args": { "k": 4 },
@@ -2156,10 +2159,7 @@
         "x": {
           "dtype": "float16",
           "shape": [4096, 32],
-          "data": {
-            "kind": "cycle",
-            "values": [0.015625, 0.03125, 0.046875, 0.0625, 0.078125, 0.09375, 0.109375, 0.125, 0.140625, 0.15625, 0.171875, 0.1875, 0.203125, 0.21875, 0.234375, 0.25, 0.265625, 0.28125, 0.296875, 0.3125, 0.328125, 0.34375, 0.359375, 0.375, 0.390625, 0.40625, 0.421875, 0.4375, 0.453125, 0.46875, 0.484375, 0.5, 0.515625]
-          }
         }
       },
       "outputs": {
@@ -2170,7 +2170,7 @@
     {
       "name": "portable_rows_smallk_1025x128_k4_f16",
       "provenance": {
-        "notes": "float16 twin for the portable small-k row rung, which takes an axis above SMALL_ROWS_MAX_AXIS (64) and within PORTABLE_ROWS_MAX_AXIS (256), and more than PORTABLE_ROWS_MIN (1024) rows. The 129-value cycle is one longer than the axis so each row is a rotation with no repeated value, and every value is a multiple of 1/256 and exact in float16."
       },
       "attrs": { "axis": -1, "largest": 1, "sorted": 1 },
       "args": { "k": 4 },
@@ -2188,6 +2188,588 @@
         "values": { "dtype": "float16", "shape": [1025, 4], "tolerance": 0 },
         "indices": { "dtype": "uint32", "shape": [1025, 4], "tolerance": 0 }
       }
     }
   ]
 }

 {
+  "fixtureArrays": {
+    "topk_small_rows_rank3_cycle33": [0.015625, 0.03125, 0.046875, 0.0625, 0.078125, 0.09375, 0.109375, 0.125, 0.140625, 0.15625, 0.171875, 0.1875, 0.203125, 0.21875, 0.234375, 0.25, 0.265625, 0.28125, 0.296875, 0.3125, 0.328125, 0.34375, 0.359375, 0.375, 0.390625, 0.40625, 0.421875, 0.4375, 0.453125, 0.46875, 0.484375, 0.5, 0.515625],
+    "order_key_extremes": ["-Infinity", -65504, -2, 0, 0, 2, 65504, "Infinity", 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
+  },
   "cases": [
     {
       "name": "onnx_default_last_axis_omitted",
       "provenance": {
         "source": "https://onnx.ai/onnx/operators/onnx__TopK.html",
         "test": "TopK axis default",
+        "notes": "Omitting axis must select the last dimension (ONNX default -1), not axis 1. Rank 3 makes those axes distinct, and both output values and indices are pinned. Representable ONNX int64 indices are stored in uint32 slots."
       },
       "args": { "k": 1 },
       "inputs": {
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/math/topk_op_test.cc",
         "test": "TopKOperator.Top1DefaultAxisLargestElements",
+        "notes": "On axis 0, a positive subnormal is strictly greater than zero; TopK must return its index rather than tie-break against a flushed zero."
       },
       "attrs": { "axis": 0, "largest": 1, "sorted": 1 },
       "args": { "k": 1 },
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/math/topk_op_test.cc",
         "test": "TopKOperator.Top3ExplicitAxisSmallestElements",
+        "notes": "On axis 0, a negative subnormal is strictly less than zero; smallest TopK must return its index."
       },
       "attrs": { "axis": 0, "largest": 0, "sorted": 1 },
       "args": { "k": 1 },
         "x": {
           "dtype": "float32",
           "shape": [2, 5],
+          "data": { "kind": "values", "values": [1.5, 2.0, 2.0, -1.0, 0.5, 0.1, 2.5, 3.0, 3.0, -2.0] }
         }
       },
       "outputs": {
         "x": {
           "dtype": "float32",
           "shape": [2, 5],
+          "data": { "kind": "values", "values": [1.0, 5.0, 3.0, 4.0, 2.0, 2.0, 7.0, 8.0, 6.0, 1.0] }
         }
       },
       "outputs": {
       },
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_top_k",
+        "notes": "The fixture represents the scalar TopK `K` input as argument `k`. This WebGPU package stores representable ONNX int64 values and indices in int32 or uint32 slots."
       }
     },
     {
       },
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_top_k_negative_axis",
+        "notes": "The fixture represents the scalar TopK `K` input as argument `k`. This WebGPU package stores representable ONNX int64 values and indices in int32 or uint32 slots."
       }
     },
     {
       },
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_top_k_smallest",
+        "notes": "The fixture represents the scalar TopK `K` input as argument `k`. This WebGPU package stores representable ONNX int64 values and indices in int32 or uint32 slots."
       }
     },
     {
       "name": "onnx_backend_top_k_same_values_largest",
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_top_k_same_values_largest",
+        "notes": "The official int64 value tensor is represented as supported float32 values with uint32 indices and an explicit `largest=1` request."
       },
       "attrs": { "axis": 0, "largest": 1, "sorted": 1 },
       "args": { "k": 3 },
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/math/topk_op_test.cc",
         "test": "TopKOperator.Top1ExplicitAxisLargestElements",
+        "notes": "An axis-0 TopK must compare and emit exact int32 values without routing them through f32."
       },
       "attrs": { "axis": 0, "largest": 1, "sorted": 1 },
       "args": { "k": 1 },
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/math/topk_op_test.cc",
         "test": "TopKOperator.Top3ExplicitAxisSmallestElements",
+        "notes": "An axis-0 smallest TopK must preserve exact int32 ordering below -2^24; an f32 round trip collapses the chosen sentinel values."
       },
       "attrs": { "axis": 0, "largest": 0, "sorted": 1 },
       "args": { "k": 1 },
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/math/topk_op_test.cc",
         "test": "TopKOperator.Top2ExplicitAxisMultiDInputSmallestElements",
+        "notes": "Sorted smallest TopK with k=2 must preserve exact int32 ordering below -2^24."
       },
       "attrs": { "axis": -1, "largest": 0, "sorted": 1 },
       "args": { "k": 2 },
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/math/topk_op_test.cc",
         "test": "TopKOperator.TopK_Int32_ExplicitAxis",
+        "notes": "A rank-3 middle-axis smallest TopK must preserve exact int32 ordering below -2^24."
       },
       "attrs": { "axis": 1, "largest": 0, "sorted": 1 },
       "args": { "k": 1 },
     {
       "name": "dispatch_cliff_rows_65537",
       "provenance": {
+        "notes": "With 65,537 rows, TopK uses a two-row workgroup dispatch. An 11-value cycle shifts every eight-element row and avoids ties, making folded row-index errors observable."
       },
       "attrs": { "axis": -1, "largest": 1, "sorted": 1 },
       "args": { "k": 4 },
     {
       "name": "dispatch_cliff_rank3_axis1_cols_65537",
       "provenance": {
+        "notes": "A rank-3 axis-1 TopK with trailing extent 65,537 uses a two-row workgroup dispatch while the z dimension carries the batch. Distinct column values make folded column-index errors observable."
       },
       "attrs": { "axis": 1, "largest": 1, "sorted": 1 },
       "args": { "k": 1 },
         "x": {
           "dtype": "float32",
           "shape": [1025, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5, "srqStep": 0.02 }
         }
       },
       "outputs": {
     {
       "name": "small_row_batched_smallest_ties_4096x8_k4",
       "provenance": {
+        "notes": "For 4,096 equal eight-element rows with largest=0, sorted output must retain the lower input indices in ascending order."
       },
       "attrs": { "axis": -1, "largest": 0, "sorted": 1 },
       "args": { "k": 4 },
     {
       "name": "rank3_axis1_2049x_k8_tournament_coverage",
       "provenance": {
+        "notes": "A non-last axis of length 2,049 gives each tournament lane nine local candidates and avoids the generic bitonic implementation."
       },
       "attrs": { "axis": 1, "largest": 1, "sorted": 1 },
       "args": { "k": 8 },
       "provenance": {
         "source": "onnxruntime/core/providers/webgpu/math/top_k.cc",
         "test": "float-typed kernel registration",
+        "notes": "Large-axis merge scratch must preserve exact int32 values above 2^24 rather than store candidates as f32."
       },
       "attrs": { "axis": -1, "largest": 1, "sorted": 1 },
       "args": { "k": 2 },
       "provenance": {
         "source": "onnxruntime/core/providers/webgpu/math/top_k.cc",
         "test": "float-typed kernel registration",
+        "notes": "Large-axis merge scratch must preserve exact negative int32 values below -2^24 rather than store candidates as f32."
       },
       "attrs": { "axis": -1, "largest": 0, "sorted": 1 },
       "args": { "k": 2 },
       "provenance": {
         "source": "onnxruntime/core/providers/webgpu/math/top_k.cc",
         "test": "float-typed kernel registration",
+        "notes": "Large-axis merge scratch must preserve exact uint32 values above 2^31 rather than store candidates as f32."
       },
       "attrs": { "axis": -1, "largest": 1, "sorted": 1 },
       "args": { "k": 3 },
       "provenance": {
         "source": "onnxruntime/core/providers/webgpu/math/top_k.cc",
         "test": "float-typed kernel registration",
+        "notes": "Two merge stages must preserve exact int32 values above 2^24 rather than store candidates as f32."
       },
       "attrs": { "axis": -1, "largest": 1, "sorted": 1 },
       "args": { "k": 64 },
     {
       "name": "two_merge_distinct_values_vocab_k16_262144",
       "provenance": {
+        "notes": "Distinct values over 262,144 candidates require two merge stages. The final stage must read the immediately preceding candidate buffer; otherwise the returned top 16 values and indices change."
       },
       "attrs": { "axis": -1, "largest": 1, "sorted": 1 },
       "args": { "k": 16 },
     {
       "name": "two_merge_monotonic_ramp_k128_8704",
       "provenance": {
+        "notes": "For a monotonic ramp, the true top 128 values occupy the final 128 indices, making stale candidate-buffer reads observable at every merge stage."
       },
       "attrs": { "axis": -1, "largest": 1, "sorted": 1 },
       "args": { "k": 128 },
       "name": "dsa_index_topk_2048_decode_16k",
       "requires": { "limits": { "maxComputeWorkgroupStorageSize": 32768 } },
       "provenance": {
+        "notes": "Selecting k=2048 from 16,384 candidates requires merge chunks larger than k; a chunk no larger than k does not reduce the candidate count. A 4,096-element chunk halves the candidates per stage and requires 32 KiB of workgroup storage."
       },
       "attrs": { "axis": -1, "largest": 1, "sorted": 1 },
       "args": { "k": 2048 },
       "name": "dsa_index_topk_2048_prefill_rows",
       "requires": { "limits": { "maxComputeWorkgroupStorageSize": 32768 } },
       "provenance": {
+        "notes": "Several query rows each select 2,048 of 4,096 candidate keys. A 4,096-element merge chunk halves the candidate count at each stage and requires the declared 32 KiB of workgroup storage."
       },
       "attrs": { "axis": -1, "largest": 1, "sorted": 1 },
       "args": { "k": 2048 },
     {
       "name": "small_rows_batched_4096x32_k4_f16",
       "provenance": {
+        "notes": "A float16 tensor with 4,096 rows, axis length 32, and k=4 exercises batched small-row selection. A 33-value cycle rotates distinct multiples of 1/64 through each row, leaving exact and unambiguous top-4 indices."
       },
       "attrs": { "axis": -1, "largest": 1, "sorted": 1 },
       "args": { "k": 4 },
         "x": {
           "dtype": "float16",
           "shape": [4096, 32],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/topk_small_rows_rank3_cycle33" } }
         }
       },
       "outputs": {
     {
       "name": "portable_rows_smallk_1025x128_k4_f16",
       "provenance": {
+        "notes": "A float16 tensor with 1,025 rows, axis length 128, and k=4 exercises the portable small-k row specialization. A 129-value cycle rotates distinct, exactly representable values through each row."
       },
       "attrs": { "axis": -1, "largest": 1, "sorted": 1 },
       "args": { "k": 4 },
         "values": { "dtype": "float16", "shape": [1025, 4], "tolerance": 0 },
         "indices": { "dtype": "uint32", "shape": [1025, 4], "tolerance": 0 }
       }
+    },
+    {
+      "name": "unsorted_small_rows_4096x8_k2_batched_route",
+      "provenance": {
+        "notes": "Exactly 4,096 eight-element rows with k=2 exercise batched small-row selection when output order is unspecified. A nine-value cycle rotates eight distinct values through each row, leaving the top-2 set unambiguous."
+      },
+      "attrs": { "axis": -1, "largest": 1, "sorted": 0 },
+      "args": { "k": 2 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [4096, 8],
+          "data": { "kind": "cycle", "values": [0.1, 0.5, 0.9, 0.3, 0.7, 0.2, 0.6, 0.4, 0.8] }
+        }
+      },
+      "outputs": {
+        "values": { "dtype": "float32", "shape": [4096, 2], "tolerance": 0.000001 },
+        "indices": { "dtype": "uint32", "shape": [4096, 2], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "small_rows_batched_rank3_4096rows_32axis_k4",
+      "provenance": {
+        "notes": "A rank-3 tensor flattens to 4,096 rows of 32 values with k=4, exercising batched small-row selection independently of tensor rank. A 33-value cycle rotates distinct values through each row and leaves the top-4 indices unambiguous."
+      },
+      "attrs": { "axis": -1, "largest": 1, "sorted": 1 },
+      "args": { "k": 4 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [128, 32, 32],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/topk_small_rows_rank3_cycle33" } }
+        }
+      },
+      "outputs": {
+        "values": { "dtype": "float32", "shape": [128, 32, 4], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [128, 32, 4] }
+      }
+    },
+    {
+      "name": "bitonic_prefix_boundary_float32_axis256_k17_largest0",
+      "provenance": {
+        "notes": "Covers both sides of the default workgroup-width pruning boundary and a non-power-of-two retained prefix. Rounded adjacent ties retain stable index ordering while selected values vary across positions; integer cases preserve exact ordering beyond float32 integer precision."
+      },
+      "attrs": { "axis": -1, "largest": 0 },
+      "args": { "k": 17 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 256],
+          "data": { "kind": "linspace", "start": 16777200.0, "end": 16777600.0 }
+        }
+      },
+      "outputs": {
+        "values": { "dtype": "float32", "shape": [2, 17], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [2, 17], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "bitonic_prefix_boundary_float32_axis256_k17_largest1",
+      "provenance": {
+        "notes": "Covers both sides of the default workgroup-width pruning boundary and a non-power-of-two retained prefix. Rounded adjacent ties retain stable index ordering while selected values vary across positions; integer cases preserve exact ordering beyond float32 integer precision."
+      },
+      "attrs": { "axis": -1, "largest": 1 },
+      "args": { "k": 17 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 256],
+          "data": { "kind": "linspace", "start": 16777200.0, "end": 16777600.0 }
+        }
+      },
+      "outputs": {
+        "values": { "dtype": "float32", "shape": [2, 17], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [2, 17], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "bitonic_prefix_boundary_float32_axis257_k17_largest0",
+      "provenance": {
+        "notes": "Covers both sides of the default workgroup-width pruning boundary and a non-power-of-two retained prefix. Rounded adjacent ties retain stable index ordering while selected values vary across positions; integer cases preserve exact ordering beyond float32 integer precision."
+      },
+      "attrs": { "axis": -1, "largest": 0 },
+      "args": { "k": 17 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 257],
+          "data": { "kind": "linspace", "start": 16777200.0, "end": 16777600.0 }
+        }
+      },
+      "outputs": {
+        "values": { "dtype": "float32", "shape": [2, 17], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [2, 17], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "bitonic_prefix_boundary_float32_axis257_k17_largest1",
+      "provenance": {
+        "notes": "Covers both sides of the default workgroup-width pruning boundary and a non-power-of-two retained prefix. Rounded adjacent ties retain stable index ordering while selected values vary across positions; integer cases preserve exact ordering beyond float32 integer precision."
+      },
+      "attrs": { "axis": -1, "largest": 1 },
+      "args": { "k": 17 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 257],
+          "data": { "kind": "linspace", "start": 16777200.0, "end": 16777600.0 }
+        }
+      },
+      "outputs": {
+        "values": { "dtype": "float32", "shape": [2, 17], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [2, 17], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "bitonic_prefix_boundary_float16_axis256_k17_largest0",
+      "provenance": {
+        "notes": "Covers both sides of the default workgroup-width pruning boundary and a non-power-of-two retained prefix. Rounded adjacent ties retain stable index ordering while selected values vary across positions; integer cases preserve exact ordering beyond float32 integer precision."
+      },
+      "attrs": { "axis": -1, "largest": 0 },
+      "args": { "k": 17 },
+      "inputs": {
+        "x": { "dtype": "float16", "shape": [2, 256], "data": { "kind": "linspace", "start": 1000.0, "end": 1200.0 } }
+      },
+      "outputs": {
+        "values": { "dtype": "float16", "shape": [2, 17], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [2, 17], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "bitonic_prefix_boundary_float16_axis256_k17_largest1",
+      "provenance": {
+        "notes": "Covers both sides of the default workgroup-width pruning boundary and a non-power-of-two retained prefix. Rounded adjacent ties retain stable index ordering while selected values vary across positions; integer cases preserve exact ordering beyond float32 integer precision."
+      },
+      "attrs": { "axis": -1, "largest": 1 },
+      "args": { "k": 17 },
+      "inputs": {
+        "x": { "dtype": "float16", "shape": [2, 256], "data": { "kind": "linspace", "start": 1000.0, "end": 1200.0 } }
+      },
+      "outputs": {
+        "values": { "dtype": "float16", "shape": [2, 17], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [2, 17], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "bitonic_prefix_boundary_float16_axis257_k17_largest0",
+      "provenance": {
+        "notes": "Covers both sides of the default workgroup-width pruning boundary and a non-power-of-two retained prefix. Rounded adjacent ties retain stable index ordering while selected values vary across positions; integer cases preserve exact ordering beyond float32 integer precision."
+      },
+      "attrs": { "axis": -1, "largest": 0 },
+      "args": { "k": 17 },
+      "inputs": {
+        "x": { "dtype": "float16", "shape": [2, 257], "data": { "kind": "linspace", "start": 1000.0, "end": 1200.0 } }
+      },
+      "outputs": {
+        "values": { "dtype": "float16", "shape": [2, 17], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [2, 17], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "bitonic_prefix_boundary_float16_axis257_k17_largest1",
+      "provenance": {
+        "notes": "Covers both sides of the default workgroup-width pruning boundary and a non-power-of-two retained prefix. Rounded adjacent ties retain stable index ordering while selected values vary across positions; integer cases preserve exact ordering beyond float32 integer precision."
+      },
+      "attrs": { "axis": -1, "largest": 1 },
+      "args": { "k": 17 },
+      "inputs": {
+        "x": { "dtype": "float16", "shape": [2, 257], "data": { "kind": "linspace", "start": 1000.0, "end": 1200.0 } }
+      },
+      "outputs": {
+        "values": { "dtype": "float16", "shape": [2, 17], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [2, 17], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "bitonic_prefix_boundary_int32_axis256_k17_largest0",
+      "provenance": {
+        "notes": "Covers both sides of the default workgroup-width pruning boundary and a non-power-of-two retained prefix. Rounded adjacent ties retain stable index ordering while selected values vary across positions; integer cases preserve exact ordering beyond float32 integer precision."
+      },
+      "attrs": { "axis": -1, "largest": 0 },
+      "args": { "k": 17 },
+      "inputs": {
+        "x": {
+          "dtype": "int32",
+          "shape": [2, 256],
+          "data": { "kind": "linspace", "start": -16777473, "end": -16777216 }
+        }
+      },
+      "outputs": {
+        "values": { "dtype": "int32", "shape": [2, 17], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [2, 17], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "bitonic_prefix_boundary_int32_axis256_k17_largest1",
+      "provenance": {
+        "notes": "Covers both sides of the default workgroup-width pruning boundary and a non-power-of-two retained prefix. Rounded adjacent ties retain stable index ordering while selected values vary across positions; integer cases preserve exact ordering beyond float32 integer precision."
+      },
+      "attrs": { "axis": -1, "largest": 1 },
+      "args": { "k": 17 },
+      "inputs": {
+        "x": {
+          "dtype": "int32",
+          "shape": [2, 256],
+          "data": { "kind": "linspace", "start": -16777473, "end": -16777216 }
+        }
+      },
+      "outputs": {
+        "values": { "dtype": "int32", "shape": [2, 17], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [2, 17], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "bitonic_prefix_boundary_int32_axis257_k17_largest0",
+      "provenance": {
+        "notes": "Covers both sides of the default workgroup-width pruning boundary and a non-power-of-two retained prefix. Rounded adjacent ties retain stable index ordering while selected values vary across positions; integer cases preserve exact ordering beyond float32 integer precision."
+      },
+      "attrs": { "axis": -1, "largest": 0 },
+      "args": { "k": 17 },
+      "inputs": {
+        "x": {
+          "dtype": "int32",
+          "shape": [2, 257],
+          "data": { "kind": "linspace", "start": -16777473, "end": -16777216 }
+        }
+      },
+      "outputs": {
+        "values": { "dtype": "int32", "shape": [2, 17], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [2, 17], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "bitonic_prefix_boundary_int32_axis257_k17_largest1",
+      "provenance": {
+        "notes": "Covers both sides of the default workgroup-width pruning boundary and a non-power-of-two retained prefix. Rounded adjacent ties retain stable index ordering while selected values vary across positions; integer cases preserve exact ordering beyond float32 integer precision."
+      },
+      "attrs": { "axis": -1, "largest": 1 },
+      "args": { "k": 17 },
+      "inputs": {
+        "x": {
+          "dtype": "int32",
+          "shape": [2, 257],
+          "data": { "kind": "linspace", "start": -16777473, "end": -16777216 }
+        }
+      },
+      "outputs": {
+        "values": { "dtype": "int32", "shape": [2, 17], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [2, 17], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "bitonic_prefix_boundary_uint32_axis256_k17_largest0",
+      "provenance": {
+        "notes": "Covers both sides of the default workgroup-width pruning boundary and a non-power-of-two retained prefix. Rounded adjacent ties retain stable index ordering while selected values vary across positions; integer cases preserve exact ordering beyond float32 integer precision."
+      },
+      "attrs": { "axis": -1, "largest": 0 },
+      "args": { "k": 17 },
+      "inputs": {
+        "x": {
+          "dtype": "uint32",
+          "shape": [2, 256],
+          "data": { "kind": "linspace", "start": 4294967040, "end": 4294967295 }
+        }
+      },
+      "outputs": {
+        "values": { "dtype": "uint32", "shape": [2, 17], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [2, 17], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "bitonic_prefix_boundary_uint32_axis256_k17_largest1",
+      "provenance": {
+        "notes": "Covers both sides of the default workgroup-width pruning boundary and a non-power-of-two retained prefix. Rounded adjacent ties retain stable index ordering while selected values vary across positions; integer cases preserve exact ordering beyond float32 integer precision."
+      },
+      "attrs": { "axis": -1, "largest": 1 },
+      "args": { "k": 17 },
+      "inputs": {
+        "x": {
+          "dtype": "uint32",
+          "shape": [2, 256],
+          "data": { "kind": "linspace", "start": 4294967040, "end": 4294967295 }
+        }
+      },
+      "outputs": {
+        "values": { "dtype": "uint32", "shape": [2, 17], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [2, 17], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "bitonic_prefix_boundary_uint32_axis257_k17_largest0",
+      "provenance": {
+        "notes": "Covers both sides of the default workgroup-width pruning boundary and a non-power-of-two retained prefix. Rounded adjacent ties retain stable index ordering while selected values vary across positions; integer cases preserve exact ordering beyond float32 integer precision."
+      },
+      "attrs": { "axis": -1, "largest": 0 },
+      "args": { "k": 17 },
+      "inputs": {
+        "x": {
+          "dtype": "uint32",
+          "shape": [2, 257],
+          "data": { "kind": "linspace", "start": 4294967040, "end": 4294967295 }
+        }
+      },
+      "outputs": {
+        "values": { "dtype": "uint32", "shape": [2, 17], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [2, 17], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "bitonic_prefix_boundary_uint32_axis257_k17_largest1",
+      "provenance": {
+        "notes": "Covers both sides of the default workgroup-width pruning boundary and a non-power-of-two retained prefix. Rounded adjacent ties retain stable index ordering while selected values vary across positions; integer cases preserve exact ordering beyond float32 integer precision."
+      },
+      "attrs": { "axis": -1, "largest": 1 },
+      "args": { "k": 17 },
+      "inputs": {
+        "x": {
+          "dtype": "uint32",
+          "shape": [2, 257],
+          "data": { "kind": "linspace", "start": 4294967040, "end": 4294967295 }
+        }
+      },
+      "outputs": {
+        "values": { "dtype": "uint32", "shape": [2, 17], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [2, 17], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "order_keys_float32_n257_axis0_largest0",
+      "args": { "k": 17 },
+      "attrs": { "axis": 0, "largest": 0, "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [257, 3],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/order_key_extremes" } }
+        }
+      },
+      "outputs": {
+        "values": { "dtype": "float32", "shape": [17, 3], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [17, 3], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "order_keys_float32_n257_axis0_largest1",
+      "args": { "k": 17 },
+      "attrs": { "axis": 0, "largest": 1, "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [257, 3],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/order_key_extremes" } }
+        }
+      },
+      "outputs": {
+        "values": { "dtype": "float32", "shape": [17, 3], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [17, 3], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "order_keys_float32_n257_axis1_largest0",
+      "args": { "k": 17 },
+      "attrs": { "axis": 1, "largest": 0, "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 257],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/order_key_extremes" } }
+        }
+      },
+      "outputs": {
+        "values": { "dtype": "float32", "shape": [3, 17], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [3, 17], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "order_keys_float32_n257_axis1_largest1",
+      "args": { "k": 17 },
+      "attrs": { "axis": 1, "largest": 1, "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 257],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/order_key_extremes" } }
+        }
+      },
+      "outputs": {
+        "values": { "dtype": "float32", "shape": [3, 17], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [3, 17], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "order_keys_float32_n4099_axis0_largest0",
+      "args": { "k": 70 },
+      "attrs": { "axis": 0, "largest": 0, "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [4099, 3],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/order_key_extremes" } }
+        }
+      },
+      "outputs": {
+        "values": { "dtype": "float32", "shape": [70, 3], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [70, 3], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "order_keys_float32_n4099_axis0_largest1",
+      "args": { "k": 70 },
+      "attrs": { "axis": 0, "largest": 1, "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [4099, 3],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/order_key_extremes" } }
+        }
+      },
+      "outputs": {
+        "values": { "dtype": "float32", "shape": [70, 3], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [70, 3], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "order_keys_float32_n4099_axis1_largest0",
+      "args": { "k": 70 },
+      "attrs": { "axis": 1, "largest": 0, "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 4099],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/order_key_extremes" } }
+        }
+      },
+      "outputs": {
+        "values": { "dtype": "float32", "shape": [3, 70], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [3, 70], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "order_keys_float32_n4099_axis1_largest1",
+      "args": { "k": 70 },
+      "attrs": { "axis": 1, "largest": 1, "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 4099],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/order_key_extremes" } }
+        }
+      },
+      "outputs": {
+        "values": { "dtype": "float32", "shape": [3, 70], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [3, 70], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "order_keys_float16_n257_axis0_largest0",
+      "args": { "k": 17 },
+      "attrs": { "axis": 0, "largest": 0, "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [257, 3],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/order_key_extremes" } }
+        }
+      },
+      "outputs": {
+        "values": { "dtype": "float16", "shape": [17, 3], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [17, 3], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "order_keys_float16_n257_axis0_largest1",
+      "args": { "k": 17 },
+      "attrs": { "axis": 0, "largest": 1, "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [257, 3],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/order_key_extremes" } }
+        }
+      },
+      "outputs": {
+        "values": { "dtype": "float16", "shape": [17, 3], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [17, 3], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "order_keys_float16_n257_axis1_largest0",
+      "args": { "k": 17 },
+      "attrs": { "axis": 1, "largest": 0, "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [3, 257],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/order_key_extremes" } }
+        }
+      },
+      "outputs": {
+        "values": { "dtype": "float16", "shape": [3, 17], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [3, 17], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "order_keys_float16_n257_axis1_largest1",
+      "args": { "k": 17 },
+      "attrs": { "axis": 1, "largest": 1, "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [3, 257],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/order_key_extremes" } }
+        }
+      },
+      "outputs": {
+        "values": { "dtype": "float16", "shape": [3, 17], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [3, 17], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "order_keys_float16_n4099_axis0_largest0",
+      "args": { "k": 70 },
+      "attrs": { "axis": 0, "largest": 0, "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [4099, 3],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/order_key_extremes" } }
+        }
+      },
+      "outputs": {
+        "values": { "dtype": "float16", "shape": [70, 3], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [70, 3], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "order_keys_float16_n4099_axis0_largest1",
+      "args": { "k": 70 },
+      "attrs": { "axis": 0, "largest": 1, "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [4099, 3],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/order_key_extremes" } }
+        }
+      },
+      "outputs": {
+        "values": { "dtype": "float16", "shape": [70, 3], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [70, 3], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "order_keys_float16_n4099_axis1_largest0",
+      "args": { "k": 70 },
+      "attrs": { "axis": 1, "largest": 0, "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [3, 4099],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/order_key_extremes" } }
+        }
+      },
+      "outputs": {
+        "values": { "dtype": "float16", "shape": [3, 70], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [3, 70], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "order_keys_float16_n4099_axis1_largest1",
+      "args": { "k": 70 },
+      "attrs": { "axis": 1, "largest": 1, "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [3, 4099],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/order_key_extremes" } }
+        }
+      },
+      "outputs": {
+        "values": { "dtype": "float16", "shape": [3, 70], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [3, 70], "tolerance": 0 }
+      }
     }
   ]
 }

build/webgpu/topk-axis.wgsl.jinja CHANGED Viewed

@@ -1,34 +1,47 @@
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 // Generic (outer, axis, inner) bitonic TopK. Rank and layout differences reduce
 // to uniform geometry, so one implementation covers every tensor rank.
 {% set is_int = scalar == "i32" or scalar == "u32" %}
-{% set val_t = scalar if is_int else "f32" %}
-{% set zero = scalar ~ "(0)" if is_int else "0.0" %}
 // Stable bitonic comparator for TopK values and their original indices.
 // Stable order key and NaN predicate for direct, tournament, and bitonic TopK.
 {% if order_scalar is not defined %}{% set order_scalar = scalar %}{% endif %}
 const PAD_INDEX: u32 = 0xffffffffu;
 fn order_key(value: {{ val_t }}) -> u32 {
-{% if is_int and order_scalar == "i32" %}
   return bitcast<u32>(value) ^ 0x80000000u;
-{% elif is_int and order_scalar == "u32" %}
   return value;
 {% else %}
-  let bits0 = bitcast<u32>(value);
-  let bits = select(bits0, 0u, bits0 == 0x80000000u);
-  return select(bits | 0x80000000u, ~bits, (bits & 0x80000000u) != 0u);
 {% endif %}
 }
 {% if not is_int %}
-fn is_nan_value(value: f32) -> bool {
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u
       && (bits & 0x007fffffu) != 0u;
 }
 {% endif %}
@@ -73,9 +86,8 @@ var<workgroup> shared_indices: array<u32, {{ sharedSize }}>;
 fn main(
   @builtin(workgroup_id) workgroup: vec3<u32>,
   @builtin(local_invocation_id) local: vec3<u32>,
-  @builtin(num_workgroups) num_workgroups: vec3<u32>,
 ) {
-  let output_position = workgroup.x + workgroup.y * num_workgroups.x;
   if (output_position >= params.outputPositions) {
     return;
   }
@@ -86,7 +98,7 @@ fn main(
   for (var axis_index = tid; axis_index < SHARED_SIZE; axis_index += WG) {
     if (axis_index < params.axis) {
-      shared_values[axis_index] = {{ val_t }}(x[input_base + axis_index * params.inner]);
       shared_indices[axis_index] = axis_index;
     } else {
       shared_values[axis_index] = {{ zero }};
@@ -95,23 +107,16 @@ fn main(
   }
   workgroupBarrier();
-  // In-place bitonic ordering of values and their original indices.
-  for (var block_size = 2u; block_size <= SHARED_SIZE; block_size <<= 1u) {
-    for (var gap = block_size >> 1u; gap > 0u; gap >>= 1u) {
-      for (var pair = tid; pair < SHARED_SIZE / 2u; pair += WG) {
-        // `gap` and `block_size` are powers of two by construction (block_size
-        // doubles from 2, gap halves from block_size >> 1), so the index math is
-        // masks and shifts. GPUs have no integer divide unit and neither
-        // value is a compile-time constant here, so a literal `/` and `%` expand
-        // to a full division per pair per stage — three of them, against a body
-        // that is otherwise four shared loads and a compare.
         let position = pair & (gap - 1u);
         let left = ((pair - position) << 1u) + position;
         let right = left + gap;
 {% if attrs.largest == 0 %}
-        let ascending = (left & block_size) == 0u;
 {% else %}
-        let ascending = (left & block_size) != 0u;
 {% endif %}
         let left_value = shared_values[left];
         let right_value = shared_values[right];
@@ -124,13 +129,25 @@ fn main(
           shared_indices[right] = left_index;
         }
       }
       workgroupBarrier();
     }
   }
   for (var rank = tid; rank < params.k; rank += WG) {
     let output_index = (outer_index * params.k + rank) * params.inner + inner_index;
-    values[output_index] = {{ scalar }}(shared_values[rank]);
     indices[output_index] = shared_indices[rank];
   }
 }

 {{ env.wgsl.resourceDeclarations }}
 // Generic (outer, axis, inner) bitonic TopK. Rank and layout differences reduce
 // to uniform geometry, so one implementation covers every tensor rank.
 {% set is_int = scalar == "i32" or scalar == "u32" %}
+{% set val_t = scalar if is_int else "u32" %}
+{% set zero = scalar ~ "(0)" if is_int else "0u" %}
+{% set order_scalar = val_t %}
 // Stable bitonic comparator for TopK values and their original indices.
 // Stable order key and NaN predicate for direct, tournament, and bitonic TopK.
 {% if order_scalar is not defined %}{% set order_scalar = scalar %}{% endif %}
 const PAD_INDEX: u32 = 0xffffffffu;
+{% macro float_order_key(value) %}
+  let bits0 = bitcast<u32>({{ value }});
+  let bits = select(bits0, 0u, bits0 == 0x80000000u);
+  return select(bits | 0x80000000u, ~bits, (bits & 0x80000000u) != 0u);
+{% endmacro %}
 fn order_key(value: {{ val_t }}) -> u32 {
+{% if order_scalar == "i32" %}
   return bitcast<u32>(value) ^ 0x80000000u;
+{% elif order_scalar == "u32" %}
   return value;
 {% else %}
+{{ float_order_key("value") }}
 {% endif %}
 }
 {% if not is_int %}
+fn is_nan_value(value: {{ val_t }}) -> bool {
+{% if floatOrderKeys is defined and floatOrderKeys %}
+  // Encoded infinities bound the finite keys; either outer range is NaN.
+  return value > 0xff800000u || value < 0x007fffffu;
+{% else %}
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u
       && (bits & 0x007fffffu) != 0u;
+{% endif %}
+}
+{% endif %}
+{% if floatOrderKeys is defined and floatOrderKeys and (stage is not defined or stage == "block") %}
+// Encode once on input. Comparators consume keys, while output gathers the
+// selected original values to preserve their representation, including -0.
+fn input_order_key(value: {{ scalar }}) -> u32 {
+{{ float_order_key("f32(value)") }}
 }
 {% endif %}
 fn main(
   @builtin(workgroup_id) workgroup: vec3<u32>,
   @builtin(local_invocation_id) local: vec3<u32>,
 ) {
+  let output_position = workgroup.x + workgroup.y * {{ DISPATCH_FOLD_WIDTH }}u;
   if (output_position >= params.outputPositions) {
     return;
   }
   for (var axis_index = tid; axis_index < SHARED_SIZE; axis_index += WG) {
     if (axis_index < params.axis) {
+      shared_values[axis_index] = {{ val_t if is_int else "input_order_key" }}(x[input_base + axis_index * params.inner]);
       shared_indices[axis_index] = axis_index;
     } else {
       shared_values[axis_index] = {{ zero }};
   }
   workgroupBarrier();
+{% macro compare_pairs(pairCount, finalStage) %}
+      for (var pair = tid; pair < {{ pairCount }}; pair += WG) {
+        // Power-of-two gaps let masks and shifts identify disjoint pairs.
         let position = pair & (gap - 1u);
         let left = ((pair - position) << 1u) + position;
         let right = left + gap;
 {% if attrs.largest == 0 %}
+        let ascending = {{ "true" if finalStage else "(left & block_size) == 0u" }};
 {% else %}
+        let ascending = {{ "false" if finalStage else "(left & block_size) != 0u" }};
 {% endif %}
         let left_value = shared_values[left];
         let right_value = shared_values[right];
           shared_indices[right] = left_index;
         }
       }
+{% endmacro %}
+{% set sortSize = blockSize if (stage is defined and stage == "block") else sharedSize %}
+{% set pruneFinal = sortSize > workgroupSize and sortRetainedSize < sortSize %}
+  for (var block_size = 2u; block_size {{ "<" if pruneFinal else "<=" }} SHARED_SIZE; block_size <<= 1u) {
+    for (var gap = block_size >> 1u; gap > 0u; gap >>= 1u) {
+{{ compare_pairs("SHARED_SIZE / 2u", false) }}
       workgroupBarrier();
     }
   }
+{% if pruneFinal %}
+  for (var gap = SHARED_SIZE >> 1u; gap > 0u; gap >>= 1u) {
+{{ compare_pairs("max(gap, " ~ ((sortRetainedSize / 2)|int) ~ "u)", true) }}
+    workgroupBarrier();
+  }
+{% endif %}
   for (var rank = tid; rank < params.k; rank += WG) {
     let output_index = (outer_index * params.k + rank) * params.inner + inner_index;
+    values[output_index] = {% if floatOrderKeys is defined and floatOrderKeys %}x[input_base + shared_indices[rank] * params.inner]{% else %}{{ scalar }}(shared_values[rank]){% endif %};
     indices[output_index] = shared_indices[rank];
   }
 }

build/webgpu/topk-large-block.wgsl.jinja CHANGED Viewed

@@ -1,38 +1,49 @@
-{% if usesF16 and source.stage != "scratch" %}
 enable f16;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
-// i32/u32 inputs sort by u32 order keys so candidate scratch remains exact; an
-// f32 scratch would corrupt integers above 2^24.
 {% set is_int = scalar == "i32" or scalar == "u32" %}
-{% if is_int %}
 {% set val_t = "u32" %}
 {% set order_scalar = "u32" %}
-{% else %}
-{% set val_t = "f32" %}
-{% endif %}
 // Stable bitonic comparator for TopK values and their original indices.
 // Stable order key and NaN predicate for direct, tournament, and bitonic TopK.
 {% if order_scalar is not defined %}{% set order_scalar = scalar %}{% endif %}
 const PAD_INDEX: u32 = 0xffffffffu;
 fn order_key(value: {{ val_t }}) -> u32 {
-{% if is_int and order_scalar == "i32" %}
   return bitcast<u32>(value) ^ 0x80000000u;
-{% elif is_int and order_scalar == "u32" %}
   return value;
 {% else %}
-  let bits0 = bitcast<u32>(value);
-  let bits = select(bits0, 0u, bits0 == 0x80000000u);
-  return select(bits | 0x80000000u, ~bits, (bits & 0x80000000u) != 0u);
 {% endif %}
 }
 {% if not is_int %}
-fn is_nan_value(value: f32) -> bool {
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u
       && (bits & 0x007fffffu) != 0u;
 }
 {% endif %}
@@ -66,7 +77,7 @@ fn should_swap(a: {{ val_t }}, a_index: u32, b: {{ val_t }}, b_index: u32, ascen
   return select(descending_swap, ascending_swap, ascending);
 }
-{% if is_int and source.stage == "block" %}
 fn input_order_key(v: {{ scalar }}) -> u32 {
 {% if scalar == "i32" %}
   return bitcast<u32>(v) ^ 0x80000000u;
@@ -75,7 +86,7 @@ fn input_order_key(v: {{ scalar }}) -> u32 {
 {% endif %}
 }
 {% endif %}
-{% if is_int and source.stage == "output" %}
 fn key_to_value(k: u32) -> {{ scalar }} {
 {% if scalar == "i32" %}
   return bitcast<i32>(k ^ 0x80000000u);
@@ -85,21 +96,20 @@ fn key_to_value(k: u32) -> {{ scalar }} {
 }
 {% endif %}
-const SHARED_SIZE: u32 = {{ blockSize if source.stage == "block" else sharedSize }}u;
 const WG: u32 = {{ workgroupSize }}u;
-var<workgroup> shared_values: array<{{ val_t }}, {{ blockSize if source.stage == "block" else sharedSize }}>;
-var<workgroup> shared_indices: array<u32, {{ blockSize if source.stage == "block" else sharedSize }}>;
 @compute @workgroup_size(WG, 1, 1)
 fn main(
   @builtin(workgroup_id) workgroup: vec3<u32>,
   @builtin(local_invocation_id) local: vec3<u32>,
-  @builtin(num_workgroups) num_workgroups: vec3<u32>,
 ) {
   let tid = local.x;
-{% if source.stage == "output" %}
-  let row = workgroup.x + workgroup.y * num_workgroups.x;
   if (row >= params.rows) {
     return;
   }
@@ -107,16 +117,16 @@ fn main(
   let input_count = params.inCandidates;
   let input_base = row * input_count;
 {% else %}
-  let flat_group = workgroup.x + workgroup.y * num_workgroups.x;
-  let groups_per_row = params.{{ "blocks" if source.stage == "block" else "outGroups" }};
   if (flat_group >= params.rows * groups_per_row) {
     return;
   }
   let row = flat_group / groups_per_row;
   let group = flat_group % groups_per_row;
-{% if source.stage == "block" %}
   let input_count = params.cols;
-{% if source.strided is defined and source.strided %}
   // Strided (outer, axis, inner) view: row = outer * inner + innerIndex, and
   // axis element j sits at (outer * cols + j) * inner + innerIndex.
   let input_base = ((row / params.inner) * params.cols) * params.inner + (row % params.inner);
@@ -132,38 +142,31 @@ fn main(
   for (var slot = tid; slot < SHARED_SIZE; slot += WG) {
     let source_index = group * SHARED_SIZE + slot;
     if (source_index < input_count) {
-{% if source.stage == "block" %}
-{% set X_INDEX = "input_base + source_index * params.inner" if (source.strided is defined and source.strided) else "input_base + slot" %}
-      shared_values[slot] = {{ "input_order_key" if is_int else "f32" }}(x[{{ X_INDEX }}]);
       shared_indices[slot] = source_index;
 {% else %}
       shared_values[slot] = candidateVals[input_base + slot];
       shared_indices[slot] = candidateIdxs[input_base + slot];
 {% endif %}
     } else {
-      shared_values[slot] = {{ "0u" if is_int else "0.0" }};
       shared_indices[slot] = PAD_INDEX;
     }
   }
   workgroupBarrier();
-  // In-place bitonic ordering of values and their original indices.
-  for (var block_size = 2u; block_size <= SHARED_SIZE; block_size <<= 1u) {
-    for (var gap = block_size >> 1u; gap > 0u; gap >>= 1u) {
-      for (var pair = tid; pair < SHARED_SIZE / 2u; pair += WG) {
-        // `gap` and `block_size` are powers of two by construction (block_size
-        // doubles from 2, gap halves from block_size >> 1), so the index math is
-        // masks and shifts. GPUs have no integer divide unit and neither
-        // value is a compile-time constant here, so a literal `/` and `%` expand
-        // to a full division per pair per stage — three of them, against a body
-        // that is otherwise four shared loads and a compare.
         let position = pair & (gap - 1u);
         let left = ((pair - position) << 1u) + position;
         let right = left + gap;
 {% if attrs.largest == 0 %}
-        let ascending = (left & block_size) == 0u;
 {% else %}
-        let ascending = (left & block_size) != 0u;
 {% endif %}
         let left_value = shared_values[left];
         let right_value = shared_values[right];
@@ -176,26 +179,46 @@ fn main(
           shared_indices[right] = left_index;
         }
       }
       workgroupBarrier();
     }
   }
   for (var rank = tid; rank < params.k; rank += WG) {
-{% if source.stage == "block" %}
     let output_index = (row * params.blocks + group) * params.k + rank;
     candidateVals[output_index] = shared_values[rank];
     candidateIdxs[output_index] = shared_indices[rank];
-{% elif source.stage == "scratch" %}
     let output_index = (row * params.outGroups + group) * params.k + rank;
     candidateVals2[output_index] = shared_values[rank];
     candidateIdxs2[output_index] = shared_indices[rank];
 {% else %}
-{% if source.strided is defined and source.strided %}
     let output_index = ((row / params.inner) * params.k + rank) * params.inner + (row % params.inner);
 {% else %}
     let output_index = row * params.k + rank;
 {% endif %}
-    values[output_index] = {{ "key_to_value" if is_int else scalar }}(shared_values[rank]);
     indices[output_index] = shared_indices[rank];
 {% endif %}
   }

+{% if usesF16 and stage != "scratch" %}
 enable f16;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
+// All inputs sort by precomputed u32 order keys. Integer keys remain exact;
+// floating outputs gather the source value after selection to retain its bits.
 {% set is_int = scalar == "i32" or scalar == "u32" %}
 {% set val_t = "u32" %}
 {% set order_scalar = "u32" %}
 // Stable bitonic comparator for TopK values and their original indices.
 // Stable order key and NaN predicate for direct, tournament, and bitonic TopK.
 {% if order_scalar is not defined %}{% set order_scalar = scalar %}{% endif %}
 const PAD_INDEX: u32 = 0xffffffffu;
+{% macro float_order_key(value) %}
+  let bits0 = bitcast<u32>({{ value }});
+  let bits = select(bits0, 0u, bits0 == 0x80000000u);
+  return select(bits | 0x80000000u, ~bits, (bits & 0x80000000u) != 0u);
+{% endmacro %}
 fn order_key(value: {{ val_t }}) -> u32 {
+{% if order_scalar == "i32" %}
   return bitcast<u32>(value) ^ 0x80000000u;
+{% elif order_scalar == "u32" %}
   return value;
 {% else %}
+{{ float_order_key("value") }}
 {% endif %}
 }
 {% if not is_int %}
+fn is_nan_value(value: {{ val_t }}) -> bool {
+{% if floatOrderKeys is defined and floatOrderKeys %}
+  // Encoded infinities bound the finite keys; either outer range is NaN.
+  return value > 0xff800000u || value < 0x007fffffu;
+{% else %}
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u
       && (bits & 0x007fffffu) != 0u;
+{% endif %}
+}
+{% endif %}
+{% if floatOrderKeys is defined and floatOrderKeys and (stage is not defined or stage == "block") %}
+// Encode once on input. Comparators consume keys, while output gathers the
+// selected original values to preserve their representation, including -0.
+fn input_order_key(value: {{ scalar }}) -> u32 {
+{{ float_order_key("f32(value)") }}
 }
 {% endif %}
   return select(descending_swap, ascending_swap, ascending);
 }
+{% if is_int and stage == "block" %}
 fn input_order_key(v: {{ scalar }}) -> u32 {
 {% if scalar == "i32" %}
   return bitcast<u32>(v) ^ 0x80000000u;
 {% endif %}
 }
 {% endif %}
+{% if is_int and stage == "output" %}
 fn key_to_value(k: u32) -> {{ scalar }} {
 {% if scalar == "i32" %}
   return bitcast<i32>(k ^ 0x80000000u);
 }
 {% endif %}
+const SHARED_SIZE: u32 = {{ blockSize if stage == "block" else sharedSize }}u;
 const WG: u32 = {{ workgroupSize }}u;
+var<workgroup> shared_values: array<{{ val_t }}, {{ blockSize if stage == "block" else sharedSize }}>;
+var<workgroup> shared_indices: array<u32, {{ blockSize if stage == "block" else sharedSize }}>;
 @compute @workgroup_size(WG, 1, 1)
 fn main(
   @builtin(workgroup_id) workgroup: vec3<u32>,
   @builtin(local_invocation_id) local: vec3<u32>,
 ) {
   let tid = local.x;
+{% if stage == "output" %}
+  let row = workgroup.x + workgroup.y * {{ DISPATCH_FOLD_WIDTH }}u;
   if (row >= params.rows) {
     return;
   }
   let input_count = params.inCandidates;
   let input_base = row * input_count;
 {% else %}
+  let flat_group = workgroup.x + workgroup.y * {{ DISPATCH_FOLD_WIDTH }}u;
+  let groups_per_row = params.{{ "blocks" if stage == "block" else "outGroups" }};
   if (flat_group >= params.rows * groups_per_row) {
     return;
   }
   let row = flat_group / groups_per_row;
   let group = flat_group % groups_per_row;
+{% if stage == "block" %}
   let input_count = params.cols;
+{% if strided is defined and strided %}
   // Strided (outer, axis, inner) view: row = outer * inner + innerIndex, and
   // axis element j sits at (outer * cols + j) * inner + innerIndex.
   let input_base = ((row / params.inner) * params.cols) * params.inner + (row % params.inner);
   for (var slot = tid; slot < SHARED_SIZE; slot += WG) {
     let source_index = group * SHARED_SIZE + slot;
     if (source_index < input_count) {
+{% if stage == "block" %}
+{% set X_INDEX = "input_base + source_index * params.inner" if (strided is defined and strided) else "input_base + slot" %}
+      shared_values[slot] = input_order_key(x[{{ X_INDEX }}]);
       shared_indices[slot] = source_index;
 {% else %}
       shared_values[slot] = candidateVals[input_base + slot];
       shared_indices[slot] = candidateIdxs[input_base + slot];
 {% endif %}
     } else {
+      shared_values[slot] = 0u;
       shared_indices[slot] = PAD_INDEX;
     }
   }
   workgroupBarrier();
+{% macro compare_pairs(pairCount, finalStage) %}
+      for (var pair = tid; pair < {{ pairCount }}; pair += WG) {
+        // Power-of-two gaps let masks and shifts identify disjoint pairs.
         let position = pair & (gap - 1u);
         let left = ((pair - position) << 1u) + position;
         let right = left + gap;
 {% if attrs.largest == 0 %}
+        let ascending = {{ "true" if finalStage else "(left & block_size) == 0u" }};
 {% else %}
+        let ascending = {{ "false" if finalStage else "(left & block_size) != 0u" }};
 {% endif %}
         let left_value = shared_values[left];
         let right_value = shared_values[right];
           shared_indices[right] = left_index;
         }
       }
+{% endmacro %}
+{% set sortSize = blockSize if (stage is defined and stage == "block") else sharedSize %}
+{% set pruneFinal = sortSize > workgroupSize and sortRetainedSize < sortSize %}
+  for (var block_size = 2u; block_size {{ "<" if pruneFinal else "<=" }} SHARED_SIZE; block_size <<= 1u) {
+    for (var gap = block_size >> 1u; gap > 0u; gap >>= 1u) {
+{{ compare_pairs("SHARED_SIZE / 2u", false) }}
       workgroupBarrier();
     }
   }
+{% if pruneFinal %}
+  for (var gap = SHARED_SIZE >> 1u; gap > 0u; gap >>= 1u) {
+{{ compare_pairs("max(gap, " ~ ((sortRetainedSize / 2)|int) ~ "u)", true) }}
+    workgroupBarrier();
+  }
+{% endif %}
   for (var rank = tid; rank < params.k; rank += WG) {
+{% if stage == "block" %}
     let output_index = (row * params.blocks + group) * params.k + rank;
     candidateVals[output_index] = shared_values[rank];
     candidateIdxs[output_index] = shared_indices[rank];
+{% elif stage == "scratch" %}
     let output_index = (row * params.outGroups + group) * params.k + rank;
     candidateVals2[output_index] = shared_values[rank];
     candidateIdxs2[output_index] = shared_indices[rank];
 {% else %}
+{% if strided is defined and strided %}
     let output_index = ((row / params.inner) * params.k + rank) * params.inner + (row % params.inner);
 {% else %}
     let output_index = row * params.k + rank;
 {% endif %}
+{% if floatOrderKeys is defined and floatOrderKeys %}
+{% if strided is defined and strided %}
+    values[output_index] = x[((row / params.inner) * params.cols + shared_indices[rank]) * params.inner + row % params.inner];
+{% else %}
+    values[output_index] = x[row * params.cols + shared_indices[rank]];
+{% endif %}
+{% else %}
+    values[output_index] = key_to_value(shared_values[rank]);
+{% endif %}
     indices[output_index] = shared_indices[rank];
 {% endif %}
   }

build/webgpu/topk-portable-rows-smallk.wgsl.jinja CHANGED Viewed

@@ -1,15 +1,11 @@
 // Portable one-workgroup-per-row TopK for many medium-width rows and small K.
-// Each lane keeps a small sorted register list for its strided columns. For each
 // output rank, lane zero scans one candidate per lane in shared memory. This
-// O(WG*K) merge avoids sorting and synchronizing the whole row while retaining
-// the subgroup route's register-local algorithmic structure.
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 // Stable TopK ordering. Native integer routes retain i32/u32 values; float and
 // split-scratch routes compare f32 values.
-{% set is_int = source.nativeValues and (scalar == "i32" or scalar == "u32") %}
 {% set val_t = scalar if is_int else "f32" %}
 {% set zero = scalar ~ "(0)" if is_int else "0.0" %}
@@ -17,24 +13,39 @@ enable f16;
 {% if order_scalar is not defined %}{% set order_scalar = scalar %}{% endif %}
 const PAD_INDEX: u32 = 0xffffffffu;
 fn order_key(value: {{ val_t }}) -> u32 {
-{% if is_int and order_scalar == "i32" %}
   return bitcast<u32>(value) ^ 0x80000000u;
-{% elif is_int and order_scalar == "u32" %}
   return value;
 {% else %}
-  // IEEE regards signed zero as equal, so normalize -0 before tie-breaking.
-  let bits0 = bitcast<u32>(value);
-  let bits = select(bits0, 0u, bits0 == 0x80000000u);
-  return select(bits | 0x80000000u, ~bits, (bits & 0x80000000u) != 0u);
 {% endif %}
 }
 {% if not is_int %}
-fn is_nan_value(value: f32) -> bool {
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u
       && (bits & 0x007fffffu) != 0u;
 }
 {% endif %}
@@ -66,9 +77,8 @@ var<workgroup> winnerIndex: u32;
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(workgroup_id) wgid: vec3<u32>,
-        @builtin(local_invocation_id) lid: vec3<u32>,
-        @builtin(num_workgroups) nwg: vec3<u32>) {
-  let row = wgid.x + wgid.y * nwg.x;
   if (row >= params.rows) { return; }
   var localValues: array<{{ val_t }}, {{ localItems }}>;

 // Portable one-workgroup-per-row TopK for many medium-width rows and small K.
+// Each lane keeps a small sorted candidate list for its strided columns. For each
 // output rank, lane zero scans one candidate per lane in shared memory. This
+// O(WG*K) merge avoids sorting and synchronizing the whole row.
 {{ env.wgsl.resourceDeclarations }}
 // Stable TopK ordering. Native integer routes retain i32/u32 values; float and
 // split-scratch routes compare f32 values.
+{% set is_int = nativeValues and (scalar == "i32" or scalar == "u32") %}
 {% set val_t = scalar if is_int else "f32" %}
 {% set zero = scalar ~ "(0)" if is_int else "0.0" %}
 {% if order_scalar is not defined %}{% set order_scalar = scalar %}{% endif %}
 const PAD_INDEX: u32 = 0xffffffffu;
+{% macro float_order_key(value) %}
+  // IEEE regards signed zero as equal, so normalize -0 before tie-breaking.
+  let bits0 = bitcast<u32>({{ value }});
+  let bits = select(bits0, 0u, bits0 == 0x80000000u);
+  return select(bits | 0x80000000u, ~bits, (bits & 0x80000000u) != 0u);
+{% endmacro %}
 fn order_key(value: {{ val_t }}) -> u32 {
+{% if order_scalar == "i32" %}
   return bitcast<u32>(value) ^ 0x80000000u;
+{% elif order_scalar == "u32" %}
   return value;
 {% else %}
+{{ float_order_key("value") }}
 {% endif %}
 }
 {% if not is_int %}
+fn is_nan_value(value: {{ val_t }}) -> bool {
+{% if floatOrderKeys is defined and floatOrderKeys %}
+  // Encoded infinities bound the finite keys; either outer range is NaN.
+  return value > 0xff800000u || value < 0x007fffffu;
+{% else %}
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u
       && (bits & 0x007fffffu) != 0u;
+{% endif %}
+}
+{% endif %}
+{% if floatOrderKeys is defined and floatOrderKeys and (stage is not defined or stage == "block") %}
+// Encode once on input. Comparators consume keys, while output gathers the
+// selected original values to preserve their representation, including -0.
+fn input_order_key(value: {{ scalar }}) -> u32 {
+{{ float_order_key("f32(value)") }}
 }
 {% endif %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(workgroup_id) wgid: vec3<u32>,
+        @builtin(local_invocation_id) lid: vec3<u32>) {
+  let row = wgid.x + wgid.y * {{ DISPATCH_FOLD_WIDTH }}u;
   if (row >= params.rows) { return; }
   var localValues: array<{{ val_t }}, {{ localItems }}>;

build/webgpu/topk-small-rows-batched.wgsl.jinja CHANGED Viewed

@@ -1,11 +1,8 @@
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 // One lane owns each independent short row and retains the native integer type.
 // Stable TopK ordering. Native integer routes retain i32/u32 values; float and
 // split-scratch routes compare f32 values.
-{% set is_int = source.nativeValues and (scalar == "i32" or scalar == "u32") %}
 {% set val_t = scalar if is_int else "f32" %}
 {% set zero = scalar ~ "(0)" if is_int else "0.0" %}
@@ -13,24 +10,39 @@ enable f16;
 {% if order_scalar is not defined %}{% set order_scalar = scalar %}{% endif %}
 const PAD_INDEX: u32 = 0xffffffffu;
 fn order_key(value: {{ val_t }}) -> u32 {
-{% if is_int and order_scalar == "i32" %}
   return bitcast<u32>(value) ^ 0x80000000u;
-{% elif is_int and order_scalar == "u32" %}
   return value;
 {% else %}
-  // IEEE regards signed zero as equal, so normalize -0 before tie-breaking.
-  let bits0 = bitcast<u32>(value);
-  let bits = select(bits0, 0u, bits0 == 0x80000000u);
-  return select(bits | 0x80000000u, ~bits, (bits & 0x80000000u) != 0u);
 {% endif %}
 }
 {% if not is_int %}
-fn is_nan_value(value: f32) -> bool {
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u
       && (bits & 0x007fffffu) != 0u;
 }
 {% endif %}
@@ -60,10 +72,9 @@ const MAX_K: u32 = {{ tunables.SMALL_ROWS_MAX_K }}u;
 @compute @workgroup_size(WG, 1, 1)
 fn main(
   @builtin(global_invocation_id) gid: vec3<u32>,
-  @builtin(num_workgroups) nwg: vec3<u32>,
 ) {
   // Rebuild the row index after a large grid folds into dispatch.y.
-  let row = gid.x + gid.y * nwg.x * WG;
   if (row >= params.rows) {
     return;
   }

 {{ env.wgsl.resourceDeclarations }}
 // One lane owns each independent short row and retains the native integer type.
 // Stable TopK ordering. Native integer routes retain i32/u32 values; float and
 // split-scratch routes compare f32 values.
+{% set is_int = nativeValues and (scalar == "i32" or scalar == "u32") %}
 {% set val_t = scalar if is_int else "f32" %}
 {% set zero = scalar ~ "(0)" if is_int else "0.0" %}
 {% if order_scalar is not defined %}{% set order_scalar = scalar %}{% endif %}
 const PAD_INDEX: u32 = 0xffffffffu;
+{% macro float_order_key(value) %}
+  // IEEE regards signed zero as equal, so normalize -0 before tie-breaking.
+  let bits0 = bitcast<u32>({{ value }});
+  let bits = select(bits0, 0u, bits0 == 0x80000000u);
+  return select(bits | 0x80000000u, ~bits, (bits & 0x80000000u) != 0u);
+{% endmacro %}
 fn order_key(value: {{ val_t }}) -> u32 {
+{% if order_scalar == "i32" %}
   return bitcast<u32>(value) ^ 0x80000000u;
+{% elif order_scalar == "u32" %}
   return value;
 {% else %}
+{{ float_order_key("value") }}
 {% endif %}
 }
 {% if not is_int %}
+fn is_nan_value(value: {{ val_t }}) -> bool {
+{% if floatOrderKeys is defined and floatOrderKeys %}
+  // Encoded infinities bound the finite keys; either outer range is NaN.
+  return value > 0xff800000u || value < 0x007fffffu;
+{% else %}
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u
       && (bits & 0x007fffffu) != 0u;
+{% endif %}
+}
+{% endif %}
+{% if floatOrderKeys is defined and floatOrderKeys and (stage is not defined or stage == "block") %}
+// Encode once on input. Comparators consume keys, while output gathers the
+// selected original values to preserve their representation, including -0.
+fn input_order_key(value: {{ scalar }}) -> u32 {
+{{ float_order_key("f32(value)") }}
 }
 {% endif %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(
   @builtin(global_invocation_id) gid: vec3<u32>,
 ) {
   // Rebuild the row index after a large grid folds into dispatch.y.
+  let row = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
   if (row >= params.rows) {
     return;
   }

build/webgpu/topk-strided-smallk.wgsl.jinja CHANGED Viewed

@@ -1,10 +1,7 @@
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 // Stable TopK ordering. Native integer routes retain i32/u32 values; float and
 // split-scratch routes compare f32 values.
-{% set is_int = source.nativeValues and (scalar == "i32" or scalar == "u32") %}
 {% set val_t = scalar if is_int else "f32" %}
 {% set zero = scalar ~ "(0)" if is_int else "0.0" %}
@@ -12,24 +9,39 @@ enable f16;
 {% if order_scalar is not defined %}{% set order_scalar = scalar %}{% endif %}
 const PAD_INDEX: u32 = 0xffffffffu;
 fn order_key(value: {{ val_t }}) -> u32 {
-{% if is_int and order_scalar == "i32" %}
   return bitcast<u32>(value) ^ 0x80000000u;
-{% elif is_int and order_scalar == "u32" %}
   return value;
 {% else %}
-  // IEEE regards signed zero as equal, so normalize -0 before tie-breaking.
-  let bits0 = bitcast<u32>(value);
-  let bits = select(bits0, 0u, bits0 == 0x80000000u);
-  return select(bits | 0x80000000u, ~bits, (bits & 0x80000000u) != 0u);
 {% endif %}
 }
 {% if not is_int %}
-fn is_nan_value(value: f32) -> bool {
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u
       && (bits & 0x007fffffu) != 0u;
 }
 {% endif %}
@@ -62,10 +74,9 @@ var<workgroup> shared_indices: array<u32, {{ workgroupSize }}>;
 @compute @workgroup_size(WG, 1, 1)
 fn main(
   @builtin(workgroup_id) workgroup: vec3<u32>,
-  @builtin(num_workgroups) nwg: vec3<u32>,
   @builtin(local_invocation_id) local: vec3<u32>,
 ) {
-  let output_position = workgroup.x + workgroup.y * nwg.x;
   let tid = local.x;
   if (output_position >= params.outputPositions) {
     return;

 {{ env.wgsl.resourceDeclarations }}
 // Stable TopK ordering. Native integer routes retain i32/u32 values; float and
 // split-scratch routes compare f32 values.
+{% set is_int = nativeValues and (scalar == "i32" or scalar == "u32") %}
 {% set val_t = scalar if is_int else "f32" %}
 {% set zero = scalar ~ "(0)" if is_int else "0.0" %}
 {% if order_scalar is not defined %}{% set order_scalar = scalar %}{% endif %}
 const PAD_INDEX: u32 = 0xffffffffu;
+{% macro float_order_key(value) %}
+  // IEEE regards signed zero as equal, so normalize -0 before tie-breaking.
+  let bits0 = bitcast<u32>({{ value }});
+  let bits = select(bits0, 0u, bits0 == 0x80000000u);
+  return select(bits | 0x80000000u, ~bits, (bits & 0x80000000u) != 0u);
+{% endmacro %}
 fn order_key(value: {{ val_t }}) -> u32 {
+{% if order_scalar == "i32" %}
   return bitcast<u32>(value) ^ 0x80000000u;
+{% elif order_scalar == "u32" %}
   return value;
 {% else %}
+{{ float_order_key("value") }}
 {% endif %}
 }
 {% if not is_int %}
+fn is_nan_value(value: {{ val_t }}) -> bool {
+{% if floatOrderKeys is defined and floatOrderKeys %}
+  // Encoded infinities bound the finite keys; either outer range is NaN.
+  return value > 0xff800000u || value < 0x007fffffu;
+{% else %}
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u
       && (bits & 0x007fffffu) != 0u;
+{% endif %}
+}
+{% endif %}
+{% if floatOrderKeys is defined and floatOrderKeys and (stage is not defined or stage == "block") %}
+// Encode once on input. Comparators consume keys, while output gathers the
+// selected original values to preserve their representation, including -0.
+fn input_order_key(value: {{ scalar }}) -> u32 {
+{{ float_order_key("f32(value)") }}
 }
 {% endif %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(
   @builtin(workgroup_id) workgroup: vec3<u32>,
   @builtin(local_invocation_id) local: vec3<u32>,
 ) {
+  let output_position = workgroup.x + workgroup.y * {{ DISPATCH_FOLD_WIDTH }}u;
   let tid = local.x;
   if (output_position >= params.outputPositions) {
     return;

build/webgpu/topk-subgroup-rows.wgsl.jinja CHANGED Viewed

@@ -1,7 +1,4 @@
 enable subgroups;
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 // One subgroup owns one contiguous row. Each lane keeps a short local TopK, then
 // log2(SG) shuffle reductions choose each row winner without workgroup memory or
@@ -9,7 +6,7 @@ enable f16;
 // variable-width routes use one minimum-width subgroup per workgroup.
 // Stable TopK ordering. Native integer routes retain i32/u32 values; float and
 // split-scratch routes compare f32 values.
-{% set is_int = source.nativeValues and (scalar == "i32" or scalar == "u32") %}
 {% set val_t = scalar if is_int else "f32" %}
 {% set zero = scalar ~ "(0)" if is_int else "0.0" %}
@@ -17,24 +14,39 @@ enable f16;
 {% if order_scalar is not defined %}{% set order_scalar = scalar %}{% endif %}
 const PAD_INDEX: u32 = 0xffffffffu;
 fn order_key(value: {{ val_t }}) -> u32 {
-{% if is_int and order_scalar == "i32" %}
   return bitcast<u32>(value) ^ 0x80000000u;
-{% elif is_int and order_scalar == "u32" %}
   return value;
 {% else %}
-  // IEEE regards signed zero as equal, so normalize -0 before tie-breaking.
-  let bits0 = bitcast<u32>(value);
-  let bits = select(bits0, 0u, bits0 == 0x80000000u);
-  return select(bits | 0x80000000u, ~bits, (bits & 0x80000000u) != 0u);
 {% endif %}
 }
 {% if not is_int %}
-fn is_nan_value(value: f32) -> bool {
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u
       && (bits & 0x007fffffu) != 0u;
 }
 {% endif %}
@@ -66,11 +78,10 @@ const LOCAL_ITEMS: u32 = {{ localItems }}u;
 @compute @workgroup_size(WG, 1, 1)
 fn main(
   @builtin(workgroup_id) workgroup: vec3<u32>,
-  @builtin(num_workgroups) nwg: vec3<u32>,
   @builtin(subgroup_id) subgroup_id: u32,
   @builtin(subgroup_invocation_id) lane: u32,
 ) {
-  let group = workgroup.x + workgroup.y * nwg.x;
   let row = group * SUBGROUPS_PER_WG + subgroup_id;
   if (row >= params.rows) {
     return;

 enable subgroups;
 {{ env.wgsl.resourceDeclarations }}
 // One subgroup owns one contiguous row. Each lane keeps a short local TopK, then
 // log2(SG) shuffle reductions choose each row winner without workgroup memory or
 // variable-width routes use one minimum-width subgroup per workgroup.
 // Stable TopK ordering. Native integer routes retain i32/u32 values; float and
 // split-scratch routes compare f32 values.
+{% set is_int = nativeValues and (scalar == "i32" or scalar == "u32") %}
 {% set val_t = scalar if is_int else "f32" %}
 {% set zero = scalar ~ "(0)" if is_int else "0.0" %}
 {% if order_scalar is not defined %}{% set order_scalar = scalar %}{% endif %}
 const PAD_INDEX: u32 = 0xffffffffu;
+{% macro float_order_key(value) %}
+  // IEEE regards signed zero as equal, so normalize -0 before tie-breaking.
+  let bits0 = bitcast<u32>({{ value }});
+  let bits = select(bits0, 0u, bits0 == 0x80000000u);
+  return select(bits | 0x80000000u, ~bits, (bits & 0x80000000u) != 0u);
+{% endmacro %}
 fn order_key(value: {{ val_t }}) -> u32 {
+{% if order_scalar == "i32" %}
   return bitcast<u32>(value) ^ 0x80000000u;
+{% elif order_scalar == "u32" %}
   return value;
 {% else %}
+{{ float_order_key("value") }}
 {% endif %}
 }
 {% if not is_int %}
+fn is_nan_value(value: {{ val_t }}) -> bool {
+{% if floatOrderKeys is defined and floatOrderKeys %}
+  // Encoded infinities bound the finite keys; either outer range is NaN.
+  return value > 0xff800000u || value < 0x007fffffu;
+{% else %}
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u
       && (bits & 0x007fffffu) != 0u;
+{% endif %}
+}
+{% endif %}
+{% if floatOrderKeys is defined and floatOrderKeys and (stage is not defined or stage == "block") %}
+// Encode once on input. Comparators consume keys, while output gathers the
+// selected original values to preserve their representation, including -0.
+fn input_order_key(value: {{ scalar }}) -> u32 {
+{{ float_order_key("f32(value)") }}
 }
 {% endif %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(
   @builtin(workgroup_id) workgroup: vec3<u32>,
   @builtin(subgroup_id) subgroup_id: u32,
   @builtin(subgroup_invocation_id) lane: u32,
 ) {
+  let group = workgroup.x + workgroup.y * {{ DISPATCH_FOLD_WIDTH }}u;
   let row = group * SUBGROUPS_PER_WG + subgroup_id;
   if (row >= params.rows) {
     return;

build/webgpu/topk-top1-last-axis.wgsl.jinja CHANGED Viewed

@@ -1,10 +1,7 @@
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 // Stable TopK ordering. Native integer routes retain i32/u32 values; float and
 // split-scratch routes compare f32 values.
-{% set is_int = source.nativeValues and (scalar == "i32" or scalar == "u32") %}
 {% set val_t = scalar if is_int else "f32" %}
 {% set zero = scalar ~ "(0)" if is_int else "0.0" %}
@@ -12,24 +9,39 @@ enable f16;
 {% if order_scalar is not defined %}{% set order_scalar = scalar %}{% endif %}
 const PAD_INDEX: u32 = 0xffffffffu;
 fn order_key(value: {{ val_t }}) -> u32 {
-{% if is_int and order_scalar == "i32" %}
   return bitcast<u32>(value) ^ 0x80000000u;
-{% elif is_int and order_scalar == "u32" %}
   return value;
 {% else %}
-  // IEEE regards signed zero as equal, so normalize -0 before tie-breaking.
-  let bits0 = bitcast<u32>(value);
-  let bits = select(bits0, 0u, bits0 == 0x80000000u);
-  return select(bits | 0x80000000u, ~bits, (bits & 0x80000000u) != 0u);
 {% endif %}
 }
 {% if not is_int %}
-fn is_nan_value(value: f32) -> bool {
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u
       && (bits & 0x007fffffu) != 0u;
 }
 {% endif %}
@@ -54,7 +66,7 @@ fn is_better(value: {{ val_t }}, index: u32, best: {{ val_t }}, best_index: u32)
 const WG: u32 = {{ workgroupSize }}u;
-{% if source.stage == "block" %}
 const BLOCK_SIZE: u32 = {{ top1BlockSize }}u;
 {% endif %}
@@ -65,11 +77,10 @@ var<workgroup> shared_indices: array<u32, {{ workgroupSize }}>;
 fn main(
   @builtin(workgroup_id) workgroup: vec3<u32>,
   @builtin(local_invocation_id) local: vec3<u32>,
-  @builtin(num_workgroups) num_workgroups: vec3<u32>,
 ) {
   let tid = local.x;
-{% if source.stage == "block" %}
-  let flat_group = workgroup.x + workgroup.y * num_workgroups.x;
   if (flat_group >= params.rows * params.blocks) {
     return;
   }
@@ -78,11 +89,11 @@ fn main(
   let begin = block * BLOCK_SIZE;
   let end = min(begin + BLOCK_SIZE, params.cols);
 {% else %}
-  let row = workgroup.x + workgroup.y * num_workgroups.x;
   if (row >= params.rows) {
     return;
   }
-{% if source.stage == "direct" %}
   let begin = 0u;
   let end = params.cols;
 {% endif %}
@@ -90,7 +101,7 @@ fn main(
   var best_value = {{ zero }};
   var best_index = PAD_INDEX;
-{% if source.stage == "output" %}
   let input_base = row * params.blocks;
   for (var candidate = tid; candidate < params.blocks; candidate += WG) {
     let value = candidateVals[input_base + candidate];
@@ -130,7 +141,7 @@ fn main(
   }
   if (tid == 0u) {
-{% if source.stage == "block" %}
     let output_index = row * params.blocks + block;
     candidateVals[output_index] = shared_values[0];
     candidateIdxs[output_index] = shared_indices[0];

 {{ env.wgsl.resourceDeclarations }}
 // Stable TopK ordering. Native integer routes retain i32/u32 values; float and
 // split-scratch routes compare f32 values.
+{% set is_int = nativeValues and (scalar == "i32" or scalar == "u32") %}
 {% set val_t = scalar if is_int else "f32" %}
 {% set zero = scalar ~ "(0)" if is_int else "0.0" %}
 {% if order_scalar is not defined %}{% set order_scalar = scalar %}{% endif %}
 const PAD_INDEX: u32 = 0xffffffffu;
+{% macro float_order_key(value) %}
+  // IEEE regards signed zero as equal, so normalize -0 before tie-breaking.
+  let bits0 = bitcast<u32>({{ value }});
+  let bits = select(bits0, 0u, bits0 == 0x80000000u);
+  return select(bits | 0x80000000u, ~bits, (bits & 0x80000000u) != 0u);
+{% endmacro %}
 fn order_key(value: {{ val_t }}) -> u32 {
+{% if order_scalar == "i32" %}
   return bitcast<u32>(value) ^ 0x80000000u;
+{% elif order_scalar == "u32" %}
   return value;
 {% else %}
+{{ float_order_key("value") }}
 {% endif %}
 }
 {% if not is_int %}
+fn is_nan_value(value: {{ val_t }}) -> bool {
+{% if floatOrderKeys is defined and floatOrderKeys %}
+  // Encoded infinities bound the finite keys; either outer range is NaN.
+  return value > 0xff800000u || value < 0x007fffffu;
+{% else %}
   let bits = bitcast<u32>(value);
   return (bits & 0x7f800000u) == 0x7f800000u
       && (bits & 0x007fffffu) != 0u;
+{% endif %}
+}
+{% endif %}
+{% if floatOrderKeys is defined and floatOrderKeys and (stage is not defined or stage == "block") %}
+// Encode once on input. Comparators consume keys, while output gathers the
+// selected original values to preserve their representation, including -0.
+fn input_order_key(value: {{ scalar }}) -> u32 {
+{{ float_order_key("f32(value)") }}
 }
 {% endif %}
 const WG: u32 = {{ workgroupSize }}u;
+{% if stage == "block" %}
 const BLOCK_SIZE: u32 = {{ top1BlockSize }}u;
 {% endif %}
 fn main(
   @builtin(workgroup_id) workgroup: vec3<u32>,
   @builtin(local_invocation_id) local: vec3<u32>,
 ) {
   let tid = local.x;
+{% if stage == "block" %}
+  let flat_group = workgroup.x + workgroup.y * {{ DISPATCH_FOLD_WIDTH }}u;
   if (flat_group >= params.rows * params.blocks) {
     return;
   }
   let begin = block * BLOCK_SIZE;
   let end = min(begin + BLOCK_SIZE, params.cols);
 {% else %}
+  let row = workgroup.x + workgroup.y * {{ DISPATCH_FOLD_WIDTH }}u;
   if (row >= params.rows) {
     return;
   }
+{% if stage == "direct" %}
   let begin = 0u;
   let end = params.cols;
 {% endif %}
   var best_value = {{ zero }};
   var best_index = PAD_INDEX;
+{% if stage == "output" %}
   let input_base = row * params.blocks;
   for (var candidate = tid; candidate < params.blocks; candidate += WG) {
     let value = candidateVals[input_base + candidate];
   }
   if (tid == 0u) {
+{% if stage == "block" %}
     let output_index = row * params.blocks + block;
     candidateVals[output_index] = shared_values[0];
     candidateIdxs[output_index] = shared_indices[0];