sync 2e7068faf55e

Browse files

Files changed (8) hide show

README.md +96 -0
build/webgpu/bench.json +231 -0
build/webgpu/manifest.json +752 -0
build/webgpu/metadata.json +20 -0
build/webgpu/sparse-attention.wgsl.jinja +425 -0
build/webgpu/sparse-kv-append.wgsl.jinja +125 -0
build/webgpu/sparse-q-rotary.wgsl.jinja +102 -0
build/webgpu/test.json +745 -0

README.md CHANGED Viewed

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

 ---
+library_name: kernels
 license: apache-2.0
+tags:
+- kernel
+- webgpu
+- wgsl
 ---
+# com.microsoft.SparseAttention
+`com.microsoft`  ·  ONNX Runtime contrib operator  ·  contrib since_version 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.
+See the [ONNX Runtime `SparseAttention` contrib-operator spec](https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.SparseAttention) for the reference semantics.
+## 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
+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
+| Variable | Allowed dtypes |
+| --- | --- |
+| `T` | `float32`, `float16` |
+| `M` | `int32` |
+## 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
+- [`sparse-attention.wgsl.jinja`](build/webgpu/sparse-attention.wgsl.jinja)
+- [`sparse-kv-append.wgsl.jinja`](build/webgpu/sparse-kv-append.wgsl.jinja)
+- [`sparse-q-rotary.wgsl.jinja`](build/webgpu/sparse-q-rotary.wgsl.jinja)
+## Use with `@huggingface/kernels`
+The loader derives every required output's shape and logical dtype from the manifest contract and this call.
+It then allocates the result tensors automatically.
+The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+Replace each `*Data` placeholder with a typed array containing the corresponding input data.
+```js
+import { getKernel } from "@huggingface/kernels";
+const kernel = await getKernel("webgpu-kernels/com.microsoft.SparseAttention", { version: 1 });
+const { pastKeyT, pastValueT, outputT } = await kernel({
+  queryT: { data: queryTData, shape: [1, 32, 8] },
+  keyT: { data: keyTData, shape: [1, 32, 8] },
+  valueT: { data: valueTData, shape: [1, 32, 8] },
+  pastKeyT: { data: pastKeyTData, shape: [1, 1, 32, 8] },
+  pastValueT: { data: pastValueTData, shape: [1, 1, 32, 8] },
+  blockRowIndicesT: { data: blockRowIndicesTData, shape: [1, 3] },
+  blockColIndicesT: { data: blockColIndicesTData, shape: [1, 3] },
+  totalSequenceLengthT: { data: totalSequenceLengthTData, shape: [1] },
+  keyTotalSequenceLengthsT: { data: keyTotalSequenceLengthsTData, shape: [1] },
+}, {
+  attrs: { num_heads: 1, kv_num_heads: 1, sparse_block_size: 16 },
+});
+```

build/webgpu/bench.json ADDED Viewed

	@@ -0,0 +1,231 @@

+{
+  "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]
+  },
+  "tunableSpace": { "WORKGROUP_SIZE": [32, 64, 128, 256], "APPEND_WORKGROUP_SIZE": [64, 128, 256] },
+  "cases": [
+    {
+      "name": "sparse-prompt-b1-s1024-h32kv8-d128-blk64",
+      "preset": "smoke",
+      "vars": {
+        "dtype": "float32",
+        "batch": 1,
+        "seq": 1024,
+        "heads": 32,
+        "kvHeads": 8,
+        "headDim": 128,
+        "qkPairs": 13254656,
+        "attendedKeys": 1024,
+        "dtypeBytes": 4
+      },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": { "shape": [1, 1024, 4096], "dtype": "float32", "dist": "normal", "seed": 9300, "scale": 1 },
+        "keyT": { "shape": [1, 1024, 1024], "dtype": "float32", "dist": "normal", "seed": 9301, "scale": 1 },
+        "valueT": { "shape": [1, 1024, 1024], "dtype": "float32", "dist": "normal", "seed": 9302, "scale": 1 },
+        "pastKeyT": { "shape": [1, 8, 1024, 128], "dtype": "float32", "dist": "normal", "seed": 9303, "scale": 1 },
+        "pastValueT": { "shape": [1, 8, 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, 1024, 4096], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "4 * args.qkPairs * args.headDim" }] }
+    },
+    {
+      "name": "sparse-decode-b8-past1023-h32kv8-d128-blk64",
+      "preset": "smoke",
+      "vars": {
+        "dtype": "float32",
+        "batch": 8,
+        "seq": 1,
+        "heads": 32,
+        "kvHeads": 8,
+        "headDim": 128,
+        "qkPairs": 163840,
+        "attendedKeys": 640,
+        "dtypeBytes": 4
+      },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": { "shape": [8, 1, 4096], "dtype": "float32", "dist": "normal", "seed": 9310, "scale": 1 },
+        "keyT": { "shape": [8, 1, 1024], "dtype": "float32", "dist": "normal", "seed": 9311, "scale": 1 },
+        "valueT": { "shape": [8, 1, 1024], "dtype": "float32", "dist": "normal", "seed": 9312, "scale": 1 },
+        "pastKeyT": { "shape": [8, 8, 1024, 128], "dtype": "float32", "dist": "normal", "seed": 9313, "scale": 1 },
+        "pastValueT": { "shape": [8, 8, 1024, 128], "dtype": "float32", "dist": "normal", "seed": 9314, "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": [8], "dtype": "int32", "dist": "constant", "value": 1024 }
+      },
+      "outputs": { "outputT": { "shape": [8, 1, 4096], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * args.batch * args.kvHeads * args.attendedKeys * args.headDim * args.dtypeBytes"
+          }
+        ]
+      }
+    },
+    {
+      "name": "sparse-decode-f16-b8-past1023",
+      "preset": "model",
+      "vars": {
+        "dtype": "float16",
+        "batch": 8,
+        "seq": 1,
+        "heads": 32,
+        "kvHeads": 8,
+        "headDim": 128,
+        "qkPairs": 163840,
+        "attendedKeys": 640,
+        "dtypeBytes": 2
+      },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": { "shape": [8, 1, 4096], "dtype": "float16", "dist": "normal", "seed": 9320, "scale": 1 },
+        "keyT": { "shape": [8, 1, 1024], "dtype": "float16", "dist": "normal", "seed": 9321, "scale": 1 },
+        "valueT": { "shape": [8, 1, 1024], "dtype": "float16", "dist": "normal", "seed": 9322, "scale": 1 },
+        "pastKeyT": { "shape": [8, 8, 1024, 128], "dtype": "float16", "dist": "normal", "seed": 9323, "scale": 1 },
+        "pastValueT": { "shape": [8, 8, 1024, 128], "dtype": "float16", "dist": "normal", "seed": 9324, "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": [8], "dtype": "int32", "dist": "constant", "value": 1024 }
+      },
+      "outputs": { "outputT": { "shape": [8, 1, 4096], "dtype": "float16" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * args.batch * args.kvHeads * args.attendedKeys * args.headDim * args.dtypeBytes"
+          }
+        ]
+      }
+    },
+    {
+      "name": "sparse-decode-packed-rotary-b8-past1023",
+      "preset": "model",
+      "vars": {
+        "dtype": "float32",
+        "batch": 8,
+        "seq": 1,
+        "heads": 32,
+        "kvHeads": 8,
+        "headDim": 128,
+        "qkPairs": 163840,
+        "attendedKeys": 640,
+        "dtypeBytes": 4
+      },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sparse_block_size": 64, "do_rotary": 1 },
+      "inputs": {
+        "queryT": { "shape": [8, 1, 6144], "dtype": "float32", "dist": "normal", "seed": 9330, "scale": 1 },
+        "pastKeyT": { "shape": [8, 8, 1024, 128], "dtype": "float32", "dist": "normal", "seed": 9333, "scale": 1 },
+        "pastValueT": { "shape": [8, 8, 1024, 128], "dtype": "float32", "dist": "normal", "seed": 9334, "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": [8], "dtype": "int32", "dist": "constant", "value": 1024 },
+        "cosCacheT": { "shape": [1024, 64], "dtype": "float32", "dist": "normal", "seed": 9335, "scale": 1 },
+        "sinCacheT": { "shape": [1024, 64], "dtype": "float32", "dist": "normal", "seed": 9336, "scale": 1 }
+      },
+      "outputs": { "outputT": { "shape": [8, 1, 4096], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * args.batch * args.kvHeads * args.attendedKeys * args.headDim * args.dtypeBytes"
+          }
+        ]
+      }
+    },
+    {
+      "name": "sparse-chunked-b2-s128-past896",
+      "preset": "model",
+      "vars": {
+        "dtype": "float32",
+        "batch": 2,
+        "seq": 128,
+        "heads": 32,
+        "kvHeads": 8,
+        "headDim": 128,
+        "qkPairs": 4919296,
+        "attendedKeys": 704,
+        "dtypeBytes": 4
+      },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": { "shape": [2, 128, 4096], "dtype": "float32", "dist": "normal", "seed": 9340, "scale": 1 },
+        "keyT": { "shape": [2, 128, 1024], "dtype": "float32", "dist": "normal", "seed": 9341, "scale": 1 },
+        "valueT": { "shape": [2, 128, 1024], "dtype": "float32", "dist": "normal", "seed": 9342, "scale": 1 },
+        "pastKeyT": { "shape": [2, 8, 1024, 128], "dtype": "float32", "dist": "normal", "seed": 9343, "scale": 1 },
+        "pastValueT": { "shape": [2, 8, 1024, 128], "dtype": "float32", "dist": "normal", "seed": 9344, "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": [2], "dtype": "int32", "dist": "constant", "value": 1024 }
+      },
+      "outputs": { "outputT": { "shape": [2, 128, 4096], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "4 * args.qkPairs * args.headDim" }] }
+    }
+  ]
+}

build/webgpu/manifest.json ADDED Viewed

	@@ -0,0 +1,752 @@

+{
+  "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 },
+    "kv_num_heads": { "required": true },
+    "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 },
+  "derive": {
+    "batchSize": "dim(shapes.queryT, 0)",
+    "seqLen": "dim(shapes.queryT, 1)",
+    "numHeads": "attrs.num_heads",
+    "kvNumHeads": "attrs.kv_num_heads",
+    "sparseBlockSize": "attrs.sparse_block_size",
+    "headSize": "dim(shapes.pastKeyT, 3)",
+    "headVec": "headSize / 4",
+    "sparseWidthBound": "max(256, tunables.WORKGROUP_SIZE)",
+    "sparseQueryTileCap": "max(1, floor((device.limits.maxComputeWorkgroupStorageSize / 4 - sparseWidthBound) / (2 * headSize + 3 * sparseWidthBound)))",
+    "sparseQueryTileWant": "min(tunables.QUERY_TILE, min(sparseBlockSize, sparseQueryTileCap))",
+    "sparseQueryTile": "1 if seqLen <= 1 else (16 if sparseQueryTileWant >= 16 and seqLen >= 16 else (8 if sparseQueryTileWant >= 8 and seqLen >= 8 else (4 if sparseQueryTileWant >= 4 and seqLen >= 4 else (2 if sparseQueryTileWant >= 2 and seqLen >= 2 else 1))))",
+    "sparseQueryTiles": "ceilDiv(seqLen, sparseQueryTile)",
+    "sparseAttnWorkgroups": "sparseQueryTiles * batchSize * numHeads",
+    "sparseAttnWorkgroup": "min(256, max(32, pow2ceil(headVec))) if sparseAttnWorkgroups >= tunables.NARROW_MIN_WORKGROUPS else tunables.WORKGROUP_SIZE",
+    "maxCacheSeq": "dim(shapes.pastKeyT, 2)",
+    "numLayout": "dim(shapes.blockRowIndicesT, 0)",
+    "maxBlocks": "dim(shapes.blockRowIndicesT, 1) - 1",
+    "maxNnz": "dim(shapes.blockColIndicesT, 1)",
+    "packedQkv": "not present.keyT",
+    "qHidden": "numHeads * headSize",
+    "kvHidden": "kvNumHeads * headSize",
+    "packedStride": "(numHeads + 2 * kvNumHeads) * headSize",
+    "doRotary": "attrs.do_rotary == 1",
+    "rotaryHalf": "dim(shapes.cosCacheT, 1) if doRotary and present.cosCacheT and ranks.cosCacheT == 2 else 0",
+    "rotaryDim": "2 * rotaryHalf",
+    "useRotary": "doRotary and rotaryDim > 0",
+    "rotaryInterleaved": "attrs.rotary_interleaved == 1",
+    "qRotaryElements": "batchSize * numHeads * seqLen * headSize",
+    "cacheShapeOk": "ranks.pastKeyT == 4 and ranks.pastValueT == 4 and dim(shapes.pastKeyT, 0) == batchSize and dim(shapes.pastKeyT, 1) == kvNumHeads and sameShape(shapes.pastValueT, shapes.pastKeyT)",
+    "queryShapeOk": "dim(shapes.queryT, 2) == (packedStride if packedQkv else qHidden)",
+    "kvShapeOk": "packedQkv or (present.valueT and ranks.keyT == 3 and ranks.valueT == 3 and dim(shapes.keyT, 0) == batchSize and dim(shapes.keyT, 1) == seqLen and dim(shapes.keyT, 2) == kvHidden and sameShape(shapes.valueT, shapes.keyT) and tensorDtypes.keyT == tensorDtypes.queryT and tensorDtypes.valueT == tensorDtypes.queryT)",
+    "kvPairOk": "present.keyT == present.valueT",
+    "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",
+    "separateContract": "contract and not packedQkv and not useRotary",
+    "separateRotaryContract": "contract and not packedQkv and useRotary"
+  },
+  "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" }
+          ]
+        }
+      }
+    ],
+    "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" }
+          ]
+        }
+      }
+    ]
+  },
+  "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" }
+        }
+      ]
+    },
+    {
+      "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" }
+        }
+      ]
+    },
+    {
+      "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" }
+        }
+      ]
+    },
+    {
+      "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" }
+        }
+      ]
+    }
+  ]
+}

build/webgpu/metadata.json ADDED Viewed

	@@ -0,0 +1,20 @@

+{
+  "name": "com.microsoft.SparseAttention",
+  "id": "_com_microsoft_sparseattention_webgpu_07e93c9",
+  "version": 1,
+  "license": "Apache-2.0",
+  "backend": { "type": "webgpu" },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "bench.json": "E0HHOiKzV0TA0/6aDZXLjmxd2z/Q5iYFBFGNYeQoQM0=",
+      "manifest.json": "L9ljQObhgAxTXE7AuOhOeG1E7zQwFIDFC+FnIgrU864=",
+      "sparse-attention.wgsl.jinja": "zYSxDV01i/wcjq9YJjdWDzRHkAbZA6Ctx7Fi/3xdrF4=",
+      "sparse-kv-append.wgsl.jinja": "RavNXJuGESPqe/Jzt5Ek6+844s5q2w/ocyUgGP3cWec=",
+      "sparse-q-rotary.wgsl.jinja": "hI7d9cx6RqRO079h4xPBMTMj+qKqIUOKW3pF3vdeWdg=",
+      "test.json": "aOQD21sQAgUaJwht0fYeGHtHCmUc7VPdhXjgKLH9YEU="
+    }
+  },
+  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
+  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.SparseAttention" }
+}

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

	@@ -0,0 +1,425 @@

+{% 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
+// length to zero; a padded prompt row must not be read as if it had history.
+fn past_sequence_length(batch: u32) -> u32 {
+  if (u32(total_sequence_length[0]) == params.seqLen) {
+    return 0u;
+  }
+  let total = u32(key_total_sequence_lengths[batch]);
+  return select(0u, total - params.seqLen, total >= params.seqLen);
+}
+{%- endmacro %}
+{% if usesF16 %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+// com.microsoft.SparseAttention, attention pass.
+// One workgroup per (batch, query-token tile, query head). The workgroup sweeps only the key
+// blocks this query's layout row selects, in tiles of WG keys with one key per thread,
+// merging each tile into a running online softmax so the score row is never
+// materialized. Repeated CSR columns are suppressed because the CSR data represents a
+// boolean block mask.
+//
+// The block mask is CSR: row q_abs / SPARSE_BLOCK of layout (head % NUM_LAYOUT) spans
+// block_col_indices[start .. end), and each of those columns names SPARSE_BLOCK
+// contiguous keys. Sweeping (end - start) * SPARSE_BLOCK virtual slots and resolving
+// each through the column list is what makes the kernel cost the SELECTED blocks rather
+// than the whole history -- the reason the operator exists.
+//
+// Causality bounds every visited row at this query's absolute position. A layout
+// with the triangular dense count and a sparse row with row + 1 entries take the
+// dense shortcut without inspecting their column values.
+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;
+const NUM_LAYOUT: u32 = {{ numLayout }}u;
+const ROW_STRIDE: u32 = {{ maxBlocks + 1 }}u;
+const COL_STRIDE: u32 = {{ maxNnz }}u;
+const Q_HIDDEN: u32 = {{ numHeads * headSize }}u;
+{% if not usesRotary %}
+const Q_STRIDE: u32 = {{ packedStride if packedQkv else numHeads * headSize }}u;
+{% endif %}
+const WG: u32 = {{ attnWorkgroup }}u;
+const Q_TILE: u32 = {{ qTile }}u;
+// FLT_MAX, not -inf, as the online (m, d) accumulator init: merges must keep
+// `m - m` finite so an empty lane / all--inf row contributes the exact
+// accumulator identity (m, d) = (-FLT_MAX, 0). Operator epilogues interpret
+// a zero final denominator according to their public semantics. Using -inf
+// here changes +inf-row behavior.
+const FLT_MAX: f32 = 3.4028234663852886e38;
+fn is_finite_f32(value: f32) -> bool {
+  return select(false, value <= FLT_MAX, value >= -FLT_MAX);
+}
+// x - m that is exactly 0 when x equals a finite m, so exp(shifted) == 1
+// exactly at the row max. `x - x` on an infinite max is a legal fast-math
+// fold to 0, which would silently turn +inf rows finite — the explicit
+// equality test keeps the NaN propagation of the serial kernels.
+fn shifted_value(value: f32, maxValue: f32) -> f32 {
+  let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
+  return select(value - maxValue, 0.0, equalFiniteMax);
+}
+fn exp_shift(value: f32, maxValue: f32) -> f32 {
+  return exp(shifted_value(value, maxValue));
+}
+var<workgroup> 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>;
+// One resolved cache row base per key of the current tile, so the value accumulation
+// re-reads a base instead of re-walking the column list per head dimension.
+var<workgroup> key_rows: array<u32, WG>;
+{% set mdStreams = qTile %}
+// Workgroup-cooperative merge of per-thread online-softmax (m, d) partials:
+// mNew = max(m1, m2)
+// dNew = d1 * exp(m1 - mNew) + d2 * exp(m2 - mNew)
+// Both the subgroup and portable barrier-tree engines return the same merged
+// pair to every invocation. Repeated merges require a workgroup barrier between
+// calls before their shared partial storage is reused.
+{% set combineSubgroups = combineSubgroups is defined and combineSubgroups %}
+{% if combineSubgroups %}
+// Per-subgroup partials are published into a deterministic slot: the subgroup's
+// ordinal index within the workgroup (lidx / sgSize). The online (m, d) merge
+// is not float-associative, so thread 0 must fold partials in a fixed order.
+// Subgroups partition a workgroup into contiguous ordinal ranges on supported
+// backends, so the ordinal slot is unique per subgroup and every slot in
+// [0, subgroupCount) is written (each subgroup elects one leader).
+// Sized for the worst case of one partial per invocation.
+var<workgroup> partialM: array<f32, WG>;
+var<workgroup> partialD: array<f32, WG>;
+var<workgroup> combinedMD: vec2<f32>;
+// When the whole workgroup is one subgroup the subgroup reduce already covers
+// it (no barriers, no shared state); otherwise subgroup leaders publish
+// partials through shared memory and thread 0 folds them in ordinal order.
+fn combine_partials(m: f32, d: f32, lidx: u32, sgSize: u32) -> vec2<f32> {
+  let sgM = subgroupMax(m);
+  // A lane with no elements contributes d == 0 (exact identity). A +inf
+  // element made exp(inf - inf) = NaN stick in that lane's d; a NaN element
+  // landed in d via exp(NaN); both survive the merge and are detected by the
+  // code after the reduction.
+  let sgD = subgroupAdd(d * exp_shift(m, sgM));
+  if (sgSize == WG) {
+    return vec2<f32>(sgM, sgD);
+  }
+  let subgroupCount = (WG + sgSize - 1u) / sgSize;
+  // Pre-seed every fold slot with the (max, denom) identity. The fold below reads a
+  // fixed subgroupCount slots in ordinal order (for determinism), but a slot whose
+  // subgroup elects no leader this call — e.g. a fully out-of-window key tile in the
+  // flash-attention loop that re-uses this shared memory each iteration — would
+  // otherwise read stale shared memory. Identity makes such a slot a no-op.
+  // (max identity = -FLT_MAX, denom identity = 0.)
+  if (lidx < subgroupCount) {
+    partialM[lidx] = -FLT_MAX;
+    partialD[lidx] = 0.0;
+  }
+  workgroupBarrier();
+  if (subgroupElect()) {
+    let slot = lidx / sgSize;
+    partialM[slot] = sgM;
+    partialD[slot] = sgD;
+  }
+  workgroupBarrier();
+  if (lidx == 0u) {
+    var accM = -FLT_MAX;
+    var accD = 0.0;
+    for (var i = 0u; i < subgroupCount; i = i + 1u) {
+      let mNew = max(accM, partialM[i]);
+      accD = accD * exp_shift(accM, mNew) + partialD[i] * exp_shift(partialM[i], mNew);
+      accM = mNew;
+    }
+    combinedMD = vec2<f32>(accM, accD);
+  }
+  workgroupBarrier();
+  return combinedMD;
+}
+{% else %}
+{% set mdStreamed = mdStreams is defined %}
+{% set mdStreams = mdStreams if mdStreams is defined else 1 %}
+{% set mdExtent = "WG" if mdStreams == 1 else "WG * " ~ mdStreams ~ "u" %}
+var<workgroup> partialM: array<f32, {{ mdExtent }}>;
+var<workgroup> partialD: array<f32, {{ mdExtent }}>;
+{% if mdStreamed %}
+// In-place fold of {{ mdStreams }} streams. The caller stores its per-thread
+// partials into partialM/partialD first and reads the merged pair of stream s
+// from slot s * WG afterwards.
+fn combine_partials_streams(lidx: u32) {
+  workgroupBarrier();
+  var stride = WG / 2u;
+  loop {
+    if (stride == 0u) {
+      break;
+    }
+    if (lidx < stride) {
+{% for s in range(mdStreams) %}
+      {
+        let slot = {{ s }}u * WG + lidx;
+        let m1 = partialM[slot];
+        let d1 = partialD[slot];
+        let m2 = partialM[slot + stride];
+        let d2 = partialD[slot + stride];
+        let mNew = max(m1, m2);
+        partialD[slot] = d1 * exp_shift(m1, mNew) + d2 * exp_shift(m2, mNew);
+        partialM[slot] = mNew;
+      }
+{% endfor %}
+    }
+    workgroupBarrier();
+    stride = stride / 2u;
+  }
+}
+{% else %}
+fn combine_partials(m: f32, d: f32, lidx: u32) -> vec2<f32> {
+  partialM[lidx] = m;
+  partialD[lidx] = d;
+  workgroupBarrier();
+  var stride = WG / 2u;
+  loop {
+    if (stride == 0u) {
+      break;
+    }
+    if (lidx < stride) {
+      let m1 = partialM[lidx];
+      let d1 = partialD[lidx];
+      let m2 = partialM[lidx + stride];
+      let d2 = partialD[lidx + stride];
+      let mNew = max(m1, m2);
+      partialD[lidx] = d1 * exp_shift(m1, mNew) + d2 * exp_shift(m2, mNew);
+      partialM[lidx] = mNew;
+    }
+    workgroupBarrier();
+    stride = stride / 2u;
+  }
+  let merged = vec2<f32>(partialM[0], partialD[0]);
+  // Trailing barrier so back-to-back calls cannot race a next call's partial
+  // stores against this call's reads of slot 0.
+  workgroupBarrier();
+  return merged;
+}
+{% endif %}
+{% endif %}
+{% if ATTN_SCALE_DIM is not defined %}{% set ATTN_SCALE_DIM = "HEAD_DIM" %}{% endif %}
+fn scale_value() -> f32 {
+  if (params.scale != 0.0) { return params.scale; }
+  return inverseSqrt(f32({{ ATTN_SCALE_DIM }}));
+}
+{{ sparse_schedule() }}
+@compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
+  let tile0 = wg.x * Q_TILE;
+  let head = wg.y % Q_HEADS;
+  let batch = wg.y / Q_HEADS;
+  let tid = lid.x;
+  let past = past_sequence_length(batch);
+  let kv_head = head / (Q_HEADS / KV_HEADS);
+  let scale = scale_value();
+  // Layouts cycle over the heads: four layouts and eight heads give S0 S1 S2 S3 S0 ...
+  let layout_id = head % NUM_LAYOUT;
+  let row_base = layout_id * ROW_STRIDE;
+  // A layout at or above the triangular dense-block count is treated as dense.
+  let dense_nnz = ROW_STRIDE * (ROW_STRIDE - 1u) / 2u;
+  let layout_is_dense = u32(block_row_indices[row_base + ROW_STRIDE - 1u]) >= dense_nnz;
+{% for j in range(qTile) %}
+  // Keys 0 .. q_abs are visible; the +1 makes the bound exclusive. A tile that runs past
+  // the last token carries inactive queries; they take no mask row and store nothing.
+  let live_{{ j }} = tile0 + {{ j }}u < params.seqLen;
+  let q_abs_{{ j }} = past + tile0 + {{ j }}u;
+  let key_bound_{{ j }} = q_abs_{{ j }} + 1u;
+  let mask_row_{{ j }} = q_abs_{{ j }} / SPARSE_BLOCK;
+{% endfor %}
+{% for j in range(qTile) %}
+  {
+{% if usesRotary %}
+    let q_base = ((batch * Q_HEADS + head) * params.seqLen + min(tile0 + {{ j }}u, params.seqLen - 1u)) * HEAD_DIM;
+{% else %}
+    let q_base = (batch * params.seqLen + min(tile0 + {{ j }}u, params.seqLen - 1u)) * Q_STRIDE + head * HEAD_DIM;
+{% endif %}
+    for (var dv: u32 = tid; dv < HEAD_VEC; dv = dv + WG) {
+      let qb = q_base + dv * 4u;
+{% if usesRotary %}
+      q_shared[{{ j }}u * HEAD_VEC + dv] = vec4<f32>(q_rotary[qb], q_rotary[qb + 1u], q_rotary[qb + 2u], q_rotary[qb + 3u]);
+{% else %}
+      q_shared[{{ j }}u * HEAD_VEC + dv] = vec4<f32>(f32(query[qb]), f32(query[qb + 1u]), f32(query[qb + 2u]), f32(query[qb + 3u]));
+{% endif %}
+      running_out[{{ j }}u * HEAD_VEC + dv] = vec4<f32>(0.0);
+    }
+  }
+{% endfor %}
+  workgroupBarrier();
+{% for j in range(qTile) %}
+  var runningMax_{{ j }} = -FLT_MAX;
+  var runningDenom_{{ j }} = 0.0;
+{% endfor %}
+  // Consecutive tokens span at most two mask rows, and every query of a row selects the
+  // same blocks, so one sweep per row covers the tile. A query contributes only to the
+  // sweep of its own row, which is why its online state is never merged across rows.
+{% if qTile > 1 %}
+  let row_first = q_abs_0 / SPARSE_BLOCK;
+  let row_last = mask_row_{{ qTile - 1 }};
+  for (var mask_row = row_first; mask_row <= row_last; mask_row = mask_row + 1u) {
+{% else %}
+  {
+    let mask_row = mask_row_0;
+{% endif %}
+    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.
+    let dense_row = layout_is_dense || (end - start == mask_row + 1u);
+    // The dense sweep runs to the furthest causal bound of the queries on this row; each
+    // query still drops the keys past its own.
+    var bound_max = 0u;
+{% for j in range(qTile) %}
+    if (live_{{ j }} && mask_row_{{ j }} == mask_row) { bound_max = max(bound_max, key_bound_{{ j }}); }
+{% endfor %}
+    let slot_count = select((end - start) * SPARSE_BLOCK, bound_max, dense_row);
+    var tileBase: u32 = 0u;
+    loop {
+      if (tileBase >= slot_count) {
+        break;
+      }
+      let slot = tileBase + tid;
+      // Resolving the slot is query-independent: the block list belongs to the row, and
+      // only the causal bound below differs between the queries sharing it.
+      var key = 0u;
+      var keyResolved = false;
+      if (slot < slot_count) {
+        if (!dense_row) {
+          // CSR is a boolean mask. Repeating a column therefore cannot repeat its
+          // probability mass, even though malformed-but-valid CSR may contain duplicates.
+          let entry = slot / SPARSE_BLOCK;
+          let block = u32(block_col_indices[layout_id * COL_STRIDE + start + entry]);
+          var duplicate = false;
+          for (var previous = 0u; previous < entry; previous = previous + 1u) {
+            duplicate = duplicate ||
+              u32(block_col_indices[layout_id * COL_STRIDE + start + previous]) == block;
+          }
+          key = block * SPARSE_BLOCK + slot % SPARSE_BLOCK;
+          keyResolved = !duplicate;
+        } else {
+          key = slot;
+          keyResolved = true;
+        }
+      }
+      // Always a row this cache holds, so an unresolved slot leaves a readable base rather
+      // than a stale or uninitialized one and the accumulation below needs no guard.
+      let base = ((batch * KV_HEADS + kv_head) * MAX_CACHE_SEQ + select(0u, key, keyResolved)) * HEAD_DIM;
+      key_rows[tid] = base;
+{% for j in range(qTile) %}
+      let allowed_{{ j }} = keyResolved && live_{{ j }} && mask_row_{{ j }} == mask_row && key < key_bound_{{ j }};
+      var acc_{{ j }} = 0.0;
+{% endfor %}
+      if (keyResolved) {
+        // One key vector, every query's score: this is the reuse the tile exists for.
+        let row_v = base / 4u;
+        for (var dv: u32 = 0u; dv < HEAD_VEC; dv = dv + 1u) {
+          let kv = vec4<f32>(present_key[row_v + dv]);
+{% for j in range(qTile) %}
+          acc_{{ j }} = acc_{{ j }} + dot(q_shared[{{ j }}u * HEAD_VEC + dv], kv);
+{% endfor %}
+        }
+      }
+      // One barrier tree folds all Q_TILE online-softmax partials.
+{% for j in range(qTile) %}
+      let score_{{ j }} = acc_{{ j }} * scale;
+      partialM[{{ j }}u * WG + tid] = select(-FLT_MAX, score_{{ j }}, allowed_{{ j }});
+      partialD[{{ j }}u * WG + tid] = select(0.0, 1.0, allowed_{{ j }});
+{% endfor %}
+      combine_partials_streams(tid);
+{% for j in range(qTile) %}
+      let tile_{{ j }} = vec2<f32>(partialM[{{ j }}u * WG], partialD[{{ j }}u * WG]);
+{% endfor %}
+{% for j in range(qTile) %}
+      let newMax_{{ j }} = max(runningMax_{{ j }}, tile_{{ j }}.x);
+      let correction_{{ j }} = exp_shift(runningMax_{{ j }}, newMax_{{ j }});
+      runningDenom_{{ j }} = runningDenom_{{ j }} * correction_{{ j }}
+        + tile_{{ j }}.y * exp_shift(tile_{{ j }}.x, newMax_{{ j }});
+      runningMax_{{ j }} = newMax_{{ j }};
+      probs[{{ j }}u * WG + tid] = select(0.0, exp_shift(score_{{ j }}, newMax_{{ j }}), allowed_{{ j }});
+{% endfor %}
+      workgroupBarrier();
+      // running_out[j][d] is owned by the same thread across every tile (tid = d mod WG),
+      // so this rescale-and-accumulate needs no further synchronization. One value vector
+      // serves every query, which is the other half of the tile's reuse; a key outside a
+      // query's causal bound carries prob 0 and is multiplied away.
+      let tileCount = min(WG, slot_count - tileBase);
+      for (var dv: u32 = tid; dv < HEAD_VEC; dv = dv + WG) {
+{% for j in range(qTile) %}
+        var vSum_{{ j }} = vec4<f32>(0.0);
+{% endfor %}
+        for (var i: u32 = 0u; i < tileCount; i = i + 1u) {
+          let vv = vec4<f32>(present_value[key_rows[i] / 4u + dv]);
+{% for j in range(qTile) %}
+          vSum_{{ j }} = vSum_{{ j }} + probs[{{ j }}u * WG + i] * vv;
+{% endfor %}
+        }
+{% for j in range(qTile) %}
+        running_out[{{ j }}u * HEAD_VEC + dv] = running_out[{{ j }}u * HEAD_VEC + dv] * correction_{{ j }} + vSum_{{ j }};
+{% endfor %}
+      }
+      // Orders this tile's probs and key_rows reads before the next tile writes them.
+      workgroupBarrier();
+      tileBase = tileBase + WG;
+    }
+  }
+  // An all-masked row becomes a uniform average over causal history rather than zero.
+{% for j in range(qTile) %}
+  if (live_{{ j }}) {
+    let hasMass_{{ j }} = runningDenom_{{ j }} > 0.0;
+    let invDenom_{{ j }} = select(0.0, 1.0 / runningDenom_{{ j }}, hasMass_{{ j }});
+    let out_base_{{ j }} = (batch * params.seqLen + tile0 + {{ j }}u) * Q_HIDDEN + head * HEAD_DIM;
+    for (var dv: u32 = tid; dv < HEAD_VEC; dv = dv + WG) {
+      var out_v: vec4<f32>;
+      if (hasMass_{{ j }}) {
+        out_v = running_out[{{ j }}u * HEAD_VEC + dv] * invDenom_{{ j }};
+      } else {
+        var total = vec4<f32>(0.0);
+        for (var key = 0u; key < key_bound_{{ j }}; key = key + 1u) {
+          let row_v = (((batch * KV_HEADS + kv_head) * MAX_CACHE_SEQ + key) * HEAD_DIM) / 4u;
+          total = total + vec4<f32>(present_value[row_v + dv]);
+        }
+        out_v = total / f32(key_bound_{{ j }});
+      }
+      let ob = out_base_{{ j }} + dv * 4u;
+      output[ob] = {{ scalar }}(out_v.x);
+      output[ob + 1u] = {{ scalar }}(out_v.y);
+      output[ob + 2u] = {{ scalar }}(out_v.z);
+      output[ob + 3u] = {{ scalar }}(out_v.w);
+    }
+  }
+{% endfor %}
+}

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

	@@ -0,0 +1,125 @@

+{% 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
+// length to zero; a padded prompt row must not be read as if it had history.
+fn past_sequence_length(batch: u32) -> u32 {
+  if (u32(total_sequence_length[0]) == params.seqLen) {
+    return 0u;
+  }
+  let total = u32(key_total_sequence_lengths[batch]);
+  return select(0u, total - params.seqLen, total >= params.seqLen);
+}
+{%- endmacro %}
+{% macro sparse_rotary(interleaved) %}
+// Which cos/sin entry a component uses, and which member of its rotation pair it is.
+// The two layouts differ only here: the NeoX split pairs d with d + ROTARY_HALF, and the
+// interleaved layout pairs the even lane with the odd one beside it.
+fn rotary_slot(d: u32) -> u32 {
+{% if interleaved %}
+  return d / 2u;
+{% else %}
+  return select(d - ROTARY_HALF, d, d < ROTARY_HALF);
+{% endif %}
+}
+fn rotary_partner(d: u32) -> u32 {
+{% if interleaved %}
+  return select(d - 1u, d + 1u, (d % 2u) == 0u);
+{% else %}
+  return select(d - ROTARY_HALF, d + ROTARY_HALF, d < ROTARY_HALF);
+{% endif %}
+}
+fn rotary_is_first(d: u32) -> bool {
+{% if interleaved %}
+  return (d % 2u) == 0u;
+{% else %}
+  return d < ROTARY_HALF;
+{% endif %}
+}
+// One component of the rotation, written from the component's own point of view so a
+// single expression covers both members of the pair: the leading one subtracts its
+// partner's sine term, the trailing one adds it.
+fn rotary_value(own: f32, partner: f32, cs: f32, sn: f32, first: bool) -> f32 {
+  return select(own * cs + partner * sn, own * cs - partner * sn, first);
+}
+{%- endmacro %}
+{% if usesF16 %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+// com.microsoft.SparseAttention, KV append pass.
+// past_key and present_key are the same allocation, so this writes only the new rows:
+// one invocation per (batch, kv head, new token, component) of the BNSH cache, landing
+// at absolute position past + t. Nothing rewrites the retained history.
+const KV_HEADS: u32 = {{ kvNumHeads }}u;
+const HEAD_DIM: u32 = {{ headSize }}u;
+const MAX_CACHE_SEQ: u32 = {{ maxCacheSeq }}u;
+{% if packedQkv %}
+// Packed layout: query carries [Q | K | V] on one row, so K starts after the query
+// heads and V after the key heads.
+const Q_STRIDE: u32 = {{ packedStride }}u;
+const PACKED_K_OFFSET: u32 = {{ numHeads * headSize }}u;
+const PACKED_V_OFFSET: u32 = {{ (numHeads + kvNumHeads) * headSize }}u;
+{% else %}
+const KV_HIDDEN: u32 = {{ kvNumHeads * headSize }}u;
+{% endif %}
+const WG: u32 = {{ appendWorkgroupSize }}u;
+{% if usesRotary %}
+const ROTARY_HALF: u32 = {{ rotaryHalf }}u;
+const ROTARY_DIM: u32 = {{ rotaryDim }}u;
+{% endif %}
+{{ sparse_schedule() }}
+{% if usesRotary %}
+{{ sparse_rotary(rotaryInterleaved) }}
+{% 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;
+  }
+  let d = index % HEAD_DIM;
+  let t = (index / HEAD_DIM) % params.seqLen;
+  let kv_head = (index / (HEAD_DIM * params.seqLen)) % KV_HEADS;
+  let batch = index / (HEAD_DIM * params.seqLen * KV_HEADS);
+  let position = past_sequence_length(batch) + t;
+  let dst = ((batch * KV_HEADS + kv_head) * MAX_CACHE_SEQ + position) * HEAD_DIM + d;
+{% if packedQkv %}
+  let row = (batch * params.seqLen + t) * Q_STRIDE;
+  let k_base = row + PACKED_K_OFFSET + kv_head * HEAD_DIM;
+  let v_base = row + PACKED_V_OFFSET + kv_head * HEAD_DIM;
+{% else %}
+  let k_base = ((batch * params.seqLen + t) * KV_HIDDEN) + kv_head * HEAD_DIM;
+  let v_base = k_base;
+{% endif %}
+{% if usesRotary %}
+  // The key is rotated at its own absolute position before it enters the cache, so the
+  // cache holds post-rotary keys and the attention pass never rotates again.
+  if (d < ROTARY_DIM) {
+    let slot = rotary_slot(d);
+    let cs = f32(cos_cache[position * ROTARY_HALF + slot]);
+    let sn = f32(sin_cache[position * ROTARY_HALF + slot]);
+    let own = f32({{ kvSource }}[k_base + d]);
+    let partner = f32({{ kvSource }}[k_base + rotary_partner(d)]);
+    present_key[dst] = {{ scalar }}(rotary_value(own, partner, cs, sn, rotary_is_first(d)));
+  } else {
+    present_key[dst] = {{ scalar }}({{ kvSource }}[k_base + d]);
+  }
+{% else %}
+  present_key[dst] = {{ scalar }}({{ kvSource }}[k_base + d]);
+{% endif %}
+  present_value[dst] = {{ scalar }}({{ vSource }}[v_base + d]);
+}

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

	@@ -0,0 +1,102 @@

+{% 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
+// length to zero; a padded prompt row must not be read as if it had history.
+fn past_sequence_length(batch: u32) -> u32 {
+  if (u32(total_sequence_length[0]) == params.seqLen) {
+    return 0u;
+  }
+  let total = u32(key_total_sequence_lengths[batch]);
+  return select(0u, total - params.seqLen, total >= params.seqLen);
+}
+{%- endmacro %}
+{% macro sparse_rotary(interleaved) %}
+// Which cos/sin entry a component uses, and which member of its rotation pair it is.
+// The two layouts differ only here: the NeoX split pairs d with d + ROTARY_HALF, and the
+// interleaved layout pairs the even lane with the odd one beside it.
+fn rotary_slot(d: u32) -> u32 {
+{% if interleaved %}
+  return d / 2u;
+{% else %}
+  return select(d - ROTARY_HALF, d, d < ROTARY_HALF);
+{% endif %}
+}
+fn rotary_partner(d: u32) -> u32 {
+{% if interleaved %}
+  return select(d - 1u, d + 1u, (d % 2u) == 0u);
+{% else %}
+  return select(d - ROTARY_HALF, d + ROTARY_HALF, d < ROTARY_HALF);
+{% endif %}
+}
+fn rotary_is_first(d: u32) -> bool {
+{% if interleaved %}
+  return (d % 2u) == 0u;
+{% else %}
+  return d < ROTARY_HALF;
+{% endif %}
+}
+// One component of the rotation, written from the component's own point of view so a
+// single expression covers both members of the pair: the leading one subtracts its
+// partner's sine term, the trailing one adds it.
+fn rotary_value(own: f32, partner: f32, cs: f32, sn: f32, first: bool) -> f32 {
+  return select(own * cs + partner * sn, own * cs - partner * sn, first);
+}
+{%- endmacro %}
+{% if usesF16 %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+// com.microsoft.SparseAttention, query rotary pass.
+// Rotates Q into a (batch, heads, seq, dim) scratch buffer so the attention pass binds
+// one query tensor instead of the query plus both caches. That keeps the attention pass
+// at the eight storage buffers WebGPU guarantees; folding the rotation in would need
+// ten.
+const Q_HEADS: u32 = {{ numHeads }}u;
+const HEAD_DIM: u32 = {{ headSize }}u;
+const ROTARY_HALF: u32 = {{ rotaryHalf }}u;
+const ROTARY_DIM: u32 = {{ rotaryDim }}u;
+const Q_STRIDE: u32 = {{ packedStride if packedQkv else numHeads * headSize }}u;
+const WG: u32 = {{ appendWorkgroupSize }}u;
+{{ sparse_schedule() }}
+{{ 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;
+  }
+  let d = index % HEAD_DIM;
+  let t = (index / HEAD_DIM) % params.seqLen;
+  let head = (index / (HEAD_DIM * params.seqLen)) % Q_HEADS;
+  let batch = index / (HEAD_DIM * params.seqLen * Q_HEADS);
+  // A query token's rotary position is the same absolute position the attention pass
+  // uses for causality: the history in front of it plus its own offset.
+  let position = past_sequence_length(batch) + t;
+  // The packed query's Q block starts at offset zero, so this base serves both layouts.
+  let q_base = (batch * params.seqLen + t) * Q_STRIDE + head * HEAD_DIM;
+  if (d < ROTARY_DIM) {
+    let slot = rotary_slot(d);
+    let cs = f32(cos_cache[position * ROTARY_HALF + slot]);
+    let sn = f32(sin_cache[position * ROTARY_HALF + slot]);
+    let own = f32(query[q_base + d]);
+    let partner = f32(query[q_base + rotary_partner(d)]);
+    // Rotary Q is materialized in the operator dtype. Round to f16 here when T
+    // is float16, then widen only for the f32 accumulation scratch buffer.
+    q_rotary[index] = f32({{ scalar }}(rotary_value(own, partner, cs, sn, rotary_is_first(d))));
+  } else {
+    q_rotary[index] = f32(query[q_base + d]);
+  }
+}

build/webgpu/test.json ADDED Viewed

	@@ -0,0 +1,745 @@

+{
+  "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]
+  },
+  "cases": [
+    {
+      "name": "prompt_separate_two_layouts",
+      "provenance": {
+        "notes": "The two-layout CSR example from the ONNX Runtime operator documentation, using separate Q/K/V tensors."
+      },
+      "attrs": { "num_heads": 4, "kv_num_heads": 2, "sparse_block_size": 16 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [2, 32, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [2, 32, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": { "dtype": "float32", "shape": [2, 2, 64, 8], "data": { "kind": "constant", "value": 0.0 } },
+        "pastValueT": { "dtype": "float32", "shape": [2, 2, 64, 8], "data": { "kind": "constant", "value": 0.0 } },
+        "blockRowIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [0, 1, 3, 5, 8, 0, 1, 3, 6, 9] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 9],
+          "data": {
+            "kind": "values",
+            "values": { "$ref": "#/fixtureArrays/prompt_separate_two_layouts_input_blockColIndicesT" }
+          }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [32] } },
+        "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [32, 32] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [2, 32, 32], "tolerance": 0.00002, "relTolerance": 0.0002 },
+        "pastKeyT": { "dtype": "float32", "shape": [2, 2, 64, 8], "tolerance": 0.00002, "relTolerance": 0.0002 },
+        "pastValueT": { "dtype": "float32", "shape": [2, 2, 64, 8], "tolerance": 0.00002, "relTolerance": 0.0002 }
+      }
+    },
+    {
+      "name": "prompt_packed_two_layouts",
+      "provenance": { "notes": "Packed [Q|K|V] form of the documented two-layout prompt case." },
+      "attrs": { "num_heads": 4, "kv_num_heads": 2, "sparse_block_size": 16 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 32, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.27, "scale": 0.5 }
+        },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 2, 64, 8], "data": { "kind": "constant", "value": 0.0 } },
+        "pastValueT": { "dtype": "float32", "shape": [1, 2, 64, 8], "data": { "kind": "constant", "value": 0.0 } },
+        "blockRowIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [0, 1, 3, 5, 8, 0, 1, 3, 6, 9] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 9],
+          "data": {
+            "kind": "values",
+            "values": { "$ref": "#/fixtureArrays/prompt_separate_two_layouts_input_blockColIndicesT" }
+          }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [32] } },
+        "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [32] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 32, 32], "tolerance": 0.00002, "relTolerance": 0.0002 },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 2, 64, 8], "tolerance": 0.00002, "relTolerance": 0.0002 },
+        "pastValueT": { "dtype": "float32", "shape": [1, 2, 64, 8], "tolerance": 0.00002, "relTolerance": 0.0002 }
+      }
+    },
+    {
+      "name": "prompt_separate_rotary_neox",
+      "provenance": {
+        "notes": "Full-width NeoX rotary with the minimum accepted head width; rotary_interleaved values other than 1 retain the NeoX layout."
+      },
+      "attrs": { "num_heads": 4, "kv_num_heads": 2, "sparse_block_size": 16, "do_rotary": 1, "rotary_interleaved": 2 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 32, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 2, 64, 16], "data": { "kind": "constant", "value": 0.0 } },
+        "pastValueT": { "dtype": "float32", "shape": [1, 2, 64, 16], "data": { "kind": "constant", "value": 0.0 } },
+        "blockRowIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [0, 1, 3, 5, 8, 0, 1, 3, 6, 9] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 9],
+          "data": {
+            "kind": "values",
+            "values": { "$ref": "#/fixtureArrays/prompt_separate_two_layouts_input_blockColIndicesT" }
+          }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [32] } },
+        "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [32] } },
+        "cosCacheT": { "dtype": "float32", "shape": [64, 8], "data": { "kind": "constant", "value": 0.8 } },
+        "sinCacheT": { "dtype": "float32", "shape": [64, 8], "data": { "kind": "constant", "value": 0.6 } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 32, 64], "tolerance": 0.00003, "relTolerance": 0.0003 },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 2, 64, 16], "tolerance": 0.00003, "relTolerance": 0.0003 },
+        "pastValueT": { "dtype": "float32", "shape": [1, 2, 64, 16], "tolerance": 0.00002, "relTolerance": 0.0002 }
+      }
+    },
+    {
+      "name": "prompt_packed_partial_rotary_interleaved",
+      "provenance": {
+        "notes": "Packed QKV with interleaved rotary over the first 16 of 32 head channels; ONNX Runtime copies the remainder through."
+      },
+      "attrs": { "num_heads": 4, "kv_num_heads": 2, "sparse_block_size": 16, "do_rotary": 1, "rotary_interleaved": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 32, 256],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.19, "scale": 0.5 }
+        },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 2, 64, 32], "data": { "kind": "constant", "value": 0.0 } },
+        "pastValueT": { "dtype": "float32", "shape": [1, 2, 64, 32], "data": { "kind": "constant", "value": 0.0 } },
+        "blockRowIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [0, 1, 3, 5, 8, 0, 1, 3, 6, 9] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 9],
+          "data": {
+            "kind": "values",
+            "values": { "$ref": "#/fixtureArrays/prompt_separate_two_layouts_input_blockColIndicesT" }
+          }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [32] } },
+        "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [32] } },
+        "cosCacheT": { "dtype": "float32", "shape": [64, 8], "data": { "kind": "constant", "value": 0.8 } },
+        "sinCacheT": { "dtype": "float32", "shape": [64, 8], "data": { "kind": "constant", "value": 0.6 } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 32, 128], "tolerance": 0.00003, "relTolerance": 0.0003 },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 2, 64, 32], "tolerance": 0.00003, "relTolerance": 0.0003 },
+        "pastValueT": { "dtype": "float32", "shape": [1, 2, 64, 32], "tolerance": 0.00002, "relTolerance": 0.0002 }
+      }
+    },
+    {
+      "name": "decode_separate_two_layouts",
+      "provenance": {
+        "notes": "Two decode rows with different history lengths, so the scalar maximum and per-row key lengths have distinct jobs."
+      },
+      "attrs": { "num_heads": 4, "kv_num_heads": 2, "sparse_block_size": 16 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 1, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [2, 1, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [2, 1, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 2, 64, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.017, "scale": 0.25 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 2, 64, 8],
+          "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, 8, 0, 1, 3, 6, 9] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 9],
+          "data": {
+            "kind": "values",
+            "values": { "$ref": "#/fixtureArrays/prompt_separate_two_layouts_input_blockColIndicesT" }
+          }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [41] } },
+        "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [41, 34] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [2, 1, 32], "tolerance": 0.00003, "relTolerance": 0.0003 },
+        "pastKeyT": { "dtype": "float32", "shape": [2, 2, 64, 8], "tolerance": 0.00002, "relTolerance": 0.0002 },
+        "pastValueT": { "dtype": "float32", "shape": [2, 2, 64, 8], "tolerance": 0.00002, "relTolerance": 0.0002 }
+      }
+    },
+    {
+      "name": "decode_packed_rotary_neox",
+      "provenance": { "notes": "Packed decode with full-width NeoX rotary at each row's absolute position." },
+      "attrs": { "num_heads": 4, "kv_num_heads": 2, "sparse_block_size": 16, "do_rotary": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 1, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 2, 64, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.017, "scale": 0.25 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 2, 64, 16],
+          "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, 8, 0, 1, 3, 6, 9] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 9],
+          "data": {
+            "kind": "values",
+            "values": { "$ref": "#/fixtureArrays/prompt_separate_two_layouts_input_blockColIndicesT" }
+          }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [41] } },
+        "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [41, 34] } },
+        "cosCacheT": { "dtype": "float32", "shape": [64, 8], "data": { "kind": "constant", "value": 0.8 } },
+        "sinCacheT": { "dtype": "float32", "shape": [64, 8], "data": { "kind": "constant", "value": 0.6 } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [2, 1, 64], "tolerance": 0.00003, "relTolerance": 0.0003 },
+        "pastKeyT": { "dtype": "float32", "shape": [2, 2, 64, 16], "tolerance": 0.00003, "relTolerance": 0.0003 },
+        "pastValueT": { "dtype": "float32", "shape": [2, 2, 64, 16], "tolerance": 0.00002, "relTolerance": 0.0002 }
+      }
+    },
+    {
+      "name": "chunked_separate_two_layouts",
+      "provenance": {
+        "notes": "Eight new tokens over 24 cached tokens, covering the non-prompt sequence_length > 1 path."
+      },
+      "attrs": { "num_heads": 4, "kv_num_heads": 2, "sparse_block_size": 16 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 8, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [2, 8, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [2, 8, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 2, 64, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.017, "scale": 0.25 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 2, 64, 8],
+          "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, 8, 0, 1, 3, 6, 9] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 9],
+          "data": {
+            "kind": "values",
+            "values": { "$ref": "#/fixtureArrays/prompt_separate_two_layouts_input_blockColIndicesT" }
+          }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [32] } },
+        "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [32, 32] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [2, 8, 32], "tolerance": 0.00003, "relTolerance": 0.0003 },
+        "pastKeyT": { "dtype": "float32", "shape": [2, 2, 64, 8], "tolerance": 0.00002, "relTolerance": 0.0002 },
+        "pastValueT": { "dtype": "float32", "shape": [2, 2, 64, 8], "tolerance": 0.00002, "relTolerance": 0.0002 }
+      }
+    },
+    {
+      "name": "prompt_all_rows_gqa_four_to_one",
+      "provenance": {
+        "notes": "A 64-token prompt reaches every documented block row while four query heads share each KV head."
+      },
+      "attrs": { "num_heads": 8, "kv_num_heads": 2, "sparse_block_size": 16 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 64, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 64, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 64, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 2, 64, 8], "data": { "kind": "constant", "value": 0.0 } },
+        "pastValueT": { "dtype": "float32", "shape": [1, 2, 64, 8], "data": { "kind": "constant", "value": 0.0 } },
+        "blockRowIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [0, 1, 3, 5, 8, 0, 1, 3, 6, 9] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 9],
+          "data": {
+            "kind": "values",
+            "values": { "$ref": "#/fixtureArrays/prompt_separate_two_layouts_input_blockColIndicesT" }
+          }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [64] } },
+        "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [64] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 64, 64], "tolerance": 0.00003, "relTolerance": 0.0003 },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 2, 64, 8], "tolerance": 0.00002, "relTolerance": 0.0002 },
+        "pastValueT": { "dtype": "float32", "shape": [1, 2, 64, 8], "tolerance": 0.00002, "relTolerance": 0.0002 }
+      }
+    },
+    {
+      "name": "prompt_block32_head_dim16",
+      "provenance": {
+        "notes": "Block size 32 and head width 16 exercise independent key-tile and output-channel tails."
+      },
+      "attrs": { "num_heads": 4, "kv_num_heads": 2, "sparse_block_size": 32 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 64, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 2, 128, 16], "data": { "kind": "constant", "value": 0.0 } },
+        "pastValueT": { "dtype": "float32", "shape": [1, 2, 128, 16], "data": { "kind": "constant", "value": 0.0 } },
+        "blockRowIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [0, 1, 3, 5, 8, 0, 1, 3, 6, 9] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 9],
+          "data": {
+            "kind": "values",
+            "values": { "$ref": "#/fixtureArrays/prompt_separate_two_layouts_input_blockColIndicesT" }
+          }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [64] } },
+        "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [64] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 64, 64], "tolerance": 0.00003, "relTolerance": 0.0003 },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 2, 128, 16], "tolerance": 0.00002, "relTolerance": 0.0002 },
+        "pastValueT": { "dtype": "float32", "shape": [1, 2, 128, 16], "tolerance": 0.00002, "relTolerance": 0.0002 }
+      }
+    },
+    {
+      "name": "prompt_explicit_scale",
+      "provenance": { "notes": "Explicit non-default score scale." },
+      "attrs": { "num_heads": 4, "kv_num_heads": 2, "sparse_block_size": 16, "scale": 0.25 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 32, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 32, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 2, 64, 8], "data": { "kind": "constant", "value": 0.0 } },
+        "pastValueT": { "dtype": "float32", "shape": [1, 2, 64, 8], "data": { "kind": "constant", "value": 0.0 } },
+        "blockRowIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [0, 1, 3, 5, 8, 0, 1, 3, 6, 9] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 9],
+          "data": {
+            "kind": "values",
+            "values": { "$ref": "#/fixtureArrays/prompt_separate_two_layouts_input_blockColIndicesT" }
+          }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [32] } },
+        "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [32] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 32, 32], "tolerance": 0.00002, "relTolerance": 0.0002 },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 2, 64, 8], "tolerance": 0.00002, "relTolerance": 0.0002 },
+        "pastValueT": { "dtype": "float32", "shape": [1, 2, 64, 8], "tolerance": 0.00002, "relTolerance": 0.0002 }
+      }
+    },
+    {
+      "name": "f16_prompt_separate",
+      "provenance": { "notes": "Meaningful float16 prompt; scores and online-softmax accumulators remain float32." },
+      "attrs": { "num_heads": 4, "kv_num_heads": 2, "sparse_block_size": 16 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [1, 32, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [1, 32, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [1, 32, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": { "dtype": "float16", "shape": [1, 2, 64, 8], "data": { "kind": "constant", "value": 0.0 } },
+        "pastValueT": { "dtype": "float16", "shape": [1, 2, 64, 8], "data": { "kind": "constant", "value": 0.0 } },
+        "blockRowIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [0, 1, 3, 5, 8, 0, 1, 3, 6, 9] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 9],
+          "data": {
+            "kind": "values",
+            "values": { "$ref": "#/fixtureArrays/prompt_separate_two_layouts_input_blockColIndicesT" }
+          }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [32] } },
+        "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [32] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float16", "shape": [1, 32, 32], "tolerance": 0.003, "relTolerance": 0.01 },
+        "pastKeyT": { "dtype": "float16", "shape": [1, 2, 64, 8], "tolerance": 0.001, "relTolerance": 0.001 },
+        "pastValueT": { "dtype": "float16", "shape": [1, 2, 64, 8], "tolerance": 0.001, "relTolerance": 0.001 }
+      }
+    },
+    {
+      "name": "f16_decode_packed_partial_rotary",
+      "provenance": { "notes": "Float16 packed decode with partial NeoX rotary and unequal per-batch history lengths." },
+      "attrs": { "num_heads": 4, "kv_num_heads": 2, "sparse_block_size": 16, "do_rotary": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [2, 1, 256],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "pastKeyT": {
+          "dtype": "float16",
+          "shape": [2, 2, 64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.017, "scale": 0.25 }
+        },
+        "pastValueT": {
+          "dtype": "float16",
+          "shape": [2, 2, 64, 32],
+          "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, 8, 0, 1, 3, 6, 9] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 9],
+          "data": {
+            "kind": "values",
+            "values": { "$ref": "#/fixtureArrays/prompt_separate_two_layouts_input_blockColIndicesT" }
+          }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [41] } },
+        "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [41, 34] } },
+        "cosCacheT": { "dtype": "float16", "shape": [64, 8], "data": { "kind": "constant", "value": 0.8 } },
+        "sinCacheT": { "dtype": "float16", "shape": [64, 8], "data": { "kind": "constant", "value": 0.6 } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float16", "shape": [2, 1, 128], "tolerance": 0.003, "relTolerance": 0.01 },
+        "pastKeyT": { "dtype": "float16", "shape": [2, 2, 64, 32], "tolerance": 0.002, "relTolerance": 0.005 },
+        "pastValueT": { "dtype": "float16", "shape": [2, 2, 64, 32], "tolerance": 0.001, "relTolerance": 0.001 }
+      }
+    },
+    {
+      "name": "dense_layout_ignores_malformed_columns",
+      "provenance": {
+        "notes": "ORT classifies a layout from its final CSR count. Once dense, repeated non-prefix column contents are ignored."
+      },
+      "attrs": { "num_heads": 1, "kv_num_heads": 1, "sparse_block_size": 16 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 32, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 32, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 32, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 1, 32, 8], "data": { "kind": "constant", "value": 0.0 } },
+        "pastValueT": { "dtype": "float32", "shape": [1, 1, 32, 8], "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": [1, 1, 1] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [32] } },
+        "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [32] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 32, 8], "tolerance": 0.00002, "relTolerance": 0.0002 },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 1, 32, 8], "tolerance": 0.00002, "relTolerance": 0.0002 },
+        "pastValueT": { "dtype": "float32", "shape": [1, 1, 32, 8], "tolerance": 0.00002, "relTolerance": 0.0002 }
+      }
+    },
+    {
+      "name": "sparse_duplicate_columns_are_idempotent",
+      "provenance": {
+        "notes": "A sparse row repeats block 0. CSR encodes a boolean mask, so the duplicate must not duplicate probability mass."
+      },
+      "attrs": { "num_heads": 1, "kv_num_heads": 1, "sparse_block_size": 16 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 48, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 48, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 48, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 1, 48, 8], "data": { "kind": "constant", "value": 0.0 } },
+        "pastValueT": { "dtype": "float32", "shape": [1, 1, 48, 8], "data": { "kind": "constant", "value": 0.0 } },
+        "blockRowIndicesT": { "dtype": "int32", "shape": [1, 4], "data": { "kind": "values", "values": [0, 1, 3, 5] } },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [1, 5],
+          "data": { "kind": "values", "values": [0, 0, 1, 0, 0] }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [48] } },
+        "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [48] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 48, 8], "tolerance": 0.00002, "relTolerance": 0.0002 },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 1, 48, 8], "tolerance": 0.00002, "relTolerance": 0.0002 },
+        "pastValueT": { "dtype": "float32", "shape": [1, 1, 48, 8], "tolerance": 0.00002, "relTolerance": 0.0002 }
+      }
+    },
+    {
+      "name": "empty_csr_row_uniform_scalar_total",
+      "provenance": {
+        "notes": "Pinned all-masked-row behavior from ORT's lowest-sentinel softmax: one causal key yields its value, not zero. Also covers scalar total_sequence_length and a zero-width column tensor."
+      },
+      "attrs": { "num_heads": 1, "kv_num_heads": 1, "sparse_block_size": 16 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 8],
+          "data": { "kind": "values", "values": [0.5, -0.5, 1.0, -1.0, 2.0, -2.0, 3.0, -3.0] }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 8],
+          "data": { "kind": "values", "values": [8.0, 7.0, 6.0, 5.0, 4.0, 3.0, 2.0, 1.0] }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 8],
+          "data": { "kind": "values", "values": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0] }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 16, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.25 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 16, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.25 }
+        },
+        "blockRowIndicesT": { "dtype": "int32", "shape": [1, 2], "data": { "kind": "values", "values": [0, 0] } },
+        "blockColIndicesT": { "dtype": "int32", "shape": [1, 0], "data": { "kind": "values", "values": [] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [], "data": { "kind": "values", "values": [1] } },
+        "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1] } }
+      },
+      "outputs": {
+        "outputT": {
+          "dtype": "float32",
+          "shape": [1, 1, 8],
+          "data": { "kind": "values", "values": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0] },
+          "tolerance": 0
+        },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 1, 16, 8], "tolerance": 0 },
+        "pastValueT": { "dtype": "float32", "shape": [1, 1, 16, 8], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "optional_rotary_input_ignored_when_attribute_is_not_one",
+      "provenance": {
+        "notes": "The CPU kernel enables do_rotary only for the exact integer value 1. Otherwise optional rotary inputs are ignored and need not be paired."
+      },
+      "attrs": { "num_heads": 1, "kv_num_heads": 1, "sparse_block_size": 16, "do_rotary": 2, "rotary_interleaved": 2 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 1, 16, 8], "data": { "kind": "constant", "value": 0.0 } },
+        "pastValueT": { "dtype": "float32", "shape": [1, 1, 16, 8], "data": { "kind": "constant", "value": 0.0 } },
+        "blockRowIndicesT": { "dtype": "int32", "shape": [1, 2], "data": { "kind": "values", "values": [0, 1] } },
+        "blockColIndicesT": { "dtype": "int32", "shape": [1, 1], "data": { "kind": "values", "values": [0] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1] } },
+        "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1] } },
+        "cosCacheT": { "dtype": "float32", "shape": [1, 1], "data": { "kind": "constant", "value": 0.0 } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 1, 8], "tolerance": 0.000001, "relTolerance": 0.000001 },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 1, 16, 8], "tolerance": 0.000001, "relTolerance": 0.000001 },
+        "pastValueT": { "dtype": "float32", "shape": [1, 1, 16, 8], "tolerance": 0.000001, "relTolerance": 0.000001 }
+      }
+    },
+    {
+      "name": "zero_width_rotary_is_a_noop",
+      "provenance": {
+        "notes": "ONNX Runtime accepts zero-width paired caches with do_rotary=1; the resulting zero-dimensional rotary transform copies Q and K through."
+      },
+      "attrs": { "num_heads": 1, "kv_num_heads": 1, "sparse_block_size": 16, "do_rotary": 1, "rotary_interleaved": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 1, 16, 16], "data": { "kind": "constant", "value": 0.0 } },
+        "pastValueT": { "dtype": "float32", "shape": [1, 1, 16, 16], "data": { "kind": "constant", "value": 0.0 } },
+        "blockRowIndicesT": { "dtype": "int32", "shape": [1, 2], "data": { "kind": "values", "values": [0, 1] } },
+        "blockColIndicesT": { "dtype": "int32", "shape": [1, 1], "data": { "kind": "values", "values": [0] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1] } },
+        "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1] } },
+        "cosCacheT": { "dtype": "float32", "shape": [1, 0], "data": { "kind": "values", "values": [] } },
+        "sinCacheT": { "dtype": "float32", "shape": [1, 0], "data": { "kind": "values", "values": [] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 1, 16], "tolerance": 0.000001, "relTolerance": 0.000001 },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 1, 16, 16], "tolerance": 0.000001, "relTolerance": 0.000001 },
+        "pastValueT": { "dtype": "float32", "shape": [1, 1, 16, 16], "tolerance": 0.000001, "relTolerance": 0.000001 }
+      }
+    }
+  ]
+}