sync 91d990483a17

Browse files

Files changed (9) hide show

README.md +26 -22
build/webgpu/bench.json +697 -1
build/webgpu/manifest.json +202 -789
build/webgpu/metadata.json +22 -10
build/webgpu/sparse-attention-sgmat.wgsl.jinja +64 -58
build/webgpu/sparse-attention.wgsl.jinja +29 -45
build/webgpu/sparse-kv-append.wgsl.jinja +3 -7
build/webgpu/sparse-q-rotary.wgsl.jinja +3 -7
build/webgpu/test.json +550 -1

README.md CHANGED Viewed

@@ -18,27 +18,27 @@ See the [ONNX Runtime `SparseAttention` contrib-operator spec](https://github.co
 ## Inputs
-| Name | Bind key | Logical dtype | WebGPU storage | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- | --- |
-| `query` | `queryT` | `T` | same as logical dtype | `3` | — | Query `(batch_size, sequence_length, num_heads * head_size)`, or packed `[Q\|K\|V]` `(batch_size, sequence_length, (num_heads + 2 * kv_num_heads) * head_size)` when `key` and `value` are omitted. | required |
-| `key` | `keyT` | `T` | same as logical dtype | `3` | — | Key `(batch_size, sequence_length, kv_num_heads * head_size)`. Omitted for packed QKV. | optional |
-| `value` | `valueT` | `T` | same as logical dtype | `3` | — | Value `(batch_size, sequence_length, kv_num_heads * head_size)`. Omitted for packed QKV. | optional |
-| `past_key` | `pastKeyT` | `T` | same as logical dtype | `4` | — | Key cache `(batch_size, kv_num_heads, max_cache_sequence_length, head_size)`, updated in place. | required |
-| `past_value` | `pastValueT` | `T` | same as logical dtype | `4` | — | Value cache with the same shape as `past_key`, updated in place. | required |
-| `block_row_indices` | `blockRowIndicesT` | `M` | `int32` | `2` | — | CSR row pointers `(num_layout, max_blocks + 1)`. Each layout starts at zero, is monotonically non-decreasing, and ends no later than that layout's `block_col_indices` width. | required |
-| `block_col_indices` | `blockColIndicesT` | `M` | `int32` | `2` | — | CSR column indices `(num_layout, max_nnz_blocks)`, right-padded past each layout's non-zero count. Every active entry is in `[0, max_blocks)`. | required |
-| `total_sequence_length` | `totalSequenceLengthT` | `M` | `int32` | — | — | Scalar or one-element vector holding the maximum total key length. Equal to `sequence_length` exactly in the prompt case, which is how the past length is decided. The value fits the cache, the sparse layout's `max_blocks * sparse_block_size` capacity, and the rotary-cache row count when rotary is enabled. | required |
-| `key_total_sequence_lengths` | `keyTotalSequenceLengthsT` | `M` | `int32` | `1` | — | Per-batch total key length excluding padding, shape `(batch_size)`. Each value is at most `total_sequence_length` and is at least 1 for a prompt or at least `sequence_length` otherwise. | required |
-| `cos_cache` | `cosCacheT` | `T` | same as logical dtype | `2` | — | Rotary cosine cache `(max_rotary_sequence_length, rotary_dimension / 2)`, where the width is a multiple of 8 no larger than `head_size / 2`. Required with `sin_cache` when `do_rotary` is 1. | optional |
-| `sin_cache` | `sinCacheT` | `T` | same as logical dtype | `2` | — | Rotary sine cache with the same shape as `cos_cache`; required with it when `do_rotary` is 1. | optional |
 ## Outputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `output` | `outputT` | `T` | `3` | derived; see description | Attention output `(batch_size, sequence_length, num_heads * head_size)`. | required |
-| `past_key` | `pastKeyT` | `T` | `4` | same as `past_key` | The key cache tensor itself after the in-place append; ONNX names this output `present_key`. | required |
-| `past_value` | `pastValueT` | `T` | `4` | same as `past_value` | The value cache tensor itself after the in-place append; ONNX names this output `present_value`. | required |
 ## Attributes
@@ -47,11 +47,11 @@ Attributes and default values (overridable per request):
 | Attribute | Default | Description |
 | --- | --- | --- |
 | `do_rotary` | `0` | Set to 1 to apply rotary embedding to Q and to K before it enters the cache; every other value disables rotary embedding. |
-| `rotary_interleaved` | `0` | Set to 1 to rotate adjacent pairs instead of using the NeoX half-split; every other value selects the NeoX layout. |
-| `num_heads` | — | Number of query heads. |
 | `kv_num_heads` | — | Number of key/value heads; must divide `num_heads`. |
-| `sparse_block_size` | — | Tokens per sparse block; one of 16, 32, 64, 128. |
 | `scale` | — | Scale applied to query-key products; omitted or zero uses `1 / sqrt(head_size)`. |
 ## Type constraints
@@ -66,7 +66,7 @@ Some implementation variants require `subgroup-matrix` and `subgroups`. These ar
 ## 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
@@ -77,10 +77,14 @@ Some implementation variants require `subgroup-matrix` and `subgroups`. These ar
 ## 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 | WebGPU storage | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- | --- |
+| `queryT` | `query` | `T` | same as logical dtype | `3` | — | Query `(batch_size, sequence_length, num_heads * head_size)`, or packed `[Q\|K\|V]` `(batch_size, sequence_length, (num_heads + 2 * kv_num_heads) * head_size)` when `key` and `value` are omitted. | required |
+| `keyT` | `key` | `T` | same as logical dtype | `3` | — | Key `(batch_size, sequence_length, kv_num_heads * head_size)`. Omitted for packed QKV. | optional |
+| `valueT` | `value` | `T` | same as logical dtype | `3` | — | Value `(batch_size, sequence_length, kv_num_heads * head_size)`. Omitted for packed QKV. | optional |
+| `pastKeyT` | `past_key` | `T` | same as logical dtype | `4` | — | Key cache `(batch_size, kv_num_heads, max_cache_sequence_length, head_size)`, updated in place. | required |
+| `pastValueT` | `past_value` | `T` | same as logical dtype | `4` | — | Value cache with the same shape as `past_key`, updated in place. | required |
+| `blockRowIndicesT` | `block_row_indices` | `M` | `int32` | `2` | — | CSR row pointers `(num_layout, max_blocks + 1)`. Each layout starts at zero, is monotonically non-decreasing, and ends no later than that layout's `block_col_indices` width. | required |
+| `blockColIndicesT` | `block_col_indices` | `M` | `int32` | `2` | — | CSR column indices `(num_layout, max_nnz_blocks)`, right-padded past each layout's non-zero count. Every active entry is in `[0, max_blocks)`. | required |
+| `totalSequenceLengthT` | `total_sequence_length` | `M` | `int32` | — | — | Scalar or one-element vector holding the maximum total key length. Equal to `sequence_length` exactly in the prompt case, which is how the past length is decided. The value fits the cache, the sparse layout's `max_blocks * sparse_block_size` capacity, and the rotary-cache row count when rotary is enabled. | required |
+| `keyTotalSequenceLengthsT` | `key_total_sequence_lengths` | `M` | `int32` | `1` | — | Per-batch total key length excluding padding, shape `(batch_size)`. Each value is at most `total_sequence_length` and is at least 1 for a prompt or at least `sequence_length` otherwise. | required |
+| `cosCacheT` | `cos_cache` | `T` | same as logical dtype | `2` | — | Rotary cosine cache `(max_rotary_sequence_length, rotary_dimension / 2)`, where the width is a multiple of 8 no larger than `head_size / 2`. Required with `sin_cache` when `do_rotary` is 1. | optional |
+| `sinCacheT` | `sin_cache` | `T` | same as logical dtype | `2` | — | Rotary sine cache with the same shape as `cos_cache`; required with it when `do_rotary` is 1. | optional |
 ## Outputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `outputT` | `output` | `T` | `3` | derived | Attention output `(batch_size, sequence_length, num_heads * head_size)`. | required |
+| `pastKeyT` | `past_key` | `T` | `4` | same as `pastKeyT` | The key cache tensor itself after the in-place append; ONNX names this output `present_key`. | required |
+| `pastValueT` | `past_value` | `T` | `4` | same as `pastValueT` | The value cache tensor itself after the in-place append; ONNX names this output `present_value`. | required |
 ## Attributes
 | Attribute | Default | Description |
 | --- | --- | --- |
 | `do_rotary` | `0` | Set to 1 to apply rotary embedding to Q and to K before it enters the cache; every other value disables rotary embedding. |
 | `kv_num_heads` | — | Number of key/value heads; must divide `num_heads`. |
+| `num_heads` | — | Number of query heads. |
+| `rotary_interleaved` | `0` | Set to 1 to rotate adjacent pairs instead of using the NeoX half-split; every other value selects the NeoX layout. |
 | `scale` | — | Scale applied to query-key products; omitted or zero uses `1 / sqrt(head_size)`. |
+| `sparse_block_size` | — | Tokens per sparse block; one of 16, 32, 64, 128. |
 ## Type constraints
 ## Files
+- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
 - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 ## Use with `@huggingface/kernels`
+```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.2
+```
+Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

build/webgpu/bench.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "com.microsoft.SparseAttention",
   "fixtureArrays": {
     "block_column_indices_t_pattern": [0, 0, 1, 0, 1, 2, 0, 1, 2, 3, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 5, 0, 1, 2, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7, 8, 0, 2, 3, 4, 5, 6, 7, 8, 9, 0, 3, 4, 5, 6, 7, 8, 9, 10, 0, 4, 5, 6, 7, 8, 9, 10, 11, 0, 4, 5, 6, 7, 8, 9, 10, 11, 12, 0, 4, 6, 7, 8, 9, 10, 11, 12, 13, 0, 4, 7, 8, 9, 10, 11, 12, 13, 14, 0, 4, 8, 9, 10, 11, 12, 13, 14, 15, 0, 0, 1, 0, 1, 2, 0, 1, 2, 3, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 5, 0, 1, 2, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7, 8, 2, 3, 4, 5, 6, 7, 8, 9, 3, 4, 5, 6, 7, 8, 9, 10, 3, 4, 5, 6, 7, 8, 9, 10, 11, 3, 5, 6, 7, 8, 9, 10, 11, 12, 3, 6, 7, 8, 9, 10, 11, 12, 13, 3, 7, 8, 9, 10, 11, 12, 13, 14, 3, 7, 8, 9, 10, 11, 12, 13, 14, 15, -1, -1, -1, -1, -1, -1, 0, 0, 1, 0, 1, 2, 0, 1, 2, 3, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 5, 0, 1, 2, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7, 8, 2, 3, 4, 5, 6, 7, 8, 9, 2, 3, 4, 5, 6, 7, 8, 9, 10, 2, 4, 5, 6, 7, 8, 9, 10, 11, 2, 5, 6, 7, 8, 9, 10, 11, 12, 2, 6, 7, 8, 9, 10, 11, 12, 13, 2, 6, 7, 8, 9, 10, 11, 12, 13, 14, 2, 6, 8, 9, 10, 11, 12, 13, 14, 15, -1, -1, -1, -1, 0, 0, 1, 0, 1, 2, 0, 1, 2, 3, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 5, 0, 1, 2, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7, 8, 1, 2, 3, 4, 5, 6, 7, 8, 9, 1, 3, 4, 5, 6, 7, 8, 9, 10, 1, 4, 5, 6, 7, 8, 9, 10, 11, 1, 5, 6, 7, 8, 9, 10, 11, 12, 1, 5, 6, 7, 8, 9, 10, 11, 12, 13, 1, 5, 7, 8, 9, 10, 11, 12, 13, 14, 1, 5, 8, 9, 10, 11, 12, 13, 14, 15, -1, -1],
     "sparse-prompt-b1-s1024-h32kv8-d128-blk64_input_blockRowIndicesT": [0, 1, 3, 6, 10, 15, 21, 28, 36, 45, 54, 63, 72, 82, 92, 102, 112, 0, 1, 3, 6, 10, 15, 21, 28, 36, 44, 52, 60, 69, 78, 87, 96, 106, 0, 1, 3, 6, 10, 15, 21, 28, 36, 44, 52, 61, 70, 79, 88, 98, 108, 0, 1, 3, 6, 10, 15, 21, 28, 36, 44, 53, 62, 71, 80, 90, 100, 110]
@@ -226,6 +225,703 @@
       },
       "outputs": { "outputT": { "shape": [2, 128, 4096], "dtype": "float32" } },
       "bench": { "metrics": [{ "type": "gflops", "value": "4 * args.qkPairs * args.headDim" }] }
     }
   ]
 }

 {
   "fixtureArrays": {
     "block_column_indices_t_pattern": [0, 0, 1, 0, 1, 2, 0, 1, 2, 3, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 5, 0, 1, 2, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7, 8, 0, 2, 3, 4, 5, 6, 7, 8, 9, 0, 3, 4, 5, 6, 7, 8, 9, 10, 0, 4, 5, 6, 7, 8, 9, 10, 11, 0, 4, 5, 6, 7, 8, 9, 10, 11, 12, 0, 4, 6, 7, 8, 9, 10, 11, 12, 13, 0, 4, 7, 8, 9, 10, 11, 12, 13, 14, 0, 4, 8, 9, 10, 11, 12, 13, 14, 15, 0, 0, 1, 0, 1, 2, 0, 1, 2, 3, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 5, 0, 1, 2, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7, 8, 2, 3, 4, 5, 6, 7, 8, 9, 3, 4, 5, 6, 7, 8, 9, 10, 3, 4, 5, 6, 7, 8, 9, 10, 11, 3, 5, 6, 7, 8, 9, 10, 11, 12, 3, 6, 7, 8, 9, 10, 11, 12, 13, 3, 7, 8, 9, 10, 11, 12, 13, 14, 3, 7, 8, 9, 10, 11, 12, 13, 14, 15, -1, -1, -1, -1, -1, -1, 0, 0, 1, 0, 1, 2, 0, 1, 2, 3, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 5, 0, 1, 2, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7, 8, 2, 3, 4, 5, 6, 7, 8, 9, 2, 3, 4, 5, 6, 7, 8, 9, 10, 2, 4, 5, 6, 7, 8, 9, 10, 11, 2, 5, 6, 7, 8, 9, 10, 11, 12, 2, 6, 7, 8, 9, 10, 11, 12, 13, 2, 6, 7, 8, 9, 10, 11, 12, 13, 14, 2, 6, 8, 9, 10, 11, 12, 13, 14, 15, -1, -1, -1, -1, 0, 0, 1, 0, 1, 2, 0, 1, 2, 3, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 5, 0, 1, 2, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7, 8, 1, 2, 3, 4, 5, 6, 7, 8, 9, 1, 3, 4, 5, 6, 7, 8, 9, 10, 1, 4, 5, 6, 7, 8, 9, 10, 11, 1, 5, 6, 7, 8, 9, 10, 11, 12, 1, 5, 6, 7, 8, 9, 10, 11, 12, 13, 1, 5, 7, 8, 9, 10, 11, 12, 13, 14, 1, 5, 8, 9, 10, 11, 12, 13, 14, 15, -1, -1],
     "sparse-prompt-b1-s1024-h32kv8-d128-blk64_input_blockRowIndicesT": [0, 1, 3, 6, 10, 15, 21, 28, 36, 45, 54, 63, 72, 82, 92, 102, 112, 0, 1, 3, 6, 10, 15, 21, 28, 36, 44, 52, 60, 69, 78, 87, 96, 106, 0, 1, 3, 6, 10, 15, 21, 28, 36, 44, 52, 61, 70, 79, 88, 98, 108, 0, 1, 3, 6, 10, 15, 21, 28, 36, 44, 53, 62, 71, 80, 90, 100, 110]
       },
       "outputs": { "outputT": { "shape": [2, 128, 4096], "dtype": "float32" } },
       "bench": { "metrics": [{ "type": "gflops", "value": "4 * args.qkPairs * args.headDim" }] }
+    },
+    {
+      "name": "boundary-sgmat_capacity_d96_s64_packed0_rotary0",
+      "attrs": { "num_heads": 2, "kv_num_heads": 2, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 64, 192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [2, 64, 192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [2, 64, 192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.017, "scale": 0.25 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.013, "scale": 0.25 }
+        },
+        "blockRowIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [0, 1, 3, 5, 7, 0, 1, 2, 4, 5] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 7],
+          "data": { "kind": "values", "values": [0, 0, 1, 0, 2, 1, 3, 0, 1, 0, 2, 3, 0, 0] }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } },
+        "keyTotalSequenceLengthsT": {
+          "dtype": "int32",
+          "shape": [2],
+          "data": { "kind": "values", "values": [164, 132] }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [2, 64, 192] } },
+      "preset": "model"
+    },
+    {
+      "name": "boundary-sgmat_capacity_d96_s64_packed0_rotary1",
+      "attrs": { "num_heads": 2, "kv_num_heads": 2, "sparse_block_size": 64, "do_rotary": 1, "rotary_interleaved": 0 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 64, 192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [2, 64, 192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [2, 64, 192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.017, "scale": 0.25 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.013, "scale": 0.25 }
+        },
+        "blockRowIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [0, 1, 3, 5, 7, 0, 1, 2, 4, 5] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 7],
+          "data": { "kind": "values", "values": [0, 0, 1, 0, 2, 1, 3, 0, 1, 0, 2, 3, 0, 0] }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } },
+        "keyTotalSequenceLengthsT": {
+          "dtype": "int32",
+          "shape": [2],
+          "data": { "kind": "values", "values": [164, 132] }
+        },
+        "cosCacheT": {
+          "dtype": "float32",
+          "shape": [256, 48],
+          "data": { "kind": "fillFloat32", "sinStep": 0.01, "cosStep": 0.013, "scale": 0.5 }
+        },
+        "sinCacheT": {
+          "dtype": "float32",
+          "shape": [256, 48],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.013, "scale": 0.5 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [2, 64, 192] } },
+      "preset": "model"
+    },
+    {
+      "name": "boundary-sgmat_capacity_d96_s64_packed1_rotary0",
+      "attrs": { "num_heads": 2, "kv_num_heads": 2, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 64, 576],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.017, "scale": 0.25 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.013, "scale": 0.25 }
+        },
+        "blockRowIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [0, 1, 3, 5, 7, 0, 1, 2, 4, 5] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 7],
+          "data": { "kind": "values", "values": [0, 0, 1, 0, 2, 1, 3, 0, 1, 0, 2, 3, 0, 0] }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } },
+        "keyTotalSequenceLengthsT": {
+          "dtype": "int32",
+          "shape": [2],
+          "data": { "kind": "values", "values": [164, 132] }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [2, 64, 192] } },
+      "preset": "model"
+    },
+    {
+      "name": "boundary-sgmat_capacity_d96_s64_packed1_rotary1",
+      "attrs": { "num_heads": 2, "kv_num_heads": 2, "sparse_block_size": 64, "do_rotary": 1, "rotary_interleaved": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 64, 576],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.017, "scale": 0.25 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.013, "scale": 0.25 }
+        },
+        "blockRowIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [0, 1, 3, 5, 7, 0, 1, 2, 4, 5] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 7],
+          "data": { "kind": "values", "values": [0, 0, 1, 0, 2, 1, 3, 0, 1, 0, 2, 3, 0, 0] }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } },
+        "keyTotalSequenceLengthsT": {
+          "dtype": "int32",
+          "shape": [2],
+          "data": { "kind": "values", "values": [164, 132] }
+        },
+        "cosCacheT": {
+          "dtype": "float32",
+          "shape": [256, 48],
+          "data": { "kind": "fillFloat32", "sinStep": 0.01, "cosStep": 0.013, "scale": 0.5 }
+        },
+        "sinCacheT": {
+          "dtype": "float32",
+          "shape": [256, 48],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.013, "scale": 0.5 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [2, 64, 192] } },
+      "preset": "model"
+    },
+    {
+      "name": "boundary-sgmat_capacity_d96_s65_packed0_rotary0",
+      "attrs": { "num_heads": 2, "kv_num_heads": 2, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 65, 192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [2, 65, 192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [2, 65, 192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.017, "scale": 0.25 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.013, "scale": 0.25 }
+        },
+        "blockRowIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [0, 1, 3, 5, 7, 0, 1, 2, 4, 5] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 7],
+          "data": { "kind": "values", "values": [0, 0, 1, 0, 2, 1, 3, 0, 1, 0, 2, 3, 0, 0] }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } },
+        "keyTotalSequenceLengthsT": {
+          "dtype": "int32",
+          "shape": [2],
+          "data": { "kind": "values", "values": [165, 133] }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [2, 65, 192] } },
+      "preset": "model"
+    },
+    {
+      "name": "boundary-sgmat_capacity_d96_s65_packed0_rotary1",
+      "attrs": { "num_heads": 2, "kv_num_heads": 2, "sparse_block_size": 64, "do_rotary": 1, "rotary_interleaved": 0 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 65, 192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [2, 65, 192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [2, 65, 192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.017, "scale": 0.25 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.013, "scale": 0.25 }
+        },
+        "blockRowIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [0, 1, 3, 5, 7, 0, 1, 2, 4, 5] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 7],
+          "data": { "kind": "values", "values": [0, 0, 1, 0, 2, 1, 3, 0, 1, 0, 2, 3, 0, 0] }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } },
+        "keyTotalSequenceLengthsT": {
+          "dtype": "int32",
+          "shape": [2],
+          "data": { "kind": "values", "values": [165, 133] }
+        },
+        "cosCacheT": {
+          "dtype": "float32",
+          "shape": [256, 48],
+          "data": { "kind": "fillFloat32", "sinStep": 0.01, "cosStep": 0.013, "scale": 0.5 }
+        },
+        "sinCacheT": {
+          "dtype": "float32",
+          "shape": [256, 48],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.013, "scale": 0.5 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [2, 65, 192] } },
+      "preset": "model"
+    },
+    {
+      "name": "boundary-sgmat_capacity_d96_s65_packed1_rotary0",
+      "attrs": { "num_heads": 2, "kv_num_heads": 2, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 65, 576],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.017, "scale": 0.25 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.013, "scale": 0.25 }
+        },
+        "blockRowIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [0, 1, 3, 5, 7, 0, 1, 2, 4, 5] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 7],
+          "data": { "kind": "values", "values": [0, 0, 1, 0, 2, 1, 3, 0, 1, 0, 2, 3, 0, 0] }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } },
+        "keyTotalSequenceLengthsT": {
+          "dtype": "int32",
+          "shape": [2],
+          "data": { "kind": "values", "values": [165, 133] }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [2, 65, 192] } },
+      "preset": "model"
+    },
+    {
+      "name": "boundary-sgmat_capacity_d96_s65_packed1_rotary1",
+      "attrs": { "num_heads": 2, "kv_num_heads": 2, "sparse_block_size": 64, "do_rotary": 1, "rotary_interleaved": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 65, 576],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.017, "scale": 0.25 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.013, "scale": 0.25 }
+        },
+        "blockRowIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [0, 1, 3, 5, 7, 0, 1, 2, 4, 5] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 7],
+          "data": { "kind": "values", "values": [0, 0, 1, 0, 2, 1, 3, 0, 1, 0, 2, 3, 0, 0] }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } },
+        "keyTotalSequenceLengthsT": {
+          "dtype": "int32",
+          "shape": [2],
+          "data": { "kind": "values", "values": [165, 133] }
+        },
+        "cosCacheT": {
+          "dtype": "float32",
+          "shape": [256, 48],
+          "data": { "kind": "fillFloat32", "sinStep": 0.01, "cosStep": 0.013, "scale": 0.5 }
+        },
+        "sinCacheT": {
+          "dtype": "float32",
+          "shape": [256, 48],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.013, "scale": 0.5 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [2, 65, 192] } },
+      "preset": "model"
+    },
+    {
+      "name": "boundary-sgmat_prompt_separate_blk64_duplicate_csr",
+      "attrs": { "num_heads": 4, "kv_num_heads": 2, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 192, 256],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 192, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 192, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 2, 192, 64], "data": { "kind": "constant", "value": 0.0 } },
+        "pastValueT": { "dtype": "float32", "shape": [1, 2, 192, 64], "data": { "kind": "constant", "value": 0.0 } },
+        "blockRowIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 4],
+          "data": { "kind": "values", "values": [0, 1, 3, 5, 0, 1, 2, 3] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [0, 0, 1, 1, 1, 0, 1, 2, 0, 0] }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [192] } },
+        "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [192] } }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [1, 192, 256] } },
+      "preset": "model"
+    },
+    {
+      "name": "boundary-sgmat_prompt_tail_s100_explicit_scale",
+      "attrs": { "num_heads": 2, "kv_num_heads": 1, "sparse_block_size": 64, "scale": 0.25 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 100, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 100, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 100, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 1, 128, 32], "data": { "kind": "constant", "value": 0.0 } },
+        "pastValueT": { "dtype": "float32", "shape": [1, 1, 128, 32], "data": { "kind": "constant", "value": 0.0 } },
+        "blockRowIndicesT": { "dtype": "int32", "shape": [1, 3], "data": { "kind": "values", "values": [0, 1, 3] } },
+        "blockColIndicesT": { "dtype": "int32", "shape": [1, 3], "data": { "kind": "values", "values": [0, 0, 1] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [100] } },
+        "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [100] } }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [1, 100, 64] } },
+      "preset": "model"
+    },
+    {
+      "name": "boundary-sgmat_prompt_blk128_subtiles",
+      "attrs": { "num_heads": 2, "kv_num_heads": 1, "sparse_block_size": 128 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 256, 256],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 256, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 256, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 1, 256, 128], "data": { "kind": "constant", "value": 0.0 } },
+        "pastValueT": { "dtype": "float32", "shape": [1, 1, 256, 128], "data": { "kind": "constant", "value": 0.0 } },
+        "blockRowIndicesT": { "dtype": "int32", "shape": [1, 3], "data": { "kind": "values", "values": [0, 1, 2] } },
+        "blockColIndicesT": { "dtype": "int32", "shape": [1, 2], "data": { "kind": "values", "values": [0, 1] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } },
+        "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [1, 256, 256] } },
+      "preset": "model"
+    },
+    {
+      "name": "boundary-sgmat_parallelism_d96_s128",
+      "attrs": { "num_heads": 4, "kv_num_heads": 2, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 128, 384],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [2, 128, 192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [2, 128, 192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.017, "scale": 0.25 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.013, "scale": 0.25 }
+        },
+        "blockRowIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [0, 1, 3, 5, 7, 0, 1, 2, 4, 5] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 7],
+          "data": { "kind": "values", "values": [0, 0, 1, 0, 2, 1, 3, 0, 1, 0, 2, 3, 0, 0] }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } },
+        "keyTotalSequenceLengthsT": {
+          "dtype": "int32",
+          "shape": [2],
+          "data": { "kind": "values", "values": [256, 256] }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [2, 128, 384] } },
+      "preset": "model"
+    },
+    {
+      "name": "boundary-sgmat_parallelism_d96_s256",
+      "attrs": { "num_heads": 4, "kv_num_heads": 2, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 256, 384],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [2, 256, 192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [2, 256, 192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.017, "scale": 0.25 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.013, "scale": 0.25 }
+        },
+        "blockRowIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [0, 1, 3, 5, 7, 0, 1, 2, 4, 5] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 7],
+          "data": { "kind": "values", "values": [0, 0, 1, 0, 2, 1, 3, 0, 1, 0, 2, 3, 0, 0] }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } },
+        "keyTotalSequenceLengthsT": {
+          "dtype": "int32",
+          "shape": [2],
+          "data": { "kind": "values", "values": [256, 256] }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [2, 256, 384] } },
+      "preset": "model"
+    },
+    {
+      "name": "sparse-low-parallelism-s64-h8-d32-past960",
+      "preset": "model",
+      "attrs": { "num_heads": 8, "kv_num_heads": 4, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": { "shape": [1, 64, 256], "dtype": "float32", "dist": "normal", "seed": 9300, "scale": 1 },
+        "keyT": { "shape": [1, 64, 128], "dtype": "float32", "dist": "normal", "seed": 9301, "scale": 1 },
+        "valueT": { "shape": [1, 64, 128], "dtype": "float32", "dist": "normal", "seed": 9302, "scale": 1 },
+        "pastKeyT": { "shape": [1, 4, 1024, 32], "dtype": "float32", "dist": "normal", "seed": 9303, "scale": 1 },
+        "pastValueT": { "shape": [1, 4, 1024, 32], "dtype": "float32", "dist": "normal", "seed": 9304, "scale": 1 },
+        "blockRowIndicesT": {
+          "shape": [4, 17],
+          "dtype": "int32",
+          "data": {
+            "kind": "values",
+            "values": { "$ref": "#/fixtureArrays/sparse-prompt-b1-s1024-h32kv8-d128-blk64_input_blockRowIndicesT" }
+          }
+        },
+        "blockColIndicesT": {
+          "shape": [4, 112],
+          "dtype": "int32",
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/block_column_indices_t_pattern" } }
+        },
+        "totalSequenceLengthT": { "shape": [1], "dtype": "int32", "data": { "kind": "values", "values": [1024] } },
+        "keyTotalSequenceLengthsT": { "shape": [1], "dtype": "int32", "dist": "constant", "value": 1024 }
+      },
+      "outputs": { "outputT": { "shape": [1, 64, 256], "dtype": "float32" } }
+    },
+    {
+      "name": "sparse-low-parallelism-s64-h8-d96-past960",
+      "preset": "model",
+      "attrs": { "num_heads": 8, "kv_num_heads": 4, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": { "shape": [1, 64, 768], "dtype": "float32", "dist": "normal", "seed": 9300, "scale": 1 },
+        "keyT": { "shape": [1, 64, 384], "dtype": "float32", "dist": "normal", "seed": 9301, "scale": 1 },
+        "valueT": { "shape": [1, 64, 384], "dtype": "float32", "dist": "normal", "seed": 9302, "scale": 1 },
+        "pastKeyT": { "shape": [1, 4, 1024, 96], "dtype": "float32", "dist": "normal", "seed": 9303, "scale": 1 },
+        "pastValueT": { "shape": [1, 4, 1024, 96], "dtype": "float32", "dist": "normal", "seed": 9304, "scale": 1 },
+        "blockRowIndicesT": {
+          "shape": [4, 17],
+          "dtype": "int32",
+          "data": {
+            "kind": "values",
+            "values": { "$ref": "#/fixtureArrays/sparse-prompt-b1-s1024-h32kv8-d128-blk64_input_blockRowIndicesT" }
+          }
+        },
+        "blockColIndicesT": {
+          "shape": [4, 112],
+          "dtype": "int32",
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/block_column_indices_t_pattern" } }
+        },
+        "totalSequenceLengthT": { "shape": [1], "dtype": "int32", "data": { "kind": "values", "values": [1024] } },
+        "keyTotalSequenceLengthsT": { "shape": [1], "dtype": "int32", "dist": "constant", "value": 1024 }
+      },
+      "outputs": { "outputT": { "shape": [1, 64, 768], "dtype": "float32" } }
+    },
+    {
+      "name": "sparse-low-parallelism-s64-h8-d128-past960",
+      "preset": "model",
+      "attrs": { "num_heads": 8, "kv_num_heads": 4, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": { "shape": [1, 64, 1024], "dtype": "float32", "dist": "normal", "seed": 9300, "scale": 1 },
+        "keyT": { "shape": [1, 64, 512], "dtype": "float32", "dist": "normal", "seed": 9301, "scale": 1 },
+        "valueT": { "shape": [1, 64, 512], "dtype": "float32", "dist": "normal", "seed": 9302, "scale": 1 },
+        "pastKeyT": { "shape": [1, 4, 1024, 128], "dtype": "float32", "dist": "normal", "seed": 9303, "scale": 1 },
+        "pastValueT": { "shape": [1, 4, 1024, 128], "dtype": "float32", "dist": "normal", "seed": 9304, "scale": 1 },
+        "blockRowIndicesT": {
+          "shape": [4, 17],
+          "dtype": "int32",
+          "data": {
+            "kind": "values",
+            "values": { "$ref": "#/fixtureArrays/sparse-prompt-b1-s1024-h32kv8-d128-blk64_input_blockRowIndicesT" }
+          }
+        },
+        "blockColIndicesT": {
+          "shape": [4, 112],
+          "dtype": "int32",
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/block_column_indices_t_pattern" } }
+        },
+        "totalSequenceLengthT": { "shape": [1], "dtype": "int32", "data": { "kind": "values", "values": [1024] } },
+        "keyTotalSequenceLengthsT": { "shape": [1], "dtype": "int32", "dist": "constant", "value": 1024 }
+      },
+      "outputs": { "outputT": { "shape": [1, 64, 1024], "dtype": "float32" } }
+    },
+    {
+      "name": "sparse-crossover-s128-h8-d96-past128",
+      "preset": "model",
+      "attrs": { "num_heads": 8, "kv_num_heads": 4, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": { "shape": [1, 128, 768], "dtype": "float32", "dist": "normal", "seed": 9300, "scale": 1 },
+        "keyT": { "shape": [1, 128, 384], "dtype": "float32", "dist": "normal", "seed": 9301, "scale": 1 },
+        "valueT": { "shape": [1, 128, 384], "dtype": "float32", "dist": "normal", "seed": 9302, "scale": 1 },
+        "pastKeyT": { "shape": [1, 4, 1024, 96], "dtype": "float32", "dist": "normal", "seed": 9303, "scale": 1 },
+        "pastValueT": { "shape": [1, 4, 1024, 96], "dtype": "float32", "dist": "normal", "seed": 9304, "scale": 1 },
+        "blockRowIndicesT": {
+          "shape": [4, 17],
+          "dtype": "int32",
+          "data": {
+            "kind": "values",
+            "values": { "$ref": "#/fixtureArrays/sparse-prompt-b1-s1024-h32kv8-d128-blk64_input_blockRowIndicesT" }
+          }
+        },
+        "blockColIndicesT": {
+          "shape": [4, 112],
+          "dtype": "int32",
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/block_column_indices_t_pattern" } }
+        },
+        "totalSequenceLengthT": { "shape": [1], "dtype": "int32", "data": { "kind": "values", "values": [256] } },
+        "keyTotalSequenceLengthsT": { "shape": [1], "dtype": "int32", "dist": "constant", "value": 256 }
+      },
+      "outputs": { "outputT": { "shape": [1, 128, 768], "dtype": "float32" } }
     }
   ]
 }

build/webgpu/manifest.json CHANGED Viewed

@@ -2,109 +2,31 @@
   "domain": "com.microsoft",
   "name": "SparseAttention",
   "sinceVersion": 1,
-  "description": "Block-sparse causal attention used by Phi-3-small. `block_row_indices` and `block_col_indices` encode one or more CSR block masks, and layouts cycle over query heads. Grouped-query heads, separate or packed `[Q|K|V]`, explicit scaling, partial or full rotary embedding in NeoX or interleaved layout, and float16 are supported. The past/present key and value tensors share allocations and are updated in place. Head sizes must be non-zero multiples of 8; bfloat16 is not implemented.",
-  "inputs": [
-    {
-      "role": "query",
-      "dtype": "T",
-      "rank": 3,
-      "description": "Query `(batch_size, sequence_length, num_heads * head_size)`, or packed `[Q|K|V]` `(batch_size, sequence_length, (num_heads + 2 * kv_num_heads) * head_size)` when `key` and `value` are omitted."
-    },
-    {
-      "role": "key",
-      "dtype": "T",
-      "rank": 3,
-      "optional": true,
-      "description": "Key `(batch_size, sequence_length, kv_num_heads * head_size)`. Omitted for packed QKV."
-    },
-    {
-      "role": "value",
-      "dtype": "T",
-      "rank": 3,
-      "optional": true,
-      "description": "Value `(batch_size, sequence_length, kv_num_heads * head_size)`. Omitted for packed QKV."
-    },
-    {
-      "role": "past_key",
-      "dtype": "T",
-      "rank": 4,
-      "description": "Key cache `(batch_size, kv_num_heads, max_cache_sequence_length, head_size)`, updated in place."
-    },
-    {
-      "role": "past_value",
-      "dtype": "T",
-      "rank": 4,
-      "description": "Value cache with the same shape as `past_key`, updated in place."
-    },
-    {
-      "role": "block_row_indices",
-      "dtype": "M",
-      "rank": 2,
-      "description": "CSR row pointers `(num_layout, max_blocks + 1)`. Each layout starts at zero, is monotonically non-decreasing, and ends no later than that layout's `block_col_indices` width."
-    },
-    {
-      "role": "block_col_indices",
-      "dtype": "M",
-      "rank": 2,
-      "description": "CSR column indices `(num_layout, max_nnz_blocks)`, right-padded past each layout's non-zero count. Every active entry is in `[0, max_blocks)`."
-    },
-    {
-      "role": "total_sequence_length",
-      "dtype": "M",
-      "description": "Scalar or one-element vector holding the maximum total key length. Equal to `sequence_length` exactly in the prompt case, which is how the past length is decided. The value fits the cache, the sparse layout's `max_blocks * sparse_block_size` capacity, and the rotary-cache row count when rotary is enabled."
-    },
-    {
-      "role": "key_total_sequence_lengths",
-      "dtype": "M",
-      "rank": 1,
-      "description": "Per-batch total key length excluding padding, shape `(batch_size)`. Each value is at most `total_sequence_length` and is at least 1 for a prompt or at least `sequence_length` otherwise."
-    },
-    {
-      "role": "cos_cache",
-      "dtype": "T",
-      "rank": 2,
-      "optional": true,
-      "description": "Rotary cosine cache `(max_rotary_sequence_length, rotary_dimension / 2)`, where the width is a multiple of 8 no larger than `head_size / 2`. Required with `sin_cache` when `do_rotary` is 1."
-    },
-    {
-      "role": "sin_cache",
-      "dtype": "T",
-      "rank": 2,
-      "optional": true,
-      "description": "Rotary sine cache with the same shape as `cos_cache`; required with it when `do_rotary` is 1."
-    }
-  ],
-  "outputs": [
-    {
-      "role": "output",
-      "dtype": "T",
-      "rank": 3,
-      "shape": "[batchSize, seqLen, numHeads * headSize]",
-      "description": "Attention output `(batch_size, sequence_length, num_heads * head_size)`."
-    },
-    {
-      "role": "past_key",
-      "dtype": "T",
-      "rank": 4,
-      "shape": "shapes.pastKeyT",
-      "description": "The key cache tensor itself after the in-place append; ONNX names this output `present_key`."
-    },
-    {
-      "role": "past_value",
-      "dtype": "T",
-      "rank": 4,
-      "shape": "shapes.pastValueT",
-      "description": "The value cache tensor itself after the in-place append; ONNX names this output `present_value`."
-    }
-  ],
-  "attributes": { "do_rotary": 0, "rotary_interleaved": 0 },
-  "attributeDescriptions": {
-    "num_heads": "Number of query heads.",
-    "kv_num_heads": "Number of key/value heads; must divide `num_heads`.",
-    "sparse_block_size": "Tokens per sparse block; one of 16, 32, 64, 128.",
-    "do_rotary": "Set to 1 to apply rotary embedding to Q and to K before it enters the cache; every other value disables rotary embedding.",
-    "rotary_interleaved": "Set to 1 to rotate adjacent pairs instead of using the NeoX half-split; every other value selects the NeoX layout.",
-    "scale": "Scale applied to query-key products; omitted or zero uses `1 / sqrt(head_size)`."
   },
   "attributeConstraints": {
     "num_heads": { "required": true },
@@ -112,31 +34,13 @@
     "sparse_block_size": { "required": true }
   },
   "typeConstraints": { "T": ["float32", "float16"], "M": ["int32"] },
-  "args": {
-    "queryT": { "kind": "tensor", "semantic": "query", "role": "input" },
-    "keyT": { "kind": "tensor", "semantic": "key", "role": "input", "required": false },
-    "valueT": { "kind": "tensor", "semantic": "value", "role": "input", "required": false },
-    "pastKeyT": { "kind": "tensor", "semantic": "past_key", "role": "inout" },
-    "pastValueT": { "kind": "tensor", "semantic": "past_value", "role": "inout" },
-    "blockRowIndicesT": { "kind": "tensor", "semantic": "block_row_indices", "role": "input", "dtype": "int32" },
-    "blockColIndicesT": { "kind": "tensor", "semantic": "block_col_indices", "role": "input", "dtype": "int32" },
-    "totalSequenceLengthT": { "kind": "tensor", "semantic": "total_sequence_length", "role": "input", "dtype": "int32" },
-    "keyTotalSequenceLengthsT": {
-      "kind": "tensor",
-      "semantic": "key_total_sequence_lengths",
-      "role": "input",
-      "dtype": "int32"
-    },
-    "cosCacheT": { "kind": "tensor", "semantic": "cos_cache", "role": "input", "required": false },
-    "sinCacheT": { "kind": "tensor", "semantic": "sin_cache", "role": "input", "required": false },
-    "outputT": { "kind": "tensor", "semantic": "output", "role": "output" }
-  },
   "tunables": {
-    "WORKGROUP_SIZE": 128,
-    "APPEND_WORKGROUP_SIZE": 256,
-    "NARROW_MIN_WORKGROUPS": 1024,
-    "QUERY_TILE": 4,
-    "V_STAGE_MAX_WORKGROUPS": 512
   },
   "derive": {
     "wave32Adapter": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 32 and device.adapterInfo.subgroupMaxSize == 32",
@@ -178,7 +82,7 @@
     "rotaryPairOk": "not doRotary or (present.cosCacheT and present.sinCacheT and ranks.cosCacheT == 2 and ranks.sinCacheT == 2 and rotaryHalf % 8 == 0 and rotaryDim <= headSize and sameShape(shapes.sinCacheT, shapes.cosCacheT) and tensorDtypes.cosCacheT == tensorDtypes.queryT and tensorDtypes.sinCacheT == tensorDtypes.queryT)",
     "blockIndexShapeOk": "ranks.blockRowIndicesT == 2 and ranks.blockColIndicesT == 2 and dim(shapes.blockColIndicesT, 0) == numLayout and maxBlocks >= 1 and maxNnz >= 0 and maxNnz <= maxBlocks * maxBlocks and tensorDtypes.blockRowIndicesT == \"int32\" and tensorDtypes.blockColIndicesT == \"int32\"",
     "scheduleShapeOk": "(ranks.totalSequenceLengthT == 0 or ranks.totalSequenceLengthT == 1) and numel(shapes.totalSequenceLengthT) == 1 and ranks.keyTotalSequenceLengthsT == 1 and dim(shapes.keyTotalSequenceLengthsT, 0) == batchSize and tensorDtypes.totalSequenceLengthT == \"int32\" and tensorDtypes.keyTotalSequenceLengthsT == \"int32\"",
-    "geometryOk": "tunables.WORKGROUP_SIZE >= 1 and floor(tunables.WORKGROUP_SIZE) == tunables.WORKGROUP_SIZE and pow2ceil(tunables.WORKGROUP_SIZE) == tunables.WORKGROUP_SIZE and tunables.WORKGROUP_SIZE <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.WORKGROUP_SIZE <= device.limits.maxComputeWorkgroupSizeX and tunables.APPEND_WORKGROUP_SIZE >= 1 and floor(tunables.APPEND_WORKGROUP_SIZE) == tunables.APPEND_WORKGROUP_SIZE and tunables.APPEND_WORKGROUP_SIZE <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.APPEND_WORKGROUP_SIZE <= device.limits.maxComputeWorkgroupSizeX and sparseQueryTiles <= device.limits.maxComputeWorkgroupsPerDimension and batchSize * numHeads <= device.limits.maxComputeWorkgroupsPerDimension and ceilDiv(ceilDiv(qRotaryElements, tunables.APPEND_WORKGROUP_SIZE), device.limits.maxComputeWorkgroupsPerDimension) <= device.limits.maxComputeWorkgroupsPerDimension and (2 * sparseQueryTile * headSize + (3 * sparseQueryTile + 1) * sparseAttnWorkgroup) * 4 <= device.limits.maxComputeWorkgroupStorageSize and sparseAttnWorkgroup <= device.limits.maxComputeInvocationsPerWorkgroup and sparseAttnWorkgroup <= device.limits.maxComputeWorkgroupSizeX",
     "contract": "ranks.queryT == 3 and ranks.outputT == 3 and (tensorDtypes.queryT == \"float32\" or tensorDtypes.queryT == \"float16\") and f16Ok(dtypes.T) and tensorDtypes.pastKeyT == tensorDtypes.queryT and tensorDtypes.pastValueT == tensorDtypes.queryT and tensorDtypes.outputT == tensorDtypes.queryT and numHeads >= 1 and kvNumHeads >= 1 and numHeads % kvNumHeads == 0 and headSize >= 8 and headSize % 8 == 0 and (not doRotary or headSize % 16 == 0) and numLayout >= 1 and numHeads % numLayout == 0 and (sparseBlockSize == 16 or sparseBlockSize == 32 or sparseBlockSize == 64 or sparseBlockSize == 128) and cacheShapeOk and queryShapeOk and kvShapeOk and kvPairOk and rotaryPairOk and blockIndexShapeOk and scheduleShapeOk and dim(shapes.outputT, 0) == batchSize and dim(shapes.outputT, 1) == seqLen and dim(shapes.outputT, 2) == qHidden",
     "packedContract": "contract and packedQkv and not useRotary",
     "packedRotaryContract": "contract and packedQkv and useRotary",
@@ -186,622 +90,112 @@
     "separateRotaryContract": "contract and not packedQkv and useRotary",
     "sparseVStageWorthIt": "sparseQueryTiles * batchSize * numHeads <= tunables.V_STAGE_MAX_WORKGROUPS",
     "sgmatQueryTiles": "ceilDiv(seqLen, 64)",
-    "sparseSgmatLdsBytes": "(64 * 32 + 64 * 64 + 64 * 2 + 128 * 2) * 4",
-    "sparseSgmatGeometryOk": "256 <= device.limits.maxComputeInvocationsPerWorkgroup and 256 <= device.limits.maxComputeWorkgroupSizeX and sgmatQueryTiles <= device.limits.maxComputeWorkgroupsPerDimension and batchSize * numHeads <= device.limits.maxComputeWorkgroupsPerDimension and sparseSgmatLdsBytes <= device.limits.maxComputeWorkgroupStorageSize",
-    "sparseSgmatOk": "tensorDtypes.queryT == \"float32\" and seqLen >= 64 and sparseBlockSize % 64 == 0 and headSize % 32 == 0 and headSize <= 128 and maxCacheSeq % 64 == 0 and device.features.has(\"subgroups\") and wave32Effective and device.features.has(\"chromium-experimental-subgroup-matrix\") and sparseSgmatGeometryOk"
-  },
-  "constants": {
     "scalar": "dtypes.T",
-    "usesF16": "dtypes.T == \"f16\"",
-    "numHeads": "numHeads",
-    "kvNumHeads": "kvNumHeads",
-    "headSize": "headSize",
-    "headVec": "headVec",
     "cacheVec": "\"vec4<f16>\" if dtypes.T == \"f16\" else \"vec4<f32>\"",
     "attnWorkgroup": "sparseAttnWorkgroup",
-    "maxCacheSeq": "maxCacheSeq",
-    "sparseBlockSize": "sparseBlockSize",
-    "numLayout": "numLayout",
-    "maxBlocks": "maxBlocks",
-    "maxNnz": "maxNnz",
-    "packedStride": "packedStride",
-    "packedQkv": "packedQkv",
     "usesRotary": "useRotary",
-    "rotaryHalf": "rotaryHalf",
-    "rotaryDim": "rotaryDim",
-    "rotaryInterleaved": "rotaryInterleaved",
     "appendWorkgroupSize": "tunables.APPEND_WORKGROUP_SIZE"
   },
-  "bindingSets": {
-    "appendSeparate": [
-      {
-        "name": "new_key",
-        "arg": "keyT",
-        "semantic": "key",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "new_value",
-        "arg": "valueT",
-        "semantic": "value",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "present_key",
-        "arg": "pastKeyT",
-        "semantic": "past_key",
-        "buffer": { "type": "storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "present_value",
-        "arg": "pastValueT",
-        "semantic": "past_value",
-        "buffer": { "type": "storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "key_total_sequence_lengths",
-        "arg": "keyTotalSequenceLengthsT",
-        "semantic": "key_total_sequence_lengths",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "total_sequence_length",
-        "arg": "totalSequenceLengthT",
-        "semantic": "total_sequence_length",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "batchSize", "type": "u32", "value": "batchSize" },
-            { "name": "seqLen", "type": "u32", "value": "seqLen" }
-          ]
-        }
-      }
-    ],
-    "appendSeparateRotary": [
-      {
-        "name": "new_key",
-        "arg": "keyT",
-        "semantic": "key",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "new_value",
-        "arg": "valueT",
-        "semantic": "value",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "present_key",
-        "arg": "pastKeyT",
-        "semantic": "past_key",
-        "buffer": { "type": "storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "present_value",
-        "arg": "pastValueT",
-        "semantic": "past_value",
-        "buffer": { "type": "storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "key_total_sequence_lengths",
-        "arg": "keyTotalSequenceLengthsT",
-        "semantic": "key_total_sequence_lengths",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "total_sequence_length",
-        "arg": "totalSequenceLengthT",
-        "semantic": "total_sequence_length",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "cos_cache",
-        "arg": "cosCacheT",
-        "semantic": "cos_cache",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "sin_cache",
-        "arg": "sinCacheT",
-        "semantic": "sin_cache",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "batchSize", "type": "u32", "value": "batchSize" },
-            { "name": "seqLen", "type": "u32", "value": "seqLen" }
-          ]
-        }
-      }
-    ],
-    "appendPacked": [
-      {
-        "name": "packed_qkv",
-        "arg": "queryT",
-        "semantic": "query",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "present_key",
-        "arg": "pastKeyT",
-        "semantic": "past_key",
-        "buffer": { "type": "storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "present_value",
-        "arg": "pastValueT",
-        "semantic": "past_value",
-        "buffer": { "type": "storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "key_total_sequence_lengths",
-        "arg": "keyTotalSequenceLengthsT",
-        "semantic": "key_total_sequence_lengths",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "total_sequence_length",
-        "arg": "totalSequenceLengthT",
-        "semantic": "total_sequence_length",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "batchSize", "type": "u32", "value": "batchSize" },
-            { "name": "seqLen", "type": "u32", "value": "seqLen" }
-          ]
-        }
-      }
-    ],
-    "appendPackedRotary": [
-      {
-        "name": "packed_qkv",
-        "arg": "queryT",
-        "semantic": "query",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "present_key",
-        "arg": "pastKeyT",
-        "semantic": "past_key",
-        "buffer": { "type": "storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "present_value",
-        "arg": "pastValueT",
-        "semantic": "past_value",
-        "buffer": { "type": "storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "key_total_sequence_lengths",
-        "arg": "keyTotalSequenceLengthsT",
-        "semantic": "key_total_sequence_lengths",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "total_sequence_length",
-        "arg": "totalSequenceLengthT",
-        "semantic": "total_sequence_length",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "cos_cache",
-        "arg": "cosCacheT",
-        "semantic": "cos_cache",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "sin_cache",
-        "arg": "sinCacheT",
-        "semantic": "sin_cache",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "batchSize", "type": "u32", "value": "batchSize" },
-            { "name": "seqLen", "type": "u32", "value": "seqLen" }
-          ]
-        }
-      }
-    ],
-    "qRotary": [
-      {
-        "name": "query",
-        "arg": "queryT",
-        "semantic": "query",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "cos_cache",
-        "arg": "cosCacheT",
-        "semantic": "cos_cache",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "sin_cache",
-        "arg": "sinCacheT",
-        "semantic": "sin_cache",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      { "name": "q_rotary", "semantic": "QRotary", "buffer": { "type": "storage" }, "elementType": "f32" },
-      {
-        "name": "key_total_sequence_lengths",
-        "arg": "keyTotalSequenceLengthsT",
-        "semantic": "key_total_sequence_lengths",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "total_sequence_length",
-        "arg": "totalSequenceLengthT",
-        "semantic": "total_sequence_length",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "batchSize", "type": "u32", "value": "batchSize" },
-            { "name": "seqLen", "type": "u32", "value": "seqLen" }
-          ]
-        }
-      }
-    ],
-    "attentionDirect": [
-      {
-        "name": "query",
-        "arg": "queryT",
-        "semantic": "query",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "present_key",
-        "arg": "pastKeyT",
-        "semantic": "past_key",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$cacheVec"
-      },
-      {
-        "name": "present_value",
-        "arg": "pastValueT",
-        "semantic": "past_value",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$cacheVec"
-      },
-      {
-        "name": "block_row_indices",
-        "arg": "blockRowIndicesT",
-        "semantic": "block_row_indices",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "block_col_indices",
-        "arg": "blockColIndicesT",
-        "semantic": "block_col_indices",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "key_total_sequence_lengths",
-        "arg": "keyTotalSequenceLengthsT",
-        "semantic": "key_total_sequence_lengths",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "total_sequence_length",
-        "arg": "totalSequenceLengthT",
-        "semantic": "total_sequence_length",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "output",
-        "arg": "outputT",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "seqLen", "type": "u32", "value": "seqLen" },
-            { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" }
-          ]
-        }
-      }
-    ],
-    "attentionSgmat": [
-      {
-        "name": "query",
-        "arg": "queryT",
-        "semantic": "query",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "present_key",
-        "arg": "pastKeyT",
-        "semantic": "past_key",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "present_value",
-        "arg": "pastValueT",
-        "semantic": "past_value",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "block_row_indices",
-        "arg": "blockRowIndicesT",
-        "semantic": "block_row_indices",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "block_col_indices",
-        "arg": "blockColIndicesT",
-        "semantic": "block_col_indices",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "key_total_sequence_lengths",
-        "arg": "keyTotalSequenceLengthsT",
-        "semantic": "key_total_sequence_lengths",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "total_sequence_length",
-        "arg": "totalSequenceLengthT",
-        "semantic": "total_sequence_length",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "output",
-        "arg": "outputT",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "seqLen", "type": "u32", "value": "seqLen" },
-            { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" }
-          ]
-        }
-      }
-    ],
-    "attentionRotary": [
-      { "name": "q_rotary", "semantic": "QRotary", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
-      {
-        "name": "present_key",
-        "arg": "pastKeyT",
-        "semantic": "past_key",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$cacheVec"
-      },
-      {
-        "name": "present_value",
-        "arg": "pastValueT",
-        "semantic": "past_value",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$cacheVec"
-      },
-      {
-        "name": "block_row_indices",
-        "arg": "blockRowIndicesT",
-        "semantic": "block_row_indices",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "block_col_indices",
-        "arg": "blockColIndicesT",
-        "semantic": "block_col_indices",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "key_total_sequence_lengths",
-        "arg": "keyTotalSequenceLengthsT",
-        "semantic": "key_total_sequence_lengths",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "total_sequence_length",
-        "arg": "totalSequenceLengthT",
-        "semantic": "total_sequence_length",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "output",
-        "arg": "outputT",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "seqLen", "type": "u32", "value": "seqLen" },
-            { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" }
-          ]
-        }
-      }
-    ],
-    "attentionSgmatRotary": [
-      { "name": "q_rotary", "semantic": "QRotary", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
-      {
-        "name": "present_key",
-        "arg": "pastKeyT",
-        "semantic": "past_key",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "present_value",
-        "arg": "pastValueT",
-        "semantic": "past_value",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "block_row_indices",
-        "arg": "blockRowIndicesT",
-        "semantic": "block_row_indices",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "block_col_indices",
-        "arg": "blockColIndicesT",
-        "semantic": "block_col_indices",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "key_total_sequence_lengths",
-        "arg": "keyTotalSequenceLengthsT",
-        "semantic": "key_total_sequence_lengths",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "total_sequence_length",
-        "arg": "totalSequenceLengthT",
-        "semantic": "total_sequence_length",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "output",
-        "arg": "outputT",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "seqLen", "type": "u32", "value": "seqLen" },
-            { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" }
-          ]
-        }
-      }
-    ]
   },
   "variants": [
     {
       "id": "separate",
       "priority": 0,
-      "when": ["separateContract", "geometryOk"],
       "passes": [
         {
           "id": "append",
           "name": "SparseAttention.Append",
           "shader": "sparse-kv-append.wgsl.jinja",
-          "bindings": "appendSeparate",
-          "constants": { "kvSource": "\"new_key\"", "vSource": "\"new_value\"" },
           "dispatch": {
-            "threads": "batchSize * kvNumHeads * seqLen * headSize",
-            "workgroupSize": "constants.appendWorkgroupSize"
           }
         },
         {
           "id": "attention",
           "name": "SparseAttention.Attention",
           "shader": "sparse-attention.wgsl.jinja",
-          "bindings": "attentionDirect",
-          "dispatch": { "x": "sparseQueryTiles", "y": "batchSize * numHeads" },
-          "constants": { "qTile": "sparseQueryTile" }
         }
-      ],
-      "constants": { "vStageWorthIt": "sparseVStageWorthIt" }
     },
     {
       "id": "separate_sgmat",
       "priority": 20,
-      "when": ["separateContract", "geometryOk", "sparseSgmatOk"],
       "requires": {
         "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
         "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
@@ -811,64 +205,69 @@
           "id": "append",
           "name": "SparseAttention.Append",
           "shader": "sparse-kv-append.wgsl.jinja",
-          "bindings": "appendSeparate",
-          "constants": { "kvSource": "\"new_key\"", "vSource": "\"new_value\"" },
           "dispatch": {
-            "threads": "batchSize * kvNumHeads * seqLen * headSize",
-            "workgroupSize": "constants.appendWorkgroupSize"
           }
         },
         {
           "id": "attention",
           "name": "SparseAttention.AttentionSgmat",
           "shader": "sparse-attention-sgmat.wgsl.jinja",
-          "bindings": "attentionSgmat",
-          "dispatch": { "x": "sgmatQueryTiles", "y": "batchSize * numHeads" }
         }
-      ]
     },
     {
       "id": "separate_rotary",
       "priority": 10,
-      "when": ["separateRotaryContract", "geometryOk"],
       "intermediates": [{ "id": "QRotary", "dtype": "float32", "shape": "[qRotaryElements]" }],
       "passes": [
         {
           "id": "append",
           "name": "SparseAttention.Append",
           "shader": "sparse-kv-append.wgsl.jinja",
-          "bindings": "appendSeparateRotary",
-          "constants": { "kvSource": "\"new_key\"", "vSource": "\"new_value\"" },
           "dispatch": {
-            "threads": "batchSize * kvNumHeads * seqLen * headSize",
-            "workgroupSize": "constants.appendWorkgroupSize"
           }
         },
         {
           "id": "qrotary",
           "name": "SparseAttention.QueryRotary",
           "shader": "sparse-q-rotary.wgsl.jinja",
-          "bindings": "qRotary",
           "dispatch": {
-            "threads": "batchSize * numHeads * seqLen * headSize",
-            "workgroupSize": "constants.appendWorkgroupSize"
           }
         },
         {
           "id": "attention",
           "name": "SparseAttention.Attention",
           "shader": "sparse-attention.wgsl.jinja",
-          "bindings": "attentionRotary",
-          "dispatch": { "x": "sparseQueryTiles", "y": "batchSize * numHeads" },
-          "constants": { "qTile": "sparseQueryTile" }
         }
-      ],
-      "constants": { "vStageWorthIt": "sparseVStageWorthIt" }
     },
     {
       "id": "separate_rotary_sgmat",
       "priority": 30,
-      "when": ["separateRotaryContract", "geometryOk", "sparseSgmatOk"],
       "requires": {
         "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
         "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
@@ -879,63 +278,68 @@
           "id": "append",
           "name": "SparseAttention.Append",
           "shader": "sparse-kv-append.wgsl.jinja",
-          "bindings": "appendSeparateRotary",
-          "constants": { "kvSource": "\"new_key\"", "vSource": "\"new_value\"" },
           "dispatch": {
-            "threads": "batchSize * kvNumHeads * seqLen * headSize",
-            "workgroupSize": "constants.appendWorkgroupSize"
           }
         },
         {
           "id": "qrotary",
           "name": "SparseAttention.QueryRotary",
           "shader": "sparse-q-rotary.wgsl.jinja",
-          "bindings": "qRotary",
           "dispatch": {
-            "threads": "batchSize * numHeads * seqLen * headSize",
-            "workgroupSize": "constants.appendWorkgroupSize"
           }
         },
         {
           "id": "attention",
           "name": "SparseAttention.AttentionSgmat",
           "shader": "sparse-attention-sgmat.wgsl.jinja",
-          "bindings": "attentionSgmatRotary",
-          "dispatch": { "x": "sgmatQueryTiles", "y": "batchSize * numHeads" }
         }
-      ]
     },
     {
       "id": "packed",
       "priority": 0,
-      "when": ["packedContract", "geometryOk"],
       "passes": [
         {
           "id": "append",
           "name": "SparseAttention.Append",
           "shader": "sparse-kv-append.wgsl.jinja",
-          "bindings": "appendPacked",
-          "constants": { "kvSource": "\"packed_qkv\"", "vSource": "\"packed_qkv\"" },
           "dispatch": {
-            "threads": "batchSize * kvNumHeads * seqLen * headSize",
-            "workgroupSize": "constants.appendWorkgroupSize"
           }
         },
         {
           "id": "attention",
           "name": "SparseAttention.Attention",
           "shader": "sparse-attention.wgsl.jinja",
-          "bindings": "attentionDirect",
-          "dispatch": { "x": "sparseQueryTiles", "y": "batchSize * numHeads" },
-          "constants": { "qTile": "sparseQueryTile" }
         }
-      ],
-      "constants": { "vStageWorthIt": "sparseVStageWorthIt" }
     },
     {
       "id": "packed_sgmat",
       "priority": 20,
-      "when": ["packedContract", "geometryOk", "sparseSgmatOk"],
       "requires": {
         "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
         "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
@@ -945,64 +349,69 @@
           "id": "append",
           "name": "SparseAttention.Append",
           "shader": "sparse-kv-append.wgsl.jinja",
-          "bindings": "appendPacked",
-          "constants": { "kvSource": "\"packed_qkv\"", "vSource": "\"packed_qkv\"" },
           "dispatch": {
-            "threads": "batchSize * kvNumHeads * seqLen * headSize",
-            "workgroupSize": "constants.appendWorkgroupSize"
           }
         },
         {
           "id": "attention",
           "name": "SparseAttention.AttentionSgmat",
           "shader": "sparse-attention-sgmat.wgsl.jinja",
-          "bindings": "attentionSgmat",
-          "dispatch": { "x": "sgmatQueryTiles", "y": "batchSize * numHeads" }
         }
-      ]
     },
     {
       "id": "packed_rotary",
       "priority": 10,
-      "when": ["packedRotaryContract", "geometryOk"],
       "intermediates": [{ "id": "QRotary", "dtype": "float32", "shape": "[qRotaryElements]" }],
       "passes": [
         {
           "id": "append",
           "name": "SparseAttention.Append",
           "shader": "sparse-kv-append.wgsl.jinja",
-          "bindings": "appendPackedRotary",
-          "constants": { "kvSource": "\"packed_qkv\"", "vSource": "\"packed_qkv\"" },
           "dispatch": {
-            "threads": "batchSize * kvNumHeads * seqLen * headSize",
-            "workgroupSize": "constants.appendWorkgroupSize"
           }
         },
         {
           "id": "qrotary",
           "name": "SparseAttention.QueryRotary",
           "shader": "sparse-q-rotary.wgsl.jinja",
-          "bindings": "qRotary",
           "dispatch": {
-            "threads": "batchSize * numHeads * seqLen * headSize",
-            "workgroupSize": "constants.appendWorkgroupSize"
           }
         },
         {
           "id": "attention",
           "name": "SparseAttention.Attention",
           "shader": "sparse-attention.wgsl.jinja",
-          "bindings": "attentionRotary",
-          "dispatch": { "x": "sparseQueryTiles", "y": "batchSize * numHeads" },
-          "constants": { "qTile": "sparseQueryTile" }
         }
-      ],
-      "constants": { "vStageWorthIt": "sparseVStageWorthIt" }
     },
     {
       "id": "packed_rotary_sgmat",
       "priority": 30,
-      "when": ["packedRotaryContract", "geometryOk", "sparseSgmatOk"],
       "requires": {
         "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
         "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
@@ -1013,31 +422,35 @@
           "id": "append",
           "name": "SparseAttention.Append",
           "shader": "sparse-kv-append.wgsl.jinja",
-          "bindings": "appendPackedRotary",
-          "constants": { "kvSource": "\"packed_qkv\"", "vSource": "\"packed_qkv\"" },
           "dispatch": {
-            "threads": "batchSize * kvNumHeads * seqLen * headSize",
-            "workgroupSize": "constants.appendWorkgroupSize"
           }
         },
         {
           "id": "qrotary",
           "name": "SparseAttention.QueryRotary",
           "shader": "sparse-q-rotary.wgsl.jinja",
-          "bindings": "qRotary",
           "dispatch": {
-            "threads": "batchSize * numHeads * seqLen * headSize",
-            "workgroupSize": "constants.appendWorkgroupSize"
           }
         },
         {
           "id": "attention",
           "name": "SparseAttention.AttentionSgmat",
           "shader": "sparse-attention-sgmat.wgsl.jinja",
-          "bindings": "attentionSgmatRotary",
-          "dispatch": { "x": "sgmatQueryTiles", "y": "batchSize * numHeads" }
         }
-      ]
     }
   ]
 }

   "domain": "com.microsoft",
   "name": "SparseAttention",
   "sinceVersion": 1,
+  "inputs": {
+    "queryT": { "onnx": "query", "dtype": "T", "rank": 3 },
+    "keyT": { "onnx": "key", "dtype": "T", "rank": 3, "optional": true },
+    "valueT": { "onnx": "value", "dtype": "T", "rank": 3, "optional": true },
+    "pastKeyT": { "onnx": "past_key", "dtype": "T", "rank": 4 },
+    "pastValueT": { "onnx": "past_value", "dtype": "T", "rank": 4 },
+    "blockRowIndicesT": { "onnx": "block_row_indices", "dtype": "M", "rank": 2, "storage": "int32" },
+    "blockColIndicesT": { "onnx": "block_col_indices", "dtype": "M", "rank": 2, "storage": "int32" },
+    "totalSequenceLengthT": { "onnx": "total_sequence_length", "dtype": "M", "storage": "int32" },
+    "keyTotalSequenceLengthsT": { "onnx": "key_total_sequence_lengths", "dtype": "M", "rank": 1, "storage": "int32" },
+    "cosCacheT": { "onnx": "cos_cache", "dtype": "T", "rank": 2, "optional": true },
+    "sinCacheT": { "onnx": "sin_cache", "dtype": "T", "rank": 2, "optional": true }
+  },
+  "outputs": {
+    "outputT": { "onnx": "output", "dtype": "T", "rank": 3, "shape": "[batchSize, seqLen, numHeads * headSize]" },
+    "pastKeyT": { "onnx": "past_key", "dtype": "T", "rank": 4, "shape": "shapes.pastKeyT" },
+    "pastValueT": { "onnx": "past_value", "dtype": "T", "rank": 4, "shape": "shapes.pastValueT" }
+  },
+  "attributes": {
+    "do_rotary": { "default": 0 },
+    "rotary_interleaved": { "default": 0 },
+    "num_heads": {},
+    "kv_num_heads": {},
+    "sparse_block_size": {},
+    "scale": {}
   },
   "attributeConstraints": {
     "num_heads": { "required": true },
     "sparse_block_size": { "required": true }
   },
   "typeConstraints": { "T": ["float32", "float16"], "M": ["int32"] },
   "tunables": {
+    "WORKGROUP_SIZE": { "default": 128 },
+    "APPEND_WORKGROUP_SIZE": { "default": 256 },
+    "NARROW_MIN_WORKGROUPS": { "default": 1024 },
+    "QUERY_TILE": { "default": 4 },
+    "V_STAGE_MAX_WORKGROUPS": { "default": 512 },
+    "MATRIX_MIN_WORKGROUPS": { "default": 16 }
   },
   "derive": {
     "wave32Adapter": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 32 and device.adapterInfo.subgroupMaxSize == 32",
     "rotaryPairOk": "not doRotary or (present.cosCacheT and present.sinCacheT and ranks.cosCacheT == 2 and ranks.sinCacheT == 2 and rotaryHalf % 8 == 0 and rotaryDim <= headSize and sameShape(shapes.sinCacheT, shapes.cosCacheT) and tensorDtypes.cosCacheT == tensorDtypes.queryT and tensorDtypes.sinCacheT == tensorDtypes.queryT)",
     "blockIndexShapeOk": "ranks.blockRowIndicesT == 2 and ranks.blockColIndicesT == 2 and dim(shapes.blockColIndicesT, 0) == numLayout and maxBlocks >= 1 and maxNnz >= 0 and maxNnz <= maxBlocks * maxBlocks and tensorDtypes.blockRowIndicesT == \"int32\" and tensorDtypes.blockColIndicesT == \"int32\"",
     "scheduleShapeOk": "(ranks.totalSequenceLengthT == 0 or ranks.totalSequenceLengthT == 1) and numel(shapes.totalSequenceLengthT) == 1 and ranks.keyTotalSequenceLengthsT == 1 and dim(shapes.keyTotalSequenceLengthsT, 0) == batchSize and tensorDtypes.totalSequenceLengthT == \"int32\" and tensorDtypes.keyTotalSequenceLengthsT == \"int32\"",
+    "geometryOk": "tunables.WORKGROUP_SIZE >= 1 and floor(tunables.WORKGROUP_SIZE) == tunables.WORKGROUP_SIZE and pow2ceil(tunables.WORKGROUP_SIZE) == tunables.WORKGROUP_SIZE and tunables.WORKGROUP_SIZE <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.WORKGROUP_SIZE <= device.limits.maxComputeWorkgroupSizeX and tunables.APPEND_WORKGROUP_SIZE >= 1 and floor(tunables.APPEND_WORKGROUP_SIZE) == tunables.APPEND_WORKGROUP_SIZE and tunables.APPEND_WORKGROUP_SIZE <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.APPEND_WORKGROUP_SIZE <= device.limits.maxComputeWorkgroupSizeX and sparseQueryTiles <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and batchSize * numHeads <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(ceilDiv(qRotaryElements, tunables.APPEND_WORKGROUP_SIZE), min(device.limits.maxComputeWorkgroupsPerDimension, 65535)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and (2 * sparseQueryTile * headSize + (3 * sparseQueryTile + 1) * sparseAttnWorkgroup) * 4 <= device.limits.maxComputeWorkgroupStorageSize and sparseAttnWorkgroup <= device.limits.maxComputeInvocationsPerWorkgroup and sparseAttnWorkgroup <= device.limits.maxComputeWorkgroupSizeX",
     "contract": "ranks.queryT == 3 and ranks.outputT == 3 and (tensorDtypes.queryT == \"float32\" or tensorDtypes.queryT == \"float16\") and f16Ok(dtypes.T) and tensorDtypes.pastKeyT == tensorDtypes.queryT and tensorDtypes.pastValueT == tensorDtypes.queryT and tensorDtypes.outputT == tensorDtypes.queryT and numHeads >= 1 and kvNumHeads >= 1 and numHeads % kvNumHeads == 0 and headSize >= 8 and headSize % 8 == 0 and (not doRotary or headSize % 16 == 0) and numLayout >= 1 and numHeads % numLayout == 0 and (sparseBlockSize == 16 or sparseBlockSize == 32 or sparseBlockSize == 64 or sparseBlockSize == 128) and cacheShapeOk and queryShapeOk and kvShapeOk and kvPairOk and rotaryPairOk and blockIndexShapeOk and scheduleShapeOk and dim(shapes.outputT, 0) == batchSize and dim(shapes.outputT, 1) == seqLen and dim(shapes.outputT, 2) == qHidden",
     "packedContract": "contract and packedQkv and not useRotary",
     "packedRotaryContract": "contract and packedQkv and useRotary",
     "separateRotaryContract": "contract and not packedQkv and useRotary",
     "sparseVStageWorthIt": "sparseQueryTiles * batchSize * numHeads <= tunables.V_STAGE_MAX_WORKGROUPS",
     "sgmatQueryTiles": "ceilDiv(seqLen, 64)",
+    "sgmatDirectQuery": "seqLen % 64 == 0",
+    "sparseSgmatTileN": "64 if (64 * 32 + 64 * 64 + 64 * 2 + 128 * 2) * 4 <= device.limits.maxComputeWorkgroupStorageSize else 32",
+    "sparseSgmatTileK": "sparseSgmatTileN / 2",
+    "sparseSgmatLdsBytes": "(64 * sparseSgmatTileK + 64 * sparseSgmatTileN + 64 * 2 + 128 * 2) * 4",
+    "sparseSgmatGeometryOk": "256 <= device.limits.maxComputeInvocationsPerWorkgroup and 256 <= device.limits.maxComputeWorkgroupSizeX and sgmatQueryTiles <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and batchSize * numHeads <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and sparseSgmatLdsBytes <= device.limits.maxComputeWorkgroupStorageSize",
+    "sparseSgmatOk": "tensorDtypes.queryT == \"float32\" and seqLen >= 64 and sparseBlockSize % 64 == 0 and headSize % 32 == 0 and headSize <= 128 and maxCacheSeq % 64 == 0 and device.features.has(\"subgroups\") and wave32Effective and device.features.has(\"chromium-experimental-subgroup-matrix\") and sparseSgmatGeometryOk",
     "scalar": "dtypes.T",
     "cacheVec": "\"vec4<f16>\" if dtypes.T == \"f16\" else \"vec4<f32>\"",
     "attnWorkgroup": "sparseAttnWorkgroup",
     "usesRotary": "useRotary",
     "appendWorkgroupSize": "tunables.APPEND_WORKGROUP_SIZE"
   },
+  "when": ["geometryOk"],
+  "bindings": {
+    "new_key": { "arg": "keyT", "buffer": "read-only-storage", "elementType": "$scalar" },
+    "new_value": { "arg": "valueT", "buffer": "read-only-storage", "elementType": "$scalar" },
+    "present_key": { "arg": "pastKeyT", "buffer": "storage", "elementType": "$scalar" },
+    "present_value": { "arg": "pastValueT", "buffer": "storage", "elementType": "$scalar" },
+    "key_total_sequence_lengths": {
+      "arg": "keyTotalSequenceLengthsT",
+      "buffer": "read-only-storage",
+      "elementType": "i32"
+    },
+    "total_sequence_length": { "arg": "totalSequenceLengthT", "buffer": "read-only-storage", "elementType": "i32" },
+    "params": {
+      "buffer": "uniform",
+      "struct": [
+        { "name": "batchSize", "type": "u32", "value": "batchSize" },
+        { "name": "seqLen", "type": "u32", "value": "seqLen" }
+      ]
+    },
+    "cos_cache": { "arg": "cosCacheT", "buffer": "read-only-storage", "elementType": "$scalar" },
+    "sin_cache": { "arg": "sinCacheT", "buffer": "read-only-storage", "elementType": "$scalar" },
+    "packed_qkv": { "arg": "queryT", "buffer": "read-only-storage", "elementType": "$scalar" },
+    "query": { "arg": "queryT", "buffer": "read-only-storage", "elementType": "$scalar" },
+    "present_key_2": {
+      "arg": "pastKeyT",
+      "name": "present_key",
+      "buffer": "read-only-storage",
+      "elementType": "$cacheVec"
+    },
+    "present_value_2": {
+      "arg": "pastValueT",
+      "name": "present_value",
+      "buffer": "read-only-storage",
+      "elementType": "$cacheVec"
+    },
+    "block_row_indices": { "arg": "blockRowIndicesT", "buffer": "read-only-storage", "elementType": "i32" },
+    "block_col_indices": { "arg": "blockColIndicesT", "buffer": "read-only-storage", "elementType": "i32" },
+    "output": { "arg": "outputT", "buffer": "storage", "elementType": "$scalar" },
+    "params_2": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "seqLen", "type": "u32", "value": "seqLen" },
+        { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" }
+      ]
+    },
+    "q_rotary": { "scratch": "QRotary", "buffer": "read-only-storage", "elementType": "f32" },
+    "present_key_3": {
+      "arg": "pastKeyT",
+      "name": "present_key",
+      "buffer": "read-only-storage",
+      "elementType": "$scalar"
+    },
+    "present_value_3": {
+      "arg": "pastValueT",
+      "name": "present_value",
+      "buffer": "read-only-storage",
+      "elementType": "$scalar"
+    },
+    "q_rotary_2": { "scratch": "QRotary", "name": "q_rotary", "buffer": "storage", "elementType": "f32" }
   },
   "variants": [
     {
       "id": "separate",
       "priority": 0,
+      "when": ["separateContract"],
+      "derive": { "vStageWorthIt": "sparseVStageWorthIt" },
       "passes": [
         {
           "id": "append",
           "name": "SparseAttention.Append",
           "shader": "sparse-kv-append.wgsl.jinja",
+          "derive": { "kvSource": "\"new_key\"", "vSource": "\"new_value\"" },
+          "bindings": ["new_key", "new_value", "present_key", "present_value", "key_total_sequence_lengths", "total_sequence_length", "params"],
           "dispatch": {
+            "x": "min(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "z": 1
           }
         },
         {
           "id": "attention",
           "name": "SparseAttention.Attention",
           "shader": "sparse-attention.wgsl.jinja",
+          "derive": { "qTile": "sparseQueryTile" },
+          "bindings": ["query", "present_key_2", "present_value_2", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_2"],
+          "dispatch": { "x": "sparseQueryTiles", "y": "batchSize * numHeads" }
         }
+      ]
     },
     {
       "id": "separate_sgmat",
       "priority": 20,
+      "when": ["separateContract", "sparseSgmatOk"],
       "requires": {
         "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
         "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
           "id": "append",
           "name": "SparseAttention.Append",
           "shader": "sparse-kv-append.wgsl.jinja",
+          "derive": { "kvSource": "\"new_key\"", "vSource": "\"new_value\"" },
+          "bindings": ["new_key", "new_value", "present_key", "present_value", "key_total_sequence_lengths", "total_sequence_length", "params"],
           "dispatch": {
+            "x": "min(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "z": 1
           }
         },
         {
           "id": "attention",
           "name": "SparseAttention.AttentionSgmat",
           "shader": "sparse-attention-sgmat.wgsl.jinja",
+          "bindings": ["query", "present_key_3", "present_value_3", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_2"],
+          "dispatch": { "x": "sgmatQueryTiles", "y": "batchSize * numHeads" },
+          "subgroupCollectivesWidth": 32
         }
+      ],
+      "demoteWhen": ["sgmatQueryTiles * batchSize * numHeads * sparseSgmatTileN < tunables.MATRIX_MIN_WORKGROUPS * 64"]
     },
     {
       "id": "separate_rotary",
       "priority": 10,
+      "when": ["separateRotaryContract"],
+      "derive": { "vStageWorthIt": "sparseVStageWorthIt" },
       "intermediates": [{ "id": "QRotary", "dtype": "float32", "shape": "[qRotaryElements]" }],
       "passes": [
         {
           "id": "append",
           "name": "SparseAttention.Append",
           "shader": "sparse-kv-append.wgsl.jinja",
+          "derive": { "kvSource": "\"new_key\"", "vSource": "\"new_value\"" },
+          "bindings": ["new_key", "new_value", "present_key", "present_value", "key_total_sequence_lengths", "total_sequence_length", "cos_cache", "sin_cache", "params"],
           "dispatch": {
+            "x": "min(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "z": 1
           }
         },
         {
           "id": "qrotary",
           "name": "SparseAttention.QueryRotary",
           "shader": "sparse-q-rotary.wgsl.jinja",
+          "bindings": ["query", "cos_cache", "sin_cache", "q_rotary_2", "key_total_sequence_lengths", "total_sequence_length", "params"],
           "dispatch": {
+            "x": "min(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "z": 1
           }
         },
         {
           "id": "attention",
           "name": "SparseAttention.Attention",
           "shader": "sparse-attention.wgsl.jinja",
+          "derive": { "qTile": "sparseQueryTile" },
+          "bindings": ["q_rotary", "present_key_2", "present_value_2", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_2"],
+          "dispatch": { "x": "sparseQueryTiles", "y": "batchSize * numHeads" }
         }
+      ]
     },
     {
       "id": "separate_rotary_sgmat",
       "priority": 30,
+      "when": ["separateRotaryContract", "sparseSgmatOk"],
       "requires": {
         "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
         "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
           "id": "append",
           "name": "SparseAttention.Append",
           "shader": "sparse-kv-append.wgsl.jinja",
+          "derive": { "kvSource": "\"new_key\"", "vSource": "\"new_value\"" },
+          "bindings": ["new_key", "new_value", "present_key", "present_value", "key_total_sequence_lengths", "total_sequence_length", "cos_cache", "sin_cache", "params"],
           "dispatch": {
+            "x": "min(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "z": 1
           }
         },
         {
           "id": "qrotary",
           "name": "SparseAttention.QueryRotary",
           "shader": "sparse-q-rotary.wgsl.jinja",
+          "bindings": ["query", "cos_cache", "sin_cache", "q_rotary_2", "key_total_sequence_lengths", "total_sequence_length", "params"],
           "dispatch": {
+            "x": "min(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "z": 1
           }
         },
         {
           "id": "attention",
           "name": "SparseAttention.AttentionSgmat",
           "shader": "sparse-attention-sgmat.wgsl.jinja",
+          "bindings": ["q_rotary", "present_key_3", "present_value_3", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_2"],
+          "dispatch": { "x": "sgmatQueryTiles", "y": "batchSize * numHeads" },
+          "subgroupCollectivesWidth": 32
         }
+      ],
+      "demoteWhen": ["sgmatQueryTiles * batchSize * numHeads * sparseSgmatTileN < tunables.MATRIX_MIN_WORKGROUPS * 64"]
     },
     {
       "id": "packed",
       "priority": 0,
+      "when": ["packedContract"],
+      "derive": { "vStageWorthIt": "sparseVStageWorthIt" },
       "passes": [
         {
           "id": "append",
           "name": "SparseAttention.Append",
           "shader": "sparse-kv-append.wgsl.jinja",
+          "derive": { "kvSource": "\"packed_qkv\"", "vSource": "\"packed_qkv\"" },
+          "bindings": ["packed_qkv", "present_key", "present_value", "key_total_sequence_lengths", "total_sequence_length", "params"],
           "dispatch": {
+            "x": "min(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "z": 1
           }
         },
         {
           "id": "attention",
           "name": "SparseAttention.Attention",
           "shader": "sparse-attention.wgsl.jinja",
+          "derive": { "qTile": "sparseQueryTile" },
+          "bindings": ["query", "present_key_2", "present_value_2", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_2"],
+          "dispatch": { "x": "sparseQueryTiles", "y": "batchSize * numHeads" }
         }
+      ]
     },
     {
       "id": "packed_sgmat",
       "priority": 20,
+      "when": ["packedContract", "sparseSgmatOk"],
       "requires": {
         "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
         "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
           "id": "append",
           "name": "SparseAttention.Append",
           "shader": "sparse-kv-append.wgsl.jinja",
+          "derive": { "kvSource": "\"packed_qkv\"", "vSource": "\"packed_qkv\"" },
+          "bindings": ["packed_qkv", "present_key", "present_value", "key_total_sequence_lengths", "total_sequence_length", "params"],
           "dispatch": {
+            "x": "min(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "z": 1
           }
         },
         {
           "id": "attention",
           "name": "SparseAttention.AttentionSgmat",
           "shader": "sparse-attention-sgmat.wgsl.jinja",
+          "bindings": ["query", "present_key_3", "present_value_3", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_2"],
+          "dispatch": { "x": "sgmatQueryTiles", "y": "batchSize * numHeads" },
+          "subgroupCollectivesWidth": 32
         }
+      ],
+      "demoteWhen": ["sgmatQueryTiles * batchSize * numHeads * sparseSgmatTileN < tunables.MATRIX_MIN_WORKGROUPS * 64"]
     },
     {
       "id": "packed_rotary",
       "priority": 10,
+      "when": ["packedRotaryContract"],
+      "derive": { "vStageWorthIt": "sparseVStageWorthIt" },
       "intermediates": [{ "id": "QRotary", "dtype": "float32", "shape": "[qRotaryElements]" }],
       "passes": [
         {
           "id": "append",
           "name": "SparseAttention.Append",
           "shader": "sparse-kv-append.wgsl.jinja",
+          "derive": { "kvSource": "\"packed_qkv\"", "vSource": "\"packed_qkv\"" },
+          "bindings": ["packed_qkv", "present_key", "present_value", "key_total_sequence_lengths", "total_sequence_length", "cos_cache", "sin_cache", "params"],
           "dispatch": {
+            "x": "min(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "z": 1
           }
         },
         {
           "id": "qrotary",
           "name": "SparseAttention.QueryRotary",
           "shader": "sparse-q-rotary.wgsl.jinja",
+          "bindings": ["query", "cos_cache", "sin_cache", "q_rotary_2", "key_total_sequence_lengths", "total_sequence_length", "params"],
           "dispatch": {
+            "x": "min(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "z": 1
           }
         },
         {
           "id": "attention",
           "name": "SparseAttention.Attention",
           "shader": "sparse-attention.wgsl.jinja",
+          "derive": { "qTile": "sparseQueryTile" },
+          "bindings": ["q_rotary", "present_key_2", "present_value_2", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_2"],
+          "dispatch": { "x": "sparseQueryTiles", "y": "batchSize * numHeads" }
         }
+      ]
     },
     {
       "id": "packed_rotary_sgmat",
       "priority": 30,
+      "when": ["packedRotaryContract", "sparseSgmatOk"],
       "requires": {
         "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
         "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
           "id": "append",
           "name": "SparseAttention.Append",
           "shader": "sparse-kv-append.wgsl.jinja",
+          "derive": { "kvSource": "\"packed_qkv\"", "vSource": "\"packed_qkv\"" },
+          "bindings": ["packed_qkv", "present_key", "present_value", "key_total_sequence_lengths", "total_sequence_length", "cos_cache", "sin_cache", "params"],
           "dispatch": {
+            "x": "min(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "z": 1
           }
         },
         {
           "id": "qrotary",
           "name": "SparseAttention.QueryRotary",
           "shader": "sparse-q-rotary.wgsl.jinja",
+          "bindings": ["query", "cos_cache", "sin_cache", "q_rotary_2", "key_total_sequence_lengths", "total_sequence_length", "params"],
           "dispatch": {
+            "x": "min(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "z": 1
           }
         },
         {
           "id": "attention",
           "name": "SparseAttention.AttentionSgmat",
           "shader": "sparse-attention-sgmat.wgsl.jinja",
+          "bindings": ["q_rotary", "present_key_3", "present_value_3", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_2"],
+          "dispatch": { "x": "sgmatQueryTiles", "y": "batchSize * numHeads" },
+          "subgroupCollectivesWidth": 32
         }
+      ],
+      "demoteWhen": ["sgmatQueryTiles * batchSize * numHeads * sparseSgmatTileN < tunables.MATRIX_MIN_WORKGROUPS * 64"]
     }
   ]
 }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,21 +1,33 @@
 {
   "name": "com.microsoft.SparseAttention",
-  "id": "_com_microsoft_sparseattention_webgpu_4a03a84",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "E0HHOiKzV0TA0/6aDZXLjmxd2z/Q5iYFBFGNYeQoQM0=",
-      "manifest.json": "gKMlEB+YkQi6M1AeKLoodBrq6h8O9hZGRuRd1MoYsak=",
-      "sparse-attention-sgmat.wgsl.jinja": "OKpr/4DCSDNAm1vzXOoPqlMGauOYb3Lw3cymXu+dXWI=",
-      "sparse-attention.wgsl.jinja": "Piza+iXd87OglOJjKORQmvbF1St7TzSgN37dCuv36K4=",
-      "sparse-kv-append.wgsl.jinja": "RavNXJuGESPqe/Jzt5Ek6+844s5q2w/ocyUgGP3cWec=",
-      "sparse-q-rotary.wgsl.jinja": "hI7d9cx6RqRO079h4xPBMTMj+qKqIUOKW3pF3vdeWdg=",
-      "test.json": "BTrN5cYcQUbwYPpPG9o1rqyvAiCPN2ePxNF+VuaQUmk="
     }
   },
-  "provenance": { "kernel": { "sha": "c928d21e6cc1310861cba3bafb75f5f679ecf5f3", "dirty": false } },
-  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.SparseAttention" }
 }

 {
   "name": "com.microsoft.SparseAttention",
+  "id": "_com_microsoft_sparseattention_webgpu_9e87250",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "nMASpXh23ARMsnca9ti0EsTpERZGnOaFX7OI9QEKmcc=",
+      "manifest.json": "pGu41bfBMFQF8i3CC/V4Elx5QtkSWNUtxk4Z6+gSFsM=",
+      "sparse-attention-sgmat.wgsl.jinja": "At5cbiFuunZkmu32JwzR5KDqNs6/D2VfJb89Op7HljM=",
+      "sparse-attention.wgsl.jinja": "4VYhFW1x7bbvfGHKcFwP/sptJwhPShWMSe383RasMEY=",
+      "sparse-kv-append.wgsl.jinja": "BHKpoS1Ekd524fT8Ft3A7OBCaCXcCe8VpJzXy5q9lA4=",
+      "sparse-q-rotary.wgsl.jinja": "1eKB4VVovC4N6CeIJW+S27RUk8VMkDUSYpHp/BXTfRY=",
+      "test.json": "lURO8tA6ovXlCiXrg3GjYlxvgcn4OdJ/A1EJ93voYq4="
     }
   },
+  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.0",
+    "variants": {
+      "separate": ["sparse-attention.wgsl.jinja", "sparse-kv-append.wgsl.jinja"],
+      "separate_sgmat": ["sparse-attention-sgmat.wgsl.jinja", "sparse-kv-append.wgsl.jinja"],
+      "separate_rotary": ["sparse-attention.wgsl.jinja", "sparse-kv-append.wgsl.jinja", "sparse-q-rotary.wgsl.jinja"],
+      "separate_rotary_sgmat": ["sparse-attention-sgmat.wgsl.jinja", "sparse-kv-append.wgsl.jinja", "sparse-q-rotary.wgsl.jinja"],
+      "packed": ["sparse-attention.wgsl.jinja", "sparse-kv-append.wgsl.jinja"],
+      "packed_sgmat": ["sparse-attention-sgmat.wgsl.jinja", "sparse-kv-append.wgsl.jinja"],
+      "packed_rotary": ["sparse-attention.wgsl.jinja", "sparse-kv-append.wgsl.jinja", "sparse-q-rotary.wgsl.jinja"],
+      "packed_rotary_sgmat": ["sparse-attention-sgmat.wgsl.jinja", "sparse-kv-append.wgsl.jinja", "sparse-q-rotary.wgsl.jinja"]
+    }
+  }
 }

build/webgpu/sparse-attention-sgmat.wgsl.jinja CHANGED Viewed

@@ -1,3 +1,4 @@
 {% macro sparse_schedule() %}
 // How much history precedes this call's tokens. Prompt mode is determined from
 // scalar total_sequence_length rather than per-batch values and forces the past
@@ -18,32 +19,19 @@ enable subgroup_size_control;
 enable chromium_experimental_subgroup_matrix;
 diagnostic(off, chromium.subgroup_matrix_uniformity);
 {{ env.wgsl.resourceDeclarations }}
-// com.microsoft.SparseAttention, subgroup-matrix attention pass (prefill band).
-// One workgroup per (batch, 64-query tile, query head). The block-sparse layout
-// already lists dense SPARSE_BLOCK-wide key tiles, so each selected block is
-// walked as 64x64 score tiles on 8x8 subgroup matrices with no gather step: the
-// K and V rows of a block are contiguous rows of the head-major cache and load
-// directly as matrix fragments. Only Q goes through shared memory, for the
-// query-tail zero padding.
-//
-// The pass runs the selected blocks twice. Sweep one folds each score tile
-// into running per-row (max, denominator) softmax statistics; sweep two
-// recomputes the same tiles (bit-identical: same operands, same order),
-// exponentiates against the now-final row max, and accumulates P.V into
-// subgroup-matrix result fragments. Recomputing QK costs half the ideal FLOPs
-// again, but it is what lets both GEMMs stay on the matrix unit: a running
-// online rescale of the output fragments would need an elementwise multiply on
-// opaque fragments, which only exists as a shared-memory round trip per key
-// tile, and the output accumulator alone (64 x head_size f32) is the whole
-// workgroup storage budget.
 //
-// Causality, CSR duplicate suppression, the dense-row shortcut, and the
-// all-masked-row uniform average reproduce the online kernel's semantics; a
-// key past a query's bound contributes probability zero and multiplies away
-// against whatever the cache holds there, exactly as the online kernel's
-// value walk does.
 const Q_HEADS: u32 = {{ numHeads }}u;
 const KV_HEADS: u32 = {{ kvNumHeads }}u;
 const HEAD_DIM: u32 = {{ headSize }}u;
@@ -56,15 +44,13 @@ const Q_HIDDEN: u32 = {{ numHeads * headSize }}u;
 {% if not usesRotary %}
 const Q_STRIDE: u32 = {{ packedStride if packedQkv else numHeads * headSize }}u;
 {% endif %}
-// 64 query rows x 64 key columns per score tile, staged over the head
-// dimension in 32-wide chunks; 8 subgroups of 32 lanes as a 4x2 grid, each
-// owning a 16x32 score sub-tile. sparse_block_size here is a multiple of 64
-// and head_size a multiple of 32, so every key tile is a whole aligned run of
-// cache rows and the head loop has no tail.
 const TILE_M: u32 = 64u;
-const TILE_N: u32 = 64u;
-const TILE_K: u32 = 32u;
-const SUB_TILES: u32 = {{ (sparseBlockSize / 64) | int }}u;
 {% set pvColBlocks = (headSize / 16) | int %}
 // Each of the two subgroup column groups owns half the output head dimension.
 const PV_COL_SPAN: u32 = {{ (headSize / 2) | int }}u;
@@ -93,11 +79,11 @@ fn exp_shift(value: f32, maxValue: f32) -> f32 {
 }
 // Q staging for the score GEMM; the score epilogues alias it as the
-// fragment-store scratch (8 subgroups x 4 banks x 64 elements).
-var<workgroup> tile_q: array<f32, 2048>;
 // Tile probabilities for the P.V GEMM; the output epilogue aliases it as the
 // result-fragment scratch once the last key tile's readers are done.
-var<workgroup> prob_tile: array<f32, 4096>;
 var<workgroup> row_m: array<f32, 64>;
 var<workgroup> row_d: array<f32, 64>;
 // Per-key-tile row partials, one slot per (row, subgroup column group).
@@ -113,15 +99,28 @@ fn scale_value() -> f32 {
 {{ sparse_schedule() }}
 {% macro score_tile() %}
-        // S = Q.K^T for this key tile, accumulated over the head dimension in
-        // TILE_K chunks. Q is staged with query tails zero-padded; K loads
-        // straight from the cache as column-major right fragments.
         for (var k_base = 0u; k_base < HEAD_DIM; k_base += TILE_K) {
           {
             let a_row = li / 4u;
-            let a_col = (li % 4u) * 8u;
-            for (var i = 0u; i < 8u; i++) {
               let k = k_base + a_col + i;
 {% if usesRotary %}
               let q_index = ((batch * Q_HEADS + head) * params.seqLen
@@ -137,26 +136,29 @@ fn scale_value() -> f32 {
             }
           }
           workgroupBarrier();
           for (var step = 0u; step < TILE_K; step += 8u) {
 {% for rb in range(2) %}
             let mat_a{{ rb }} = subgroupMatrixLoad<subgroup_matrix_left<f32, 8, 8>, row_major>(
-              &tile_q, (base_a + {{ rb * 8 }}u) * TILE_K + step, TILE_K
             );
 {% endfor %}
-{% for cb in range(4) %}
             let mat_b{{ cb }} = subgroupMatrixLoad<subgroup_matrix_right<f32, 8, 8>, col_major>(
               &present_key, key_row_elem + (base_b + {{ cb * 8 }}u) * HEAD_DIM + k_base + step, HEAD_DIM
             );
 {% endfor %}
 {% for rb in range(2) %}
-{% for cb in range(4) %}
             mat_s{{ rb }}{{ cb }} = subgroupMatrixMultiplyAccumulate(
               mat_a{{ rb }}, mat_b{{ cb }}, mat_s{{ rb }}{{ cb }}
             );
 {% endfor %}
 {% endfor %}
           }
           workgroupBarrier();
         }
 {% endmacro %}
@@ -167,8 +169,8 @@ fn scale_value() -> f32 {
   for (var mask_row = row_first; mask_row <= row_last; mask_row++) {
     let start = u32(block_row_indices[row_base + mask_row]);
     let end = u32(block_row_indices[row_base + mask_row + 1u]);
-    // A row listing every causal block is also dense, so its column values are
-    // ignored, matching the online kernel and the pinned dense-layout fixture.
     let dense_row = layout_is_dense || (end - start == mask_row + 1u);
     // Furthest causal bound (exclusive) of the tile's queries on this row.
     let bound_max = min(q_abs0 + rows_live, (mask_row + 1u) * SPARSE_BLOCK);
@@ -197,7 +199,7 @@ fn scale_value() -> f32 {
       if (tile_ok && key_base < bound_max) {
         let key_row_elem = ((batch * KV_HEADS + kv_head) * MAX_CACHE_SEQ + key_base) * HEAD_DIM;
 {% for rb in range(2) %}
-{% for cb in range(4) %}
         var mat_s{{ rb }}{{ cb }}: subgroup_matrix_result<f32, 8, 8>;
 {% endfor %}
 {% endfor %}
@@ -215,13 +217,13 @@ fn scale_value() -> f32 {
         var tile_stat_m{{ rb }} = -FLT_MAX;
         var tile_stat_d{{ rb }} = 0.0;
 {% endif %}
-{% for cb in range(4) %}
         subgroupMatrixStore<row_major>(
-          &tile_q, (subgroup * 4u + {{ cb }}u) * 64u, mat_s{{ rb }}{{ cb }}, 8u
         );
 {% endfor %}
         workgroupBarrier();
-{% for cb in range(4) %}
         for (var pair = 0u; pair < 2u; pair++) {
           let r = base_a + {{ rb * 8 }}u + row_in_block;
           let key = key_base + base_b + {{ cb * 8 }}u + col_in_block + pair;
@@ -230,7 +232,7 @@ fn scale_value() -> f32 {
 {% if phase == "stats" %}
           if (allowed) {
             let scored = tile_q[
-              (subgroup * 4u + {{ cb }}u) * 64u + row_in_block * 8u + col_in_block + pair
             ] * scale;
             let prev_m = tile_stat_m{{ rb }};
             tile_stat_m{{ rb }} = max(tile_stat_m{{ rb }}, scored);
@@ -241,7 +243,7 @@ fn scale_value() -> f32 {
           var prob = 0.0;
           if (allowed) {
             prob = exp_shift(tile_q[
-              (subgroup * 4u + {{ cb }}u) * 64u + row_in_block * 8u + col_in_block + pair
             ] * scale, row_m[r]);
           }
           prob_tile[r * TILE_N + base_b + {{ cb * 8 }}u + col_in_block + pair] = prob;
@@ -348,7 +350,7 @@ fn main(
   let subtile_idy = subgroup % 4u;
   let subtile_idx = subgroup / 4u;
   let base_a = subtile_idy * 16u;
-  let base_b = subtile_idx * 32u;
   let row_in_block = lane / 4u;
   let col_in_block = (lane % 4u) * 2u;
@@ -374,30 +376,34 @@ fn main(
 {{ sweep("stats") }}
 {{ sweep("apply") }}
-  // Normalize by the final denominators and store; the result fragments
-  // round-trip through prob_tile, which no reader needs any more.
 {% for rb in range(2) %}
-{% if rb > 0 %}
   workgroupBarrier();
 {% endif %}
-{% for cb in range(pvColBlocks) %}
   subgroupMatrixStore<row_major>(
-    &prob_tile, (subgroup * {{ pvColBlocks }}u + {{ cb }}u) * 64u, mat_o{{ rb }}{{ cb }}, 8u
   );
 {% endfor %}
   workgroupBarrier();
-{% for cb in range(pvColBlocks) %}
   for (var pair = 0u; pair < 2u; pair++) {
     let r = base_a + {{ rb * 8 }}u + row_in_block;
     if (r < rows_live && row_d[r] > 0.0) {
       let dcol = subtile_idx * PV_COL_SPAN + {{ cb * 8 }}u + col_in_block + pair;
       let out_index = (batch * params.seqLen + tile0 + r) * Q_HIDDEN + head * HEAD_DIM + dcol;
       output[out_index] = {{ scalar }}(prob_tile[
-        (subgroup * {{ pvColBlocks }}u + {{ cb }}u) * 64u + row_in_block * 8u + col_in_block + pair
       ] / row_d[r]);
     }
   }
 {% endfor %}
 {% endfor %}
   // An all-masked row becomes a uniform average over causal history rather

+{% set scoreColBlocks = (sparseSgmatTileN / 16) | int %}
 {% macro sparse_schedule() %}
 // How much history precedes this call's tokens. Prompt mode is determined from
 // scalar total_sequence_length rather than per-batch values and forces the past
 enable chromium_experimental_subgroup_matrix;
 diagnostic(off, chromium.subgroup_matrix_uniformity);
 {{ env.wgsl.resourceDeclarations }}
+// Subgroup-matrix attention over 64-query tiles. Each workgroup processes one
+// `(batch, query tile, query head)` tuple. Selected sparse blocks are traversed
+// as 64-query score tiles; their contiguous K/V cache rows load directly as matrix
+// fragments. Sequences of complete query tiles also load Q directly; sequences
+// with a partial final tile stage Q with zero padding. The manifest bounds tile widths by workgroup storage.
 //
+// The first sweep folds score tiles into per-row `(max, denominator)` softmax
+// statistics. The second recomputes the same scores, applies the completed
+// normalization, and accumulates P.V into result fragments. Causal bounds,
+// duplicate CSR columns, dense rows, and all-masked rows are handled explicitly.
 const Q_HEADS: u32 = {{ numHeads }}u;
 const KV_HEADS: u32 = {{ kvNumHeads }}u;
 const HEAD_DIM: u32 = {{ headSize }}u;
 {% if not usesRotary %}
 const Q_STRIDE: u32 = {{ packedStride if packedQkv else numHeads * headSize }}u;
 {% endif %}
+// Eight 32-lane subgroups form a 4x2 grid. The device storage budget chooses
+// 64 or 32 key columns and a matching head-dimension staging width. Both divide
+// the admitted sparse-block and head dimensions without a key or head tail.
 const TILE_M: u32 = 64u;
+const TILE_N: u32 = {{ sparseSgmatTileN }}u;
+const TILE_K: u32 = {{ sparseSgmatTileK }}u;
+const SUB_TILES: u32 = {{ (sparseBlockSize / sparseSgmatTileN) | int }}u;
 {% set pvColBlocks = (headSize / 16) | int %}
 // Each of the two subgroup column groups owns half the output head dimension.
 const PV_COL_SPAN: u32 = {{ (headSize / 2) | int }}u;
 }
 // Q staging for the score GEMM; the score epilogues alias it as the
+// fragment-store scratch (8 subgroups x scoreColBlocks banks x 64 elements).
+var<workgroup> tile_q: array<f32, {{ 64 * sparseSgmatTileK }}>;
 // Tile probabilities for the P.V GEMM; the output epilogue aliases it as the
 // result-fragment scratch once the last key tile's readers are done.
+var<workgroup> prob_tile: array<f32, {{ 64 * sparseSgmatTileN }}>;
 var<workgroup> row_m: array<f32, 64>;
 var<workgroup> row_d: array<f32, 64>;
 // Per-key-tile row partials, one slot per (row, subgroup column group).
 {{ sparse_schedule() }}
+{% set queryMatrixSource = "tile_q" if not sgmatDirectQuery else ("q_rotary" if usesRotary else "query") %}
+{% set queryMatrixStride = "TILE_K" if not sgmatDirectQuery else ("HEAD_DIM" if usesRotary else "Q_STRIDE") %}
+{% macro query_matrix_offset(rb) -%}
+{% if not sgmatDirectQuery -%}
+(base_a + {{ rb * 8 }}u) * TILE_K + step
+{%- elif usesRotary -%}
+((batch * Q_HEADS + head) * params.seqLen + tile0 + base_a + {{ rb * 8 }}u) * HEAD_DIM + k_base + step
+{%- else -%}
+(batch * params.seqLen + tile0 + base_a + {{ rb * 8 }}u) * Q_STRIDE + head * HEAD_DIM + k_base + step
+{%- endif %}
+{%- endmacro %}
 {% macro score_tile() %}
+        // S = Q.K^T, retaining the same sequence of 8-wide matrix operations.
+        // Fully populated query tiles need neither staging nor K-loop barriers.
+        // A sequence with a partial query tile retains zero-padded staging.
         for (var k_base = 0u; k_base < HEAD_DIM; k_base += TILE_K) {
+{% if not sgmatDirectQuery %}
           {
             let a_row = li / 4u;
+            let a_col = (li % 4u) * {{ (sparseSgmatTileK / 4) | int }}u;
+            for (var i = 0u; i < {{ (sparseSgmatTileK / 4) | int }}u; i++) {
               let k = k_base + a_col + i;
 {% if usesRotary %}
               let q_index = ((batch * Q_HEADS + head) * params.seqLen
             }
           }
           workgroupBarrier();
+{% endif %}
           for (var step = 0u; step < TILE_K; step += 8u) {
 {% for rb in range(2) %}
             let mat_a{{ rb }} = subgroupMatrixLoad<subgroup_matrix_left<f32, 8, 8>, row_major>(
+              &{{ queryMatrixSource }}, {{ query_matrix_offset(rb) }}, {{ queryMatrixStride }}
             );
 {% endfor %}
+{% for cb in range(scoreColBlocks) %}
             let mat_b{{ cb }} = subgroupMatrixLoad<subgroup_matrix_right<f32, 8, 8>, col_major>(
               &present_key, key_row_elem + (base_b + {{ cb * 8 }}u) * HEAD_DIM + k_base + step, HEAD_DIM
             );
 {% endfor %}
 {% for rb in range(2) %}
+{% for cb in range(scoreColBlocks) %}
             mat_s{{ rb }}{{ cb }} = subgroupMatrixMultiplyAccumulate(
               mat_a{{ rb }}, mat_b{{ cb }}, mat_s{{ rb }}{{ cb }}
             );
 {% endfor %}
 {% endfor %}
           }
+{% if not sgmatDirectQuery %}
           workgroupBarrier();
+{% endif %}
         }
 {% endmacro %}
   for (var mask_row = row_first; mask_row <= row_last; mask_row++) {
     let start = u32(block_row_indices[row_base + mask_row]);
     let end = u32(block_row_indices[row_base + mask_row + 1u]);
+    // A row listing every causal block is dense; its column-index values are not
+    // read.
     let dense_row = layout_is_dense || (end - start == mask_row + 1u);
     // Furthest causal bound (exclusive) of the tile's queries on this row.
     let bound_max = min(q_abs0 + rows_live, (mask_row + 1u) * SPARSE_BLOCK);
       if (tile_ok && key_base < bound_max) {
         let key_row_elem = ((batch * KV_HEADS + kv_head) * MAX_CACHE_SEQ + key_base) * HEAD_DIM;
 {% for rb in range(2) %}
+{% for cb in range(scoreColBlocks) %}
         var mat_s{{ rb }}{{ cb }}: subgroup_matrix_result<f32, 8, 8>;
 {% endfor %}
 {% endfor %}
         var tile_stat_m{{ rb }} = -FLT_MAX;
         var tile_stat_d{{ rb }} = 0.0;
 {% endif %}
+{% for cb in range(scoreColBlocks) %}
         subgroupMatrixStore<row_major>(
+          &tile_q, (subgroup * {{ scoreColBlocks }}u + {{ cb }}u) * 64u, mat_s{{ rb }}{{ cb }}, 8u
         );
 {% endfor %}
         workgroupBarrier();
+{% for cb in range(scoreColBlocks) %}
         for (var pair = 0u; pair < 2u; pair++) {
           let r = base_a + {{ rb * 8 }}u + row_in_block;
           let key = key_base + base_b + {{ cb * 8 }}u + col_in_block + pair;
 {% if phase == "stats" %}
           if (allowed) {
             let scored = tile_q[
+              (subgroup * {{ scoreColBlocks }}u + {{ cb }}u) * 64u + row_in_block * 8u + col_in_block + pair
             ] * scale;
             let prev_m = tile_stat_m{{ rb }};
             tile_stat_m{{ rb }} = max(tile_stat_m{{ rb }}, scored);
           var prob = 0.0;
           if (allowed) {
             prob = exp_shift(tile_q[
+              (subgroup * {{ scoreColBlocks }}u + {{ cb }}u) * 64u + row_in_block * 8u + col_in_block + pair
             ] * scale, row_m[r]);
           }
           prob_tile[r * TILE_N + base_b + {{ cb * 8 }}u + col_in_block + pair] = prob;
   let subtile_idy = subgroup % 4u;
   let subtile_idx = subgroup / 4u;
   let base_a = subtile_idy * 16u;
+  let base_b = subtile_idx * {{ (sparseSgmatTileN / 2) | int }}u;
   let row_in_block = lane / 4u;
   let col_in_block = (lane % 4u) * 2u;
 {{ sweep("stats") }}
 {{ sweep("apply") }}
+  // Normalize by the final denominators and store. Publish only as many
+  // fragment columns per batch as fit prob_tile, keeping the smaller key tile's
+  // storage bound even when the output head is wider.
+{% set storeColBlocks = pvColBlocks if pvColBlocks <= sparseSgmatTileN / 8 else (sparseSgmatTileN / 8) | int %}
 {% for rb in range(2) %}
+{% for colBase in range(0, pvColBlocks, storeColBlocks) %}
+{% if rb > 0 or colBase > 0 %}
   workgroupBarrier();
 {% endif %}
+{% for cb in range(colBase, pvColBlocks if pvColBlocks < colBase + storeColBlocks else colBase + storeColBlocks) %}
   subgroupMatrixStore<row_major>(
+    &prob_tile, (subgroup * {{ storeColBlocks }}u + {{ cb - colBase }}u) * 64u, mat_o{{ rb }}{{ cb }}, 8u
   );
 {% endfor %}
   workgroupBarrier();
+{% for cb in range(colBase, pvColBlocks if pvColBlocks < colBase + storeColBlocks else colBase + storeColBlocks) %}
   for (var pair = 0u; pair < 2u; pair++) {
     let r = base_a + {{ rb * 8 }}u + row_in_block;
     if (r < rows_live && row_d[r] > 0.0) {
       let dcol = subtile_idx * PV_COL_SPAN + {{ cb * 8 }}u + col_in_block + pair;
       let out_index = (batch * params.seqLen + tile0 + r) * Q_HIDDEN + head * HEAD_DIM + dcol;
       output[out_index] = {{ scalar }}(prob_tile[
+        (subgroup * {{ storeColBlocks }}u + {{ cb - colBase }}u) * 64u + row_in_block * 8u + col_in_block + pair
       ] / row_d[r]);
     }
   }
 {% endfor %}
+{% endfor %}
 {% endfor %}
   // An all-masked row becomes a uniform average over causal history rather

build/webgpu/sparse-attention.wgsl.jinja CHANGED Viewed

@@ -11,9 +11,6 @@ fn past_sequence_length(batch: u32) -> u32 {
 }
 {%- endmacro %}
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.SparseAttention, attention pass.
@@ -35,11 +32,8 @@ enable f16;
 const Q_HEADS: u32 = {{ numHeads }}u;
 const KV_HEADS: u32 = {{ kvNumHeads }}u;
 const HEAD_DIM: u32 = {{ headSize }}u;
-// The cache is read four elements at a time. The sweep was one narrow dependent load
-// per head dimension per key, and the score phase alone issues HEAD_DIM of them for
-// every key a thread owns; a vec4 view quarters that count. The operator's contract
-// requires head_size to be a multiple of eight, so there is no shape where the
-// quartering does not divide and no scalar form to fall back to.
 const HEAD_VEC: u32 = {{ headVec }}u;
 const MAX_CACHE_SEQ: u32 = {{ maxCacheSeq }}u;
 const SPARSE_BLOCK: u32 = {{ sparseBlockSize }}u;
@@ -80,12 +74,8 @@ var<workgroup> q_shared: array<vec4<f32>, Q_TILE * HEAD_VEC>;
 var<workgroup> running_out: array<vec4<f32>, Q_TILE * HEAD_VEC>;
 var<workgroup> probs: array<f32, Q_TILE * WG>;
 {% if vStageWorthIt %}
-// Staged value rows for the running-output walk (small-dispatch shapes only):
-// the walk itself occupies just HEAD_VEC lanes, so the whole workgroup
-// prefetches each chunk of value rows at full width and the walkers read
-// shared memory instead of chaining dependent global loads. Large dispatches
-// already hide that latency across resident workgroups, and the chunk
-// barriers only cost them, so they keep the direct walk.
 const V_STAGE_KEYS: u32 = 16u;
 var<workgroup> v_stage: array<vec4<f32>, V_STAGE_KEYS * HEAD_VEC>;
 {% endif %}
@@ -102,20 +92,22 @@ var<workgroup> key_rows: array<u32, WG>;
 // calls before their shared partial storage is reused.
 {% set combineSubgroups = combineSubgroups is defined and combineSubgroups %}
 {% if combineSubgroups %}
-// Per-subgroup partials are published into a deterministic slot: the subgroup's
-// ordinal index within the workgroup (lidx / sgSize). The online (m, d) merge
-// is not float-associative, so thread 0 must fold partials in a fixed order.
-// Subgroups partition a workgroup into contiguous ordinal ranges on supported
-// backends, so the ordinal slot is unique per subgroup and every slot in
-// [0, subgroupCount) is written (each subgroup elects one leader).
-// Sized for the worst case of one partial per invocation.
 var<workgroup> partialM: array<f32, WG>;
 var<workgroup> partialD: array<f32, WG>;
 var<workgroup> combinedMD: vec2<f32>;
 // When the whole workgroup is one subgroup the subgroup reduce already covers
-// it (no barriers, no shared state); otherwise subgroup leaders publish
-// partials through shared memory and thread 0 folds them in ordinal order.
 fn combine_partials(m: f32, d: f32, lidx: u32, sgSize: u32) -> vec2<f32> {
   let sgM = subgroupMax(m);
   // A lane with no elements contributes d == 0 (exact identity). A +inf
@@ -126,31 +118,24 @@ fn combine_partials(m: f32, d: f32, lidx: u32, sgSize: u32) -> vec2<f32> {
   if (sgSize == WG) {
     return vec2<f32>(sgM, sgD);
   }
-  let subgroupCount = (WG + sgSize - 1u) / sgSize;
-  // Pre-seed every fold slot with the (max, denom) identity. The fold below reads a
-  // fixed subgroupCount slots in ordinal order (for determinism), but a slot whose
-  // subgroup elects no leader this call — e.g. a fully out-of-window key tile in the
-  // flash-attention loop that re-uses this shared memory each iteration — would
-  // otherwise read stale shared memory. Identity makes such a slot a no-op.
-  // (max identity = -FLT_MAX, denom identity = 0.)
-  if (lidx < subgroupCount) {
-    partialM[lidx] = -FLT_MAX;
-    partialD[lidx] = 0.0;
-  }
-  workgroupBarrier();
   if (subgroupElect()) {
-    let slot = lidx / sgSize;
-    partialM[slot] = sgM;
-    partialD[slot] = sgD;
   }
   workgroupBarrier();
   if (lidx == 0u) {
     var accM = -FLT_MAX;
     var accD = 0.0;
-    for (var i = 0u; i < subgroupCount; i = i + 1u) {
-      let mNew = max(accM, partialM[i]);
-      accD = accD * exp_shift(accM, mNew) + partialD[i] * exp_shift(partialM[i], mNew);
-      accM = mNew;
     }
     combinedMD = vec2<f32>(accM, accD);
   }
@@ -165,9 +150,8 @@ var<workgroup> partialM: array<f32, {{ mdExtent }}>;
 var<workgroup> partialD: array<f32, {{ mdExtent }}>;
 {% if mdStreamed %}
-// In-place fold of {{ mdStreams }} streams. The caller stores its per-thread
-// partials into partialM/partialD first and reads the merged pair of stream s
-// from slot s * WG afterwards.
 fn combine_partials_streams(lidx: u32) {
   workgroupBarrier();
   var stride = WG / 2u;

 }
 {%- endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.SparseAttention, attention pass.
 const Q_HEADS: u32 = {{ numHeads }}u;
 const KV_HEADS: u32 = {{ kvNumHeads }}u;
 const HEAD_DIM: u32 = {{ headSize }}u;
+// The cache is read four elements at a time. `head_size` is a multiple of eight,
+// so every row divides into whole vec4 values and requires no scalar tail.
 const HEAD_VEC: u32 = {{ headVec }}u;
 const MAX_CACHE_SEQ: u32 = {{ maxCacheSeq }}u;
 const SPARSE_BLOCK: u32 = {{ sparseBlockSize }}u;
 var<workgroup> running_out: array<vec4<f32>, Q_TILE * HEAD_VEC>;
 var<workgroup> probs: array<f32, Q_TILE * WG>;
 {% if vStageWorthIt %}
+// The workgroup stages each chunk of value rows, then HEAD_VEC lanes consume
+// those rows from shared memory during the running-output update.
 const V_STAGE_KEYS: u32 = 16u;
 var<workgroup> v_stage: array<vec4<f32>, V_STAGE_KEYS * HEAD_VEC>;
 {% endif %}
 // calls before their shared partial storage is reused.
 {% set combineSubgroups = combineSubgroups is defined and combineSubgroups %}
 {% if combineSubgroups %}
+// Cross-subgroup merge that assumes nothing about which invocations share a
+// subgroup or how many subgroups there are: each subgroup's elected lane
+// publishes the subgroup pair in the slot at its OWN invocation index and sets
+// that index's bit in a workgroup bitmask; thread 0 then folds exactly the
+// published slots, in ascending index order (the online (m, d) merge is not
+// float-associative, so the order is fixed), and clears the mask for the next
+// call as it reads it. Workgroup memory starts zeroed, so the mask needs no
+// setup. Same three collectives as a single-subgroup reduce, two barriers.
 var<workgroup> partialM: array<f32, WG>;
 var<workgroup> partialD: array<f32, WG>;
+var<workgroup> leaderMask: array<atomic<u32>, (WG + 31u) / 32u>;
 var<workgroup> combinedMD: vec2<f32>;
 // When the whole workgroup is one subgroup the subgroup reduce already covers
+// it (no barriers, no shared state). `subgroup_size` is the size of the current
+// subgroup and uniform, so the test is exact and may guard the barriers below.
 fn combine_partials(m: f32, d: f32, lidx: u32, sgSize: u32) -> vec2<f32> {
   let sgM = subgroupMax(m);
   // A lane with no elements contributes d == 0 (exact identity). A +inf
   if (sgSize == WG) {
     return vec2<f32>(sgM, sgD);
   }
   if (subgroupElect()) {
+    partialM[lidx] = sgM;
+    partialD[lidx] = sgD;
+    atomicOr(&leaderMask[lidx / 32u], 1u << (lidx % 32u));
   }
   workgroupBarrier();
   if (lidx == 0u) {
     var accM = -FLT_MAX;
     var accD = 0.0;
+    for (var w = 0u; w < (WG + 31u) / 32u; w = w + 1u) {
+      var bits = atomicExchange(&leaderMask[w], 0u);
+      while (bits != 0u) {
+        let slot = w * 32u + firstTrailingBit(bits);
+        bits = bits & (bits - 1u);
+        let mNew = max(accM, partialM[slot]);
+        accD = accD * exp_shift(accM, mNew) + partialD[slot] * exp_shift(partialM[slot], mNew);
+        accM = mNew;
+      }
     }
     combinedMD = vec2<f32>(accM, accD);
   }
 var<workgroup> partialD: array<f32, {{ mdExtent }}>;
 {% if mdStreamed %}
+// In-place fold of {{ mdStreams }} streams. Input partials occupy
+// partialM/partialD; stream s returns its merged pair in slot s * WG.
 fn combine_partials_streams(lidx: u32) {
   workgroupBarrier();
   var stride = WG / 2u;

build/webgpu/sparse-kv-append.wgsl.jinja CHANGED Viewed

@@ -46,9 +46,6 @@ fn rotary_value(own: f32, partner: f32, cs: f32, sn: f32, first: bool) -> f32 {
 }
 {%- endmacro %}
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.SparseAttention, KV append pass.
@@ -80,11 +77,10 @@ const ROTARY_DIM: u32 = {{ rotaryDim }}u;
 {% endif %}
 @compute @workgroup_size(WG, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>,
-        @builtin(num_workgroups) nwg: vec3<u32>) {
   // 2D-folded flat index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension limit. Reduces to gid.x when nwg.y == 1.
-  let index = gid.x + gid.y * nwg.x * WG;
   let count = params.batchSize * KV_HEADS * params.seqLen * HEAD_DIM;
   if (index >= count) {
     return;

 }
 {%- endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.SparseAttention, KV append pass.
 {% endif %}
 @compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
   // 2D-folded flat index: gid.y carries the high bits past the
+  // per-axis dispatch fold width. Reduces to gid.x when the dispatch does not fold.
+  let index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
   let count = params.batchSize * KV_HEADS * params.seqLen * HEAD_DIM;
   if (index >= count) {
     return;

build/webgpu/sparse-q-rotary.wgsl.jinja CHANGED Viewed

@@ -46,9 +46,6 @@ fn rotary_value(own: f32, partner: f32, cs: f32, sn: f32, first: bool) -> f32 {
 }
 {%- endmacro %}
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.SparseAttention, query rotary pass.
@@ -68,11 +65,10 @@ const WG: u32 = {{ appendWorkgroupSize }}u;
 {{ sparse_rotary(rotaryInterleaved) }}
 @compute @workgroup_size(WG, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>,
-        @builtin(num_workgroups) nwg: vec3<u32>) {
   // 2D-folded flat index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension limit. Reduces to gid.x when nwg.y == 1.
-  let index = gid.x + gid.y * nwg.x * WG;
   let count = params.batchSize * Q_HEADS * params.seqLen * HEAD_DIM;
   if (index >= count) {
     return;

 }
 {%- endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.SparseAttention, query rotary pass.
 {{ sparse_rotary(rotaryInterleaved) }}
 @compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
   // 2D-folded flat index: gid.y carries the high bits past the
+  // per-axis dispatch fold width. Reduces to gid.x when the dispatch does not fold.
+  let index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
   let count = params.batchSize * Q_HEADS * params.seqLen * HEAD_DIM;
   if (index >= count) {
     return;

build/webgpu/test.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "com.microsoft.SparseAttention",
   "fixtureArrays": {
     "prompt_separate_two_layouts_input_blockColIndicesT": [0, 0, 1, 1, 2, 1, 2, 3, -1, 0, 0, 1, 0, 1, 2, 0, 2, 3]
   },
@@ -1056,6 +1055,556 @@
         "pastKeyT": { "dtype": "float32", "shape": [1, 1, 128, 64], "tolerance": 0.00002, "relTolerance": 0.0002 },
         "pastValueT": { "dtype": "float32", "shape": [1, 1, 128, 64], "tolerance": 0.00002, "relTolerance": 0.0002 }
       }
     }
   ]
 }

 {
   "fixtureArrays": {
     "prompt_separate_two_layouts_input_blockColIndicesT": [0, 0, 1, 1, 2, 1, 2, 3, -1, 0, 0, 1, 0, 1, 2, 0, 2, 3]
   },
         "pastKeyT": { "dtype": "float32", "shape": [1, 1, 128, 64], "tolerance": 0.00002, "relTolerance": 0.0002 },
         "pastValueT": { "dtype": "float32", "shape": [1, 1, 128, 64], "tolerance": 0.00002, "relTolerance": 0.0002 }
       }
+    },
+    {
+      "name": "sgmat_capacity_d96_s64_packed0_rotary0",
+      "provenance": {
+        "notes": "Three 32-wide head blocks exercise the final partial publication bank on a constrained matrix tile; per-batch history crosses sparse rows, with complete and padded query tiles."
+      },
+      "attrs": { "num_heads": 2, "kv_num_heads": 2, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 64, 192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [2, 64, 192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [2, 64, 192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.017, "scale": 0.25 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.013, "scale": 0.25 }
+        },
+        "blockRowIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [0, 1, 3, 5, 7, 0, 1, 2, 4, 5] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 7],
+          "data": { "kind": "values", "values": [0, 0, 1, 0, 2, 1, 3, 0, 1, 0, 2, 3, 0, 0] }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } },
+        "keyTotalSequenceLengthsT": {
+          "dtype": "int32",
+          "shape": [2],
+          "data": { "kind": "values", "values": [164, 132] }
+        }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [2, 64, 192], "tolerance": 0.0001, "relTolerance": 0.0005 },
+        "pastKeyT": { "dtype": "float32", "shape": [2, 2, 256, 96], "tolerance": 0.00002, "relTolerance": 0.0002 },
+        "pastValueT": { "dtype": "float32", "shape": [2, 2, 256, 96], "tolerance": 0.00002, "relTolerance": 0.0002 }
+      }
+    },
+    {
+      "name": "sgmat_capacity_d96_s64_packed0_rotary1",
+      "provenance": {
+        "notes": "Three 32-wide head blocks exercise the final partial publication bank on a constrained matrix tile; per-batch history crosses sparse rows, with complete and padded query tiles."
+      },
+      "attrs": { "num_heads": 2, "kv_num_heads": 2, "sparse_block_size": 64, "do_rotary": 1, "rotary_interleaved": 0 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 64, 192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [2, 64, 192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [2, 64, 192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.017, "scale": 0.25 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.013, "scale": 0.25 }
+        },
+        "blockRowIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [0, 1, 3, 5, 7, 0, 1, 2, 4, 5] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 7],
+          "data": { "kind": "values", "values": [0, 0, 1, 0, 2, 1, 3, 0, 1, 0, 2, 3, 0, 0] }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } },
+        "keyTotalSequenceLengthsT": {
+          "dtype": "int32",
+          "shape": [2],
+          "data": { "kind": "values", "values": [164, 132] }
+        },
+        "cosCacheT": {
+          "dtype": "float32",
+          "shape": [256, 48],
+          "data": { "kind": "fillFloat32", "sinStep": 0.01, "cosStep": 0.013, "scale": 0.5 }
+        },
+        "sinCacheT": {
+          "dtype": "float32",
+          "shape": [256, 48],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.013, "scale": 0.5 }
+        }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [2, 64, 192], "tolerance": 0.0001, "relTolerance": 0.0005 },
+        "pastKeyT": { "dtype": "float32", "shape": [2, 2, 256, 96], "tolerance": 0.00002, "relTolerance": 0.0002 },
+        "pastValueT": { "dtype": "float32", "shape": [2, 2, 256, 96], "tolerance": 0.00002, "relTolerance": 0.0002 }
+      }
+    },
+    {
+      "name": "sgmat_capacity_d96_s64_packed1_rotary0",
+      "provenance": {
+        "notes": "Three 32-wide head blocks exercise the final partial publication bank on a constrained matrix tile; per-batch history crosses sparse rows, with complete and padded query tiles."
+      },
+      "attrs": { "num_heads": 2, "kv_num_heads": 2, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 64, 576],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.017, "scale": 0.25 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.013, "scale": 0.25 }
+        },
+        "blockRowIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [0, 1, 3, 5, 7, 0, 1, 2, 4, 5] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 7],
+          "data": { "kind": "values", "values": [0, 0, 1, 0, 2, 1, 3, 0, 1, 0, 2, 3, 0, 0] }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } },
+        "keyTotalSequenceLengthsT": {
+          "dtype": "int32",
+          "shape": [2],
+          "data": { "kind": "values", "values": [164, 132] }
+        }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [2, 64, 192], "tolerance": 0.0001, "relTolerance": 0.0005 },
+        "pastKeyT": { "dtype": "float32", "shape": [2, 2, 256, 96], "tolerance": 0.00002, "relTolerance": 0.0002 },
+        "pastValueT": { "dtype": "float32", "shape": [2, 2, 256, 96], "tolerance": 0.00002, "relTolerance": 0.0002 }
+      }
+    },
+    {
+      "name": "sgmat_capacity_d96_s64_packed1_rotary1",
+      "provenance": {
+        "notes": "Three 32-wide head blocks exercise the final partial publication bank on a constrained matrix tile; per-batch history crosses sparse rows, with complete and padded query tiles."
+      },
+      "attrs": { "num_heads": 2, "kv_num_heads": 2, "sparse_block_size": 64, "do_rotary": 1, "rotary_interleaved": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 64, 576],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.017, "scale": 0.25 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.013, "scale": 0.25 }
+        },
+        "blockRowIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [0, 1, 3, 5, 7, 0, 1, 2, 4, 5] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 7],
+          "data": { "kind": "values", "values": [0, 0, 1, 0, 2, 1, 3, 0, 1, 0, 2, 3, 0, 0] }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } },
+        "keyTotalSequenceLengthsT": {
+          "dtype": "int32",
+          "shape": [2],
+          "data": { "kind": "values", "values": [164, 132] }
+        },
+        "cosCacheT": {
+          "dtype": "float32",
+          "shape": [256, 48],
+          "data": { "kind": "fillFloat32", "sinStep": 0.01, "cosStep": 0.013, "scale": 0.5 }
+        },
+        "sinCacheT": {
+          "dtype": "float32",
+          "shape": [256, 48],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.013, "scale": 0.5 }
+        }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [2, 64, 192], "tolerance": 0.0001, "relTolerance": 0.0005 },
+        "pastKeyT": { "dtype": "float32", "shape": [2, 2, 256, 96], "tolerance": 0.00002, "relTolerance": 0.0002 },
+        "pastValueT": { "dtype": "float32", "shape": [2, 2, 256, 96], "tolerance": 0.00002, "relTolerance": 0.0002 }
+      }
+    },
+    {
+      "name": "sgmat_capacity_d96_s65_packed0_rotary0",
+      "provenance": {
+        "notes": "Three 32-wide head blocks exercise the final partial publication bank on a constrained matrix tile; per-batch history crosses sparse rows, with complete and padded query tiles."
+      },
+      "attrs": { "num_heads": 2, "kv_num_heads": 2, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 65, 192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [2, 65, 192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [2, 65, 192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.017, "scale": 0.25 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.013, "scale": 0.25 }
+        },
+        "blockRowIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [0, 1, 3, 5, 7, 0, 1, 2, 4, 5] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 7],
+          "data": { "kind": "values", "values": [0, 0, 1, 0, 2, 1, 3, 0, 1, 0, 2, 3, 0, 0] }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } },
+        "keyTotalSequenceLengthsT": {
+          "dtype": "int32",
+          "shape": [2],
+          "data": { "kind": "values", "values": [165, 133] }
+        }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [2, 65, 192], "tolerance": 0.0001, "relTolerance": 0.0005 },
+        "pastKeyT": { "dtype": "float32", "shape": [2, 2, 256, 96], "tolerance": 0.00002, "relTolerance": 0.0002 },
+        "pastValueT": { "dtype": "float32", "shape": [2, 2, 256, 96], "tolerance": 0.00002, "relTolerance": 0.0002 }
+      }
+    },
+    {
+      "name": "sgmat_capacity_d96_s65_packed0_rotary1",
+      "provenance": {
+        "notes": "Three 32-wide head blocks exercise the final partial publication bank on a constrained matrix tile; per-batch history crosses sparse rows, with complete and padded query tiles."
+      },
+      "attrs": { "num_heads": 2, "kv_num_heads": 2, "sparse_block_size": 64, "do_rotary": 1, "rotary_interleaved": 0 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 65, 192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [2, 65, 192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [2, 65, 192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.017, "scale": 0.25 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.013, "scale": 0.25 }
+        },
+        "blockRowIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [0, 1, 3, 5, 7, 0, 1, 2, 4, 5] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 7],
+          "data": { "kind": "values", "values": [0, 0, 1, 0, 2, 1, 3, 0, 1, 0, 2, 3, 0, 0] }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } },
+        "keyTotalSequenceLengthsT": {
+          "dtype": "int32",
+          "shape": [2],
+          "data": { "kind": "values", "values": [165, 133] }
+        },
+        "cosCacheT": {
+          "dtype": "float32",
+          "shape": [256, 48],
+          "data": { "kind": "fillFloat32", "sinStep": 0.01, "cosStep": 0.013, "scale": 0.5 }
+        },
+        "sinCacheT": {
+          "dtype": "float32",
+          "shape": [256, 48],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.013, "scale": 0.5 }
+        }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [2, 65, 192], "tolerance": 0.0001, "relTolerance": 0.0005 },
+        "pastKeyT": { "dtype": "float32", "shape": [2, 2, 256, 96], "tolerance": 0.00002, "relTolerance": 0.0002 },
+        "pastValueT": { "dtype": "float32", "shape": [2, 2, 256, 96], "tolerance": 0.00002, "relTolerance": 0.0002 }
+      }
+    },
+    {
+      "name": "sgmat_capacity_d96_s65_packed1_rotary0",
+      "provenance": {
+        "notes": "Three 32-wide head blocks exercise the final partial publication bank on a constrained matrix tile; per-batch history crosses sparse rows, with complete and padded query tiles."
+      },
+      "attrs": { "num_heads": 2, "kv_num_heads": 2, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 65, 576],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.017, "scale": 0.25 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.013, "scale": 0.25 }
+        },
+        "blockRowIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [0, 1, 3, 5, 7, 0, 1, 2, 4, 5] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 7],
+          "data": { "kind": "values", "values": [0, 0, 1, 0, 2, 1, 3, 0, 1, 0, 2, 3, 0, 0] }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } },
+        "keyTotalSequenceLengthsT": {
+          "dtype": "int32",
+          "shape": [2],
+          "data": { "kind": "values", "values": [165, 133] }
+        }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [2, 65, 192], "tolerance": 0.0001, "relTolerance": 0.0005 },
+        "pastKeyT": { "dtype": "float32", "shape": [2, 2, 256, 96], "tolerance": 0.00002, "relTolerance": 0.0002 },
+        "pastValueT": { "dtype": "float32", "shape": [2, 2, 256, 96], "tolerance": 0.00002, "relTolerance": 0.0002 }
+      }
+    },
+    {
+      "name": "sgmat_capacity_d96_s65_packed1_rotary1",
+      "provenance": {
+        "notes": "Three 32-wide head blocks exercise the final partial publication bank on a constrained matrix tile; per-batch history crosses sparse rows, with complete and padded query tiles."
+      },
+      "attrs": { "num_heads": 2, "kv_num_heads": 2, "sparse_block_size": 64, "do_rotary": 1, "rotary_interleaved": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 65, 576],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.017, "scale": 0.25 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.013, "scale": 0.25 }
+        },
+        "blockRowIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [0, 1, 3, 5, 7, 0, 1, 2, 4, 5] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 7],
+          "data": { "kind": "values", "values": [0, 0, 1, 0, 2, 1, 3, 0, 1, 0, 2, 3, 0, 0] }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } },
+        "keyTotalSequenceLengthsT": {
+          "dtype": "int32",
+          "shape": [2],
+          "data": { "kind": "values", "values": [165, 133] }
+        },
+        "cosCacheT": {
+          "dtype": "float32",
+          "shape": [256, 48],
+          "data": { "kind": "fillFloat32", "sinStep": 0.01, "cosStep": 0.013, "scale": 0.5 }
+        },
+        "sinCacheT": {
+          "dtype": "float32",
+          "shape": [256, 48],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.013, "scale": 0.5 }
+        }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [2, 65, 192], "tolerance": 0.0001, "relTolerance": 0.0005 },
+        "pastKeyT": { "dtype": "float32", "shape": [2, 2, 256, 96], "tolerance": 0.00002, "relTolerance": 0.0002 },
+        "pastValueT": { "dtype": "float32", "shape": [2, 2, 256, 96], "tolerance": 0.00002, "relTolerance": 0.0002 }
+      }
+    },
+    {
+      "name": "sgmat_parallelism_d96_s128",
+      "provenance": {
+        "notes": "The matrix scheduling boundary: 16 or 32 workgroups with identical batch/head geometry. A 96-wide head exercises partial result-bank publication on the smaller capacity tier."
+      },
+      "attrs": { "num_heads": 4, "kv_num_heads": 2, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 128, 384],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [2, 128, 192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [2, 128, 192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.017, "scale": 0.25 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.013, "scale": 0.25 }
+        },
+        "blockRowIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [0, 1, 3, 5, 7, 0, 1, 2, 4, 5] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 7],
+          "data": { "kind": "values", "values": [0, 0, 1, 0, 2, 1, 3, 0, 1, 0, 2, 3, 0, 0] }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } },
+        "keyTotalSequenceLengthsT": {
+          "dtype": "int32",
+          "shape": [2],
+          "data": { "kind": "values", "values": [256, 256] }
+        }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [2, 128, 384], "tolerance": 0.0001, "relTolerance": 0.0005 },
+        "pastKeyT": { "dtype": "float32", "shape": [2, 2, 256, 96], "tolerance": 0.00002, "relTolerance": 0.0002 },
+        "pastValueT": { "dtype": "float32", "shape": [2, 2, 256, 96], "tolerance": 0.00002, "relTolerance": 0.0002 }
+      }
+    },
+    {
+      "name": "sgmat_parallelism_d96_s256",
+      "provenance": {
+        "notes": "The matrix scheduling boundary: 16 or 32 workgroups with identical batch/head geometry. A 96-wide head exercises partial result-bank publication on the smaller capacity tier."
+      },
+      "attrs": { "num_heads": 4, "kv_num_heads": 2, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 256, 384],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [2, 256, 192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [2, 256, 192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.017, "scale": 0.25 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.013, "scale": 0.25 }
+        },
+        "blockRowIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [0, 1, 3, 5, 7, 0, 1, 2, 4, 5] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 7],
+          "data": { "kind": "values", "values": [0, 0, 1, 0, 2, 1, 3, 0, 1, 0, 2, 3, 0, 0] }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } },
+        "keyTotalSequenceLengthsT": {
+          "dtype": "int32",
+          "shape": [2],
+          "data": { "kind": "values", "values": [256, 256] }
+        }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [2, 256, 384], "tolerance": 0.0001, "relTolerance": 0.0005 },
+        "pastKeyT": { "dtype": "float32", "shape": [2, 2, 256, 96], "tolerance": 0.00002, "relTolerance": 0.0002 },
+        "pastValueT": { "dtype": "float32", "shape": [2, 2, 256, 96], "tolerance": 0.00002, "relTolerance": 0.0002 }
+      }
     }
   ]
 }