sync 6fdf6301e2bb

Browse files

Files changed (9) hide show

README.md +19 -2
build/webgpu/bench.json +1277 -1
build/webgpu/manifest.json +277 -59
build/webgpu/metadata.json +15 -11
build/webgpu/sparse-attention-sgmat.wgsl.jinja +87 -61
build/webgpu/sparse-attention.wgsl.jinja +53 -67
build/webgpu/sparse-kv-append.wgsl.jinja +8 -8
build/webgpu/sparse-q-rotary.wgsl.jinja +8 -8
build/webgpu/test.json +0 -0

README.md CHANGED Viewed

@@ -60,6 +60,23 @@ Attributes and default values (overridable per request):
 | `T` | `float32`, `float16` |
 | `M` | `int32` |
 ## Device requirements
 Some implementation variants require `subgroup-matrix` and `subgroups`. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.
@@ -69,7 +86,7 @@ Some implementation variants require `subgroup-matrix` and `subgroups`. These ar
 - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
-- [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 - [`sparse-attention-sgmat.wgsl.jinja`](build/webgpu/sparse-attention-sgmat.wgsl.jinja)
 - [`sparse-attention.wgsl.jinja`](build/webgpu/sparse-attention.wgsl.jinja)
 - [`sparse-kv-append.wgsl.jinja`](build/webgpu/sparse-kv-append.wgsl.jinja)
@@ -78,7 +95,7 @@ Some implementation variants require `subgroup-matrix` and `subgroups`. These ar
 ## Use with `@huggingface/kernels`
 ```sh
-npm install --save-exact @huggingface/kernels@0.0.1-preview.2
 ```
 Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.

 | `T` | `float32`, `float16` |
 | `M` | `int32` |
+## Implementation variants
+One implementation is selected per call from the device capabilities, the request shapes and the dtypes; these notes say what each one covers.
+- `separate` — Portable online attention shares each key/value load across query rows. Spare lanes partition value accumulation, with the partition count bounded by the head width, workgroup size, and available workgroup storage. Single-partition staging requires one output column per lane and includes the staging array in its storage budget.
+- `separate_sgmat` — Float32 subgroup-matrix online attention computes each score tile once, rescales output fragments through existing score scratch, and directly loads complete query tiles. Only a partial query tile is staged; key-tile width and fragment-bank capacity follow the reported workgroup-storage limit.
+- `separate_sgmat_tail` — For a prompt, a thin query tail uses the shared portable template after complete matrix tiles. History calls retain the full matrix computation. Tail size and value partitions follow workgroup capacity. Heads narrower than its query tile use the full matrix path, since their shorter value contraction does not repay a separate portable dispatch.
+- `separate_rotary` — Portable online attention shares each key/value load across query rows. Spare lanes partition value accumulation, with the partition count bounded by the head width, workgroup size, and available workgroup storage. Single-partition staging requires one output column per lane and includes the staging array in its storage budget.
+- `separate_rotary_sgmat` — Float32 subgroup-matrix online attention computes each score tile once, rescales output fragments through existing score scratch, and directly loads complete query tiles. Only a partial query tile is staged; key-tile width and fragment-bank capacity follow the reported workgroup-storage limit.
+- `separate_rotary_sgmat_tail` — For a prompt, a thin query tail uses the shared portable template after complete matrix tiles. History calls retain the full matrix computation. Tail size and value partitions follow workgroup capacity. Heads narrower than its query tile use the full matrix path, since their shorter value contraction does not repay a separate portable dispatch.
+- `packed` — Portable online attention shares each key/value load across query rows. Spare lanes partition value accumulation, with the partition count bounded by the head width, workgroup size, and available workgroup storage. Single-partition staging requires one output column per lane and includes the staging array in its storage budget.
+- `packed_sgmat` — Float32 subgroup-matrix online attention computes each score tile once, rescales output fragments through existing score scratch, and directly loads complete query tiles. Only a partial query tile is staged; key-tile width and fragment-bank capacity follow the reported workgroup-storage limit.
+- `packed_sgmat_tail` — For a prompt, a thin query tail uses the shared portable template after complete matrix tiles. History calls retain the full matrix computation. Tail size and value partitions follow workgroup capacity. Heads narrower than its query tile use the full matrix path, since their shorter value contraction does not repay a separate portable dispatch.
+- `packed_rotary` — Portable online attention shares each key/value load across query rows. Spare lanes partition value accumulation, with the partition count bounded by the head width, workgroup size, and available workgroup storage. Single-partition staging requires one output column per lane and includes the staging array in its storage budget.
+- `packed_rotary_sgmat` — Float32 subgroup-matrix online attention computes each score tile once, rescales output fragments through existing score scratch, and directly loads complete query tiles. Only a partial query tile is staged; key-tile width and fragment-bank capacity follow the reported workgroup-storage limit.
+- `packed_rotary_sgmat_tail` — For a prompt, a thin query tail uses the shared portable template after complete matrix tiles. History calls retain the full matrix computation. Tail size and value partitions follow workgroup capacity. Heads narrower than its query tile use the full matrix path, since their shorter value contraction does not repay a separate portable dispatch.
 ## Device requirements
 Some implementation variants require `subgroup-matrix` and `subgroups`. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.
 - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
+- [`bench.json`](build/webgpu/bench.json) — benchmark cases
 - [`sparse-attention-sgmat.wgsl.jinja`](build/webgpu/sparse-attention-sgmat.wgsl.jinja)
 - [`sparse-attention.wgsl.jinja`](build/webgpu/sparse-attention.wgsl.jinja)
 - [`sparse-kv-append.wgsl.jinja`](build/webgpu/sparse-kv-append.wgsl.jinja)
 ## Use with `@huggingface/kernels`
 ```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.3
 ```
 Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.

build/webgpu/bench.json CHANGED Viewed

@@ -1,7 +1,8 @@
 {
   "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": [
@@ -922,6 +923,1281 @@
         "keyTotalSequenceLengthsT": { "shape": [1], "dtype": "int32", "dist": "constant", "value": 256 }
       },
       "outputs": { "outputT": { "shape": [1, 128, 768], "dtype": "float32" } }
     }
   ]
 }

 {
   "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],
+    "boundary-value-partition-float32-d520-s1_input_blockColIndicesT": [0, 0, 1, 1, 2, 1, 2, 3, -1, 0, 0, 1, 0, 1, 2, 0, 2, 3]
   },
   "tunableSpace": { "WORKGROUP_SIZE": [32, 64, 128, 256], "APPEND_WORKGROUP_SIZE": [64, 128, 256] },
   "cases": [
         "keyTotalSequenceLengthsT": { "shape": [1], "dtype": "int32", "dist": "constant", "value": 256 }
       },
       "outputs": { "outputT": { "shape": [1, 128, 768], "dtype": "float32" } }
+    },
+    {
+      "name": "boundary-value-partition-float32-d520-s1",
+      "attrs": { "num_heads": 4, "kv_num_heads": 2, "sparse_block_size": 16 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 1, 2080],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [2, 1, 1040],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [2, 1, 1040],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 2, 64, 520],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.017, "scale": 0.25 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 2, 64, 520],
+          "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/boundary-value-partition-float32-d520-s1_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, 2080] } },
+      "preset": "model"
+    },
+    {
+      "name": "boundary-value-partition-float32-d520-s5",
+      "attrs": { "num_heads": 4, "kv_num_heads": 2, "sparse_block_size": 16 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 5, 2080],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [2, 5, 1040],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [2, 5, 1040],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 2, 64, 520],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.017, "scale": 0.25 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 2, 64, 520],
+          "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/boundary-value-partition-float32-d520-s1_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, 5, 2080] } },
+      "preset": "model"
+    },
+    {
+      "name": "boundary-value-owner-float32-d144-wg32-s1",
+      "attrs": { "num_heads": 4, "kv_num_heads": 2, "sparse_block_size": 16 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 1, 576],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [2, 1, 288],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [2, 1, 288],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 2, 64, 144],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.017, "scale": 0.25 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 2, 64, 144],
+          "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/boundary-value-partition-float32-d520-s1_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, 576] } },
+      "tunables": { "WORKGROUP_SIZE": 32 },
+      "preset": "model"
+    },
+    {
+      "name": "boundary-value-owner-float32-d144-wg32-s5",
+      "attrs": { "num_heads": 4, "kv_num_heads": 2, "sparse_block_size": 16 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 5, 576],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [2, 5, 288],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [2, 5, 288],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 2, 64, 144],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.017, "scale": 0.25 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 2, 64, 144],
+          "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/boundary-value-partition-float32-d520-s1_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, 5, 576] } },
+      "tunables": { "WORKGROUP_SIZE": 32 },
+      "preset": "model"
+    },
+    {
+      "name": "sparse-prompt-tail-b1-s127-h32kv8-d128-blk64",
+      "preset": "model",
+      "vars": {
+        "dtype": "float32",
+        "batch": 1,
+        "seq": 127,
+        "heads": 32,
+        "kvHeads": 8,
+        "headDim": 128,
+        "qkPairs": 260096,
+        "attendedKeys": 1024,
+        "dtypeBytes": 4
+      },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": { "shape": [1, 127, 4096], "dtype": "float32", "dist": "normal", "seed": 9300, "scale": 1 },
+        "keyT": { "shape": [1, 127, 1024], "dtype": "float32", "dist": "normal", "seed": 9301, "scale": 1 },
+        "valueT": { "shape": [1, 127, 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": [127] } },
+        "keyTotalSequenceLengthsT": { "shape": [1], "dtype": "int32", "dist": "constant", "value": 127 }
+      },
+      "outputs": { "outputT": { "shape": [1, 127, 4096], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "4 * args.qkPairs * args.headDim" }] }
+    },
+    {
+      "name": "sparse-prompt-tail-b1-s129-h32kv8-d128-blk64",
+      "preset": "model",
+      "vars": {
+        "dtype": "float32",
+        "batch": 1,
+        "seq": 129,
+        "heads": 32,
+        "kvHeads": 8,
+        "headDim": 128,
+        "qkPairs": 268320,
+        "attendedKeys": 1024,
+        "dtypeBytes": 4
+      },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": { "shape": [1, 129, 4096], "dtype": "float32", "dist": "normal", "seed": 9300, "scale": 1 },
+        "keyT": { "shape": [1, 129, 1024], "dtype": "float32", "dist": "normal", "seed": 9301, "scale": 1 },
+        "valueT": { "shape": [1, 129, 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": [129] } },
+        "keyTotalSequenceLengthsT": { "shape": [1], "dtype": "int32", "dist": "constant", "value": 129 }
+      },
+      "outputs": { "outputT": { "shape": [1, 129, 4096], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "4 * args.qkPairs * args.headDim" }] }
+    },
+    {
+      "name": "sparse-prompt-tail-b1-s511-h32kv8-d128-blk64",
+      "preset": "model",
+      "vars": {
+        "dtype": "float32",
+        "batch": 1,
+        "seq": 511,
+        "heads": 32,
+        "kvHeads": 8,
+        "headDim": 128,
+        "qkPairs": 4186112,
+        "attendedKeys": 1024,
+        "dtypeBytes": 4
+      },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": { "shape": [1, 511, 4096], "dtype": "float32", "dist": "normal", "seed": 9300, "scale": 1 },
+        "keyT": { "shape": [1, 511, 1024], "dtype": "float32", "dist": "normal", "seed": 9301, "scale": 1 },
+        "valueT": { "shape": [1, 511, 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": [511] } },
+        "keyTotalSequenceLengthsT": { "shape": [1], "dtype": "int32", "dist": "constant", "value": 511 }
+      },
+      "outputs": { "outputT": { "shape": [1, 511, 4096], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "4 * args.qkPairs * args.headDim" }] }
+    },
+    {
+      "name": "sparse-prompt-tail-b1-s513-h32kv8-d128-blk64",
+      "preset": "model",
+      "vars": {
+        "dtype": "float32",
+        "batch": 1,
+        "seq": 513,
+        "heads": 32,
+        "kvHeads": 8,
+        "headDim": 128,
+        "qkPairs": 4217376,
+        "attendedKeys": 1024,
+        "dtypeBytes": 4
+      },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": { "shape": [1, 513, 4096], "dtype": "float32", "dist": "normal", "seed": 9300, "scale": 1 },
+        "keyT": { "shape": [1, 513, 1024], "dtype": "float32", "dist": "normal", "seed": 9301, "scale": 1 },
+        "valueT": { "shape": [1, 513, 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": [513] } },
+        "keyTotalSequenceLengthsT": { "shape": [1], "dtype": "int32", "dist": "constant", "value": 513 }
+      },
+      "outputs": { "outputT": { "shape": [1, 513, 4096], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "4 * args.qkPairs * args.headDim" }] }
+    },
+    {
+      "name": "sparse-prompt-tail-b1-s1023-h32kv8-d128-blk64",
+      "preset": "model",
+      "vars": {
+        "dtype": "float32",
+        "batch": 1,
+        "seq": 1023,
+        "heads": 32,
+        "kvHeads": 8,
+        "headDim": 128,
+        "qkPairs": 13234176,
+        "attendedKeys": 1024,
+        "dtypeBytes": 4
+      },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": { "shape": [1, 1023, 4096], "dtype": "float32", "dist": "normal", "seed": 9300, "scale": 1 },
+        "keyT": { "shape": [1, 1023, 1024], "dtype": "float32", "dist": "normal", "seed": 9301, "scale": 1 },
+        "valueT": { "shape": [1, 1023, 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": [1023] } },
+        "keyTotalSequenceLengthsT": { "shape": [1], "dtype": "int32", "dist": "constant", "value": 1023 }
+      },
+      "outputs": { "outputT": { "shape": [1, 1023, 4096], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "4 * args.qkPairs * args.headDim" }] }
+    },
+    {
+      "name": "hybrid-neighbor-s65-h32kv8-d128",
+      "preset": "model",
+      "vars": {
+        "dtype": "float32",
+        "batch": 1,
+        "seq": 65,
+        "heads": 32,
+        "kvHeads": 8,
+        "headDim": 128,
+        "qkPairs": 68640,
+        "attendedKeys": 65,
+        "dtypeBytes": 4
+      },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": { "shape": [1, 65, 4096], "dtype": "float32", "dist": "normal", "seed": 9300, "scale": 1 },
+        "keyT": { "shape": [1, 65, 1024], "dtype": "float32", "dist": "normal", "seed": 9301, "scale": 1 },
+        "valueT": { "shape": [1, 65, 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": [65] } },
+        "keyTotalSequenceLengthsT": { "shape": [1], "dtype": "int32", "dist": "constant", "value": 65 }
+      },
+      "outputs": { "outputT": { "shape": [1, 65, 4096], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "4 * args.qkPairs * args.headDim" }] }
+    },
+    {
+      "name": "hybrid-neighbor-s66-h32kv8-d128",
+      "preset": "model",
+      "vars": {
+        "dtype": "float32",
+        "batch": 1,
+        "seq": 66,
+        "heads": 32,
+        "kvHeads": 8,
+        "headDim": 128,
+        "qkPairs": 70752,
+        "attendedKeys": 66,
+        "dtypeBytes": 4
+      },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": { "shape": [1, 66, 4096], "dtype": "float32", "dist": "normal", "seed": 9300, "scale": 1 },
+        "keyT": { "shape": [1, 66, 1024], "dtype": "float32", "dist": "normal", "seed": 9301, "scale": 1 },
+        "valueT": { "shape": [1, 66, 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": [66] } },
+        "keyTotalSequenceLengthsT": { "shape": [1], "dtype": "int32", "dist": "constant", "value": 66 }
+      },
+      "outputs": { "outputT": { "shape": [1, 66, 4096], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "4 * args.qkPairs * args.headDim" }] }
+    },
+    {
+      "name": "hybrid-neighbor-s67-h32kv8-d128",
+      "preset": "model",
+      "vars": {
+        "dtype": "float32",
+        "batch": 1,
+        "seq": 67,
+        "heads": 32,
+        "kvHeads": 8,
+        "headDim": 128,
+        "qkPairs": 72896,
+        "attendedKeys": 67,
+        "dtypeBytes": 4
+      },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": { "shape": [1, 67, 4096], "dtype": "float32", "dist": "normal", "seed": 9300, "scale": 1 },
+        "keyT": { "shape": [1, 67, 1024], "dtype": "float32", "dist": "normal", "seed": 9301, "scale": 1 },
+        "valueT": { "shape": [1, 67, 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": [67] } },
+        "keyTotalSequenceLengthsT": { "shape": [1], "dtype": "int32", "dist": "constant", "value": 67 }
+      },
+      "outputs": { "outputT": { "shape": [1, 67, 4096], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "4 * args.qkPairs * args.headDim" }] }
+    },
+    {
+      "name": "hybrid-neighbor-s68-h32kv8-d128",
+      "preset": "model",
+      "vars": {
+        "dtype": "float32",
+        "batch": 1,
+        "seq": 68,
+        "heads": 32,
+        "kvHeads": 8,
+        "headDim": 128,
+        "qkPairs": 75072,
+        "attendedKeys": 68,
+        "dtypeBytes": 4
+      },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": { "shape": [1, 68, 4096], "dtype": "float32", "dist": "normal", "seed": 9300, "scale": 1 },
+        "keyT": { "shape": [1, 68, 1024], "dtype": "float32", "dist": "normal", "seed": 9301, "scale": 1 },
+        "valueT": { "shape": [1, 68, 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": [68] } },
+        "keyTotalSequenceLengthsT": { "shape": [1], "dtype": "int32", "dist": "constant", "value": 68 }
+      },
+      "outputs": { "outputT": { "shape": [1, 68, 4096], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "4 * args.qkPairs * args.headDim" }] }
+    },
+    {
+      "name": "hybrid-neighbor-s130-h32kv8-d128",
+      "preset": "model",
+      "vars": {
+        "dtype": "float32",
+        "batch": 1,
+        "seq": 130,
+        "heads": 32,
+        "kvHeads": 8,
+        "headDim": 128,
+        "qkPairs": 272480,
+        "attendedKeys": 130,
+        "dtypeBytes": 4
+      },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": { "shape": [1, 130, 4096], "dtype": "float32", "dist": "normal", "seed": 9300, "scale": 1 },
+        "keyT": { "shape": [1, 130, 1024], "dtype": "float32", "dist": "normal", "seed": 9301, "scale": 1 },
+        "valueT": { "shape": [1, 130, 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": [130] } },
+        "keyTotalSequenceLengthsT": { "shape": [1], "dtype": "int32", "dist": "constant", "value": 130 }
+      },
+      "outputs": { "outputT": { "shape": [1, 130, 4096], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "4 * args.qkPairs * args.headDim" }] }
+    },
+    {
+      "name": "hybrid-neighbor-s131-h32kv8-d128",
+      "preset": "model",
+      "vars": {
+        "dtype": "float32",
+        "batch": 1,
+        "seq": 131,
+        "heads": 32,
+        "kvHeads": 8,
+        "headDim": 128,
+        "qkPairs": 276672,
+        "attendedKeys": 131,
+        "dtypeBytes": 4
+      },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": { "shape": [1, 131, 4096], "dtype": "float32", "dist": "normal", "seed": 9300, "scale": 1 },
+        "keyT": { "shape": [1, 131, 1024], "dtype": "float32", "dist": "normal", "seed": 9301, "scale": 1 },
+        "valueT": { "shape": [1, 131, 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": [131] } },
+        "keyTotalSequenceLengthsT": { "shape": [1], "dtype": "int32", "dist": "constant", "value": 131 }
+      },
+      "outputs": { "outputT": { "shape": [1, 131, 4096], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "4 * args.qkPairs * args.headDim" }] }
+    },
+    {
+      "name": "hybrid-neighbor-s132-h32kv8-d128",
+      "preset": "model",
+      "vars": {
+        "dtype": "float32",
+        "batch": 1,
+        "seq": 132,
+        "heads": 32,
+        "kvHeads": 8,
+        "headDim": 128,
+        "qkPairs": 280896,
+        "attendedKeys": 132,
+        "dtypeBytes": 4
+      },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": { "shape": [1, 132, 4096], "dtype": "float32", "dist": "normal", "seed": 9300, "scale": 1 },
+        "keyT": { "shape": [1, 132, 1024], "dtype": "float32", "dist": "normal", "seed": 9301, "scale": 1 },
+        "valueT": { "shape": [1, 132, 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": [132] } },
+        "keyTotalSequenceLengthsT": { "shape": [1], "dtype": "int32", "dist": "constant", "value": 132 }
+      },
+      "outputs": { "outputT": { "shape": [1, 132, 4096], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "4 * args.qkPairs * args.headDim" }] }
+    },
+    {
+      "name": "hybrid-neighbor-s512-h32kv8-d128",
+      "preset": "model",
+      "vars": {
+        "dtype": "float32",
+        "batch": 1,
+        "seq": 512,
+        "heads": 32,
+        "kvHeads": 8,
+        "headDim": 128,
+        "qkPairs": 4202496,
+        "attendedKeys": 512,
+        "dtypeBytes": 4
+      },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": { "shape": [1, 512, 4096], "dtype": "float32", "dist": "normal", "seed": 9300, "scale": 1 },
+        "keyT": { "shape": [1, 512, 1024], "dtype": "float32", "dist": "normal", "seed": 9301, "scale": 1 },
+        "valueT": { "shape": [1, 512, 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": [512] } },
+        "keyTotalSequenceLengthsT": { "shape": [1], "dtype": "int32", "dist": "constant", "value": 512 }
+      },
+      "outputs": { "outputT": { "shape": [1, 512, 4096], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "4 * args.qkPairs * args.headDim" }] }
+    },
+    {
+      "name": "hybrid-neighbor-s514-h32kv8-d128",
+      "preset": "model",
+      "vars": {
+        "dtype": "float32",
+        "batch": 1,
+        "seq": 514,
+        "heads": 32,
+        "kvHeads": 8,
+        "headDim": 128,
+        "qkPairs": 4232288,
+        "attendedKeys": 514,
+        "dtypeBytes": 4
+      },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": { "shape": [1, 514, 4096], "dtype": "float32", "dist": "normal", "seed": 9300, "scale": 1 },
+        "keyT": { "shape": [1, 514, 1024], "dtype": "float32", "dist": "normal", "seed": 9301, "scale": 1 },
+        "valueT": { "shape": [1, 514, 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": [514] } },
+        "keyTotalSequenceLengthsT": { "shape": [1], "dtype": "int32", "dist": "constant", "value": 514 }
+      },
+      "outputs": { "outputT": { "shape": [1, 514, 4096], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "4 * args.qkPairs * args.headDim" }] }
+    },
+    {
+      "name": "hybrid-neighbor-s515-h32kv8-d128",
+      "preset": "model",
+      "vars": {
+        "dtype": "float32",
+        "batch": 1,
+        "seq": 515,
+        "heads": 32,
+        "kvHeads": 8,
+        "headDim": 128,
+        "qkPairs": 4247232,
+        "attendedKeys": 515,
+        "dtypeBytes": 4
+      },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": { "shape": [1, 515, 4096], "dtype": "float32", "dist": "normal", "seed": 9300, "scale": 1 },
+        "keyT": { "shape": [1, 515, 1024], "dtype": "float32", "dist": "normal", "seed": 9301, "scale": 1 },
+        "valueT": { "shape": [1, 515, 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": [515] } },
+        "keyTotalSequenceLengthsT": { "shape": [1], "dtype": "int32", "dist": "constant", "value": 515 }
+      },
+      "outputs": { "outputT": { "shape": [1, 515, 4096], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "4 * args.qkPairs * args.headDim" }] }
+    },
+    {
+      "name": "hybrid-neighbor-s516-h32kv8-d128",
+      "preset": "model",
+      "vars": {
+        "dtype": "float32",
+        "batch": 1,
+        "seq": 516,
+        "heads": 32,
+        "kvHeads": 8,
+        "headDim": 128,
+        "qkPairs": 4262208,
+        "attendedKeys": 516,
+        "dtypeBytes": 4
+      },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": { "shape": [1, 516, 4096], "dtype": "float32", "dist": "normal", "seed": 9300, "scale": 1 },
+        "keyT": { "shape": [1, 516, 1024], "dtype": "float32", "dist": "normal", "seed": 9301, "scale": 1 },
+        "valueT": { "shape": [1, 516, 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": [516] } },
+        "keyTotalSequenceLengthsT": { "shape": [1], "dtype": "int32", "dist": "constant", "value": 516 }
+      },
+      "outputs": { "outputT": { "shape": [1, 516, 4096], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "4 * args.qkPairs * args.headDim" }] }
+    },
+    {
+      "name": "hybrid-neighbor-s517-h32kv8-d128",
+      "preset": "model",
+      "vars": {
+        "dtype": "float32",
+        "batch": 1,
+        "seq": 517,
+        "heads": 32,
+        "kvHeads": 8,
+        "headDim": 128,
+        "qkPairs": 4277216,
+        "attendedKeys": 517,
+        "dtypeBytes": 4
+      },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": { "shape": [1, 517, 4096], "dtype": "float32", "dist": "normal", "seed": 9300, "scale": 1 },
+        "keyT": { "shape": [1, 517, 1024], "dtype": "float32", "dist": "normal", "seed": 9301, "scale": 1 },
+        "valueT": { "shape": [1, 517, 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": [517] } },
+        "keyTotalSequenceLengthsT": { "shape": [1], "dtype": "int32", "dist": "constant", "value": 517 }
+      },
+      "outputs": { "outputT": { "shape": [1, 517, 4096], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "4 * args.qkPairs * args.headDim" }] }
+    },
+    {
+      "name": "hybrid-neighbor-s577-h32kv8-d128",
+      "preset": "model",
+      "vars": {
+        "dtype": "float32",
+        "batch": 1,
+        "seq": 577,
+        "heads": 32,
+        "kvHeads": 8,
+        "headDim": 128,
+        "qkPairs": 5234720,
+        "attendedKeys": 577,
+        "dtypeBytes": 4
+      },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": { "shape": [1, 577, 4096], "dtype": "float32", "dist": "normal", "seed": 9300, "scale": 1 },
+        "keyT": { "shape": [1, 577, 1024], "dtype": "float32", "dist": "normal", "seed": 9301, "scale": 1 },
+        "valueT": { "shape": [1, 577, 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": [577] } },
+        "keyTotalSequenceLengthsT": { "shape": [1], "dtype": "int32", "dist": "constant", "value": 577 }
+      },
+      "outputs": { "outputT": { "shape": [1, 577, 4096], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "4 * args.qkPairs * args.headDim" }] }
+    },
+    {
+      "name": "isolate-packed-rotary-prompt-s65",
+      "preset": "model",
+      "vars": {
+        "dtype": "float32",
+        "batch": 1,
+        "seq": 65,
+        "heads": 32,
+        "kvHeads": 8,
+        "headDim": 128,
+        "qkPairs": 68640,
+        "attendedKeys": 65,
+        "dtypeBytes": 4
+      },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sparse_block_size": 64, "do_rotary": 1 },
+      "inputs": {
+        "queryT": { "shape": [1, 65, 6144], "dtype": "float32", "dist": "normal", "seed": 9330, "scale": 1 },
+        "pastKeyT": { "shape": [1, 8, 1024, 128], "dtype": "float32", "dist": "normal", "seed": 9333, "scale": 1 },
+        "pastValueT": { "shape": [1, 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": [65] } },
+        "keyTotalSequenceLengthsT": { "shape": [1], "dtype": "int32", "dist": "constant", "value": 65 },
+        "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": [1, 65, 4096], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "4 * args.qkPairs * args.headDim" }] }
+    },
+    {
+      "name": "isolate-separate-history-s65",
+      "preset": "model",
+      "vars": {
+        "dtype": "float32",
+        "batch": 1,
+        "seq": 65,
+        "heads": 32,
+        "kvHeads": 8,
+        "headDim": 128,
+        "qkPairs": 1264000,
+        "attendedKeys": 1020,
+        "dtypeBytes": 4
+      },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": { "shape": [1, 65, 4096], "dtype": "float32", "dist": "normal", "seed": 9300, "scale": 1 },
+        "keyT": { "shape": [1, 65, 1024], "dtype": "float32", "dist": "normal", "seed": 9301, "scale": 1 },
+        "valueT": { "shape": [1, 65, 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": [1020] } },
+        "keyTotalSequenceLengthsT": { "shape": [1], "dtype": "int32", "dist": "constant", "value": 1020 }
+      },
+      "outputs": { "outputT": { "shape": [1, 65, 4096], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "4 * args.qkPairs * args.headDim" }] }
+    },
+    {
+      "name": "isolate-separate-history-s129",
+      "preset": "model",
+      "vars": {
+        "dtype": "float32",
+        "batch": 1,
+        "seq": 129,
+        "heads": 32,
+        "kvHeads": 8,
+        "headDim": 128,
+        "qkPairs": 2474880,
+        "attendedKeys": 1020,
+        "dtypeBytes": 4
+      },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": { "shape": [1, 129, 4096], "dtype": "float32", "dist": "normal", "seed": 9300, "scale": 1 },
+        "keyT": { "shape": [1, 129, 1024], "dtype": "float32", "dist": "normal", "seed": 9301, "scale": 1 },
+        "valueT": { "shape": [1, 129, 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": [1020] } },
+        "keyTotalSequenceLengthsT": { "shape": [1], "dtype": "int32", "dist": "constant", "value": 1020 }
+      },
+      "outputs": { "outputT": { "shape": [1, 129, 4096], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "4 * args.qkPairs * args.headDim" }] }
+    },
+    {
+      "name": "isolate-separate-history-s513",
+      "preset": "model",
+      "vars": {
+        "dtype": "float32",
+        "batch": 1,
+        "seq": 513,
+        "heads": 32,
+        "kvHeads": 8,
+        "headDim": 128,
+        "qkPairs": 9052032,
+        "attendedKeys": 1020,
+        "dtypeBytes": 4
+      },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": { "shape": [1, 513, 4096], "dtype": "float32", "dist": "normal", "seed": 9300, "scale": 1 },
+        "keyT": { "shape": [1, 513, 1024], "dtype": "float32", "dist": "normal", "seed": 9301, "scale": 1 },
+        "valueT": { "shape": [1, 513, 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": [1020] } },
+        "keyTotalSequenceLengthsT": { "shape": [1], "dtype": "int32", "dist": "constant", "value": 1020 }
+      },
+      "outputs": { "outputT": { "shape": [1, 513, 4096], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "4 * args.qkPairs * args.headDim" }] }
+    },
+    {
+      "name": "hybrid-packed-rotary-s65-past955",
+      "preset": "model",
+      "vars": {
+        "dtype": "float32",
+        "batch": 1,
+        "seq": 65,
+        "heads": 32,
+        "kvHeads": 8,
+        "headDim": 128,
+        "qkPairs": 1264000,
+        "attendedKeys": 1020,
+        "dtypeBytes": 4
+      },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sparse_block_size": 64, "do_rotary": 1 },
+      "inputs": {
+        "queryT": { "shape": [1, 65, 6144], "dtype": "float32", "dist": "normal", "seed": 9330, "scale": 1 },
+        "pastKeyT": { "shape": [1, 8, 1024, 128], "dtype": "float32", "dist": "normal", "seed": 9333, "scale": 1 },
+        "pastValueT": { "shape": [1, 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": [1020] } },
+        "keyTotalSequenceLengthsT": { "shape": [1], "dtype": "int32", "dist": "constant", "value": 1020 },
+        "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": [1, 65, 4096], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "4 * args.qkPairs * args.headDim" }] }
+    },
+    {
+      "name": "hybrid-d96-s131",
+      "preset": "model",
+      "vars": {
+        "dtype": "float32",
+        "batch": 1,
+        "seq": 131,
+        "heads": 32,
+        "kvHeads": 8,
+        "headDim": 96,
+        "qkPairs": 276672,
+        "attendedKeys": 131,
+        "dtypeBytes": 4
+      },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": { "shape": [1, 131, 3072], "dtype": "float32", "dist": "normal", "seed": 9300, "scale": 1 },
+        "keyT": { "shape": [1, 131, 768], "dtype": "float32", "dist": "normal", "seed": 9301, "scale": 1 },
+        "valueT": { "shape": [1, 131, 768], "dtype": "float32", "dist": "normal", "seed": 9302, "scale": 1 },
+        "pastKeyT": { "shape": [1, 8, 1024, 96], "dtype": "float32", "dist": "normal", "seed": 9303, "scale": 1 },
+        "pastValueT": { "shape": [1, 8, 1024, 96], "dtype": "float32", "dist": "normal", "seed": 9304, "scale": 1 },
+        "blockRowIndicesT": {
+          "shape": [4, 17],
+          "dtype": "int32",
+          "data": {
+            "kind": "values",
+            "values": { "$ref": "#/fixtureArrays/sparse-prompt-b1-s1024-h32kv8-d128-blk64_input_blockRowIndicesT" }
+          }
+        },
+        "blockColIndicesT": {
+          "shape": [4, 112],
+          "dtype": "int32",
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/block_column_indices_t_pattern" } }
+        },
+        "totalSequenceLengthT": { "shape": [1], "dtype": "int32", "data": { "kind": "values", "values": [131] } },
+        "keyTotalSequenceLengthsT": { "shape": [1], "dtype": "int32", "dist": "constant", "value": 131 }
+      },
+      "outputs": { "outputT": { "shape": [1, 131, 3072], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "4 * args.qkPairs * args.headDim" }] }
+    },
+    {
+      "name": "hybrid-d32-s131",
+      "preset": "model",
+      "vars": {
+        "dtype": "float32",
+        "batch": 1,
+        "seq": 131,
+        "heads": 32,
+        "kvHeads": 8,
+        "headDim": 32,
+        "qkPairs": 276672,
+        "attendedKeys": 131,
+        "dtypeBytes": 4
+      },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": { "shape": [1, 131, 1024], "dtype": "float32", "dist": "normal", "seed": 9300, "scale": 1 },
+        "keyT": { "shape": [1, 131, 256], "dtype": "float32", "dist": "normal", "seed": 9301, "scale": 1 },
+        "valueT": { "shape": [1, 131, 256], "dtype": "float32", "dist": "normal", "seed": 9302, "scale": 1 },
+        "pastKeyT": { "shape": [1, 8, 1024, 32], "dtype": "float32", "dist": "normal", "seed": 9303, "scale": 1 },
+        "pastValueT": { "shape": [1, 8, 1024, 32], "dtype": "float32", "dist": "normal", "seed": 9304, "scale": 1 },
+        "blockRowIndicesT": {
+          "shape": [4, 17],
+          "dtype": "int32",
+          "data": {
+            "kind": "values",
+            "values": { "$ref": "#/fixtureArrays/sparse-prompt-b1-s1024-h32kv8-d128-blk64_input_blockRowIndicesT" }
+          }
+        },
+        "blockColIndicesT": {
+          "shape": [4, 112],
+          "dtype": "int32",
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/block_column_indices_t_pattern" } }
+        },
+        "totalSequenceLengthT": { "shape": [1], "dtype": "int32", "data": { "kind": "values", "values": [131] } },
+        "keyTotalSequenceLengthsT": { "shape": [1], "dtype": "int32", "dist": "constant", "value": 131 }
+      },
+      "outputs": { "outputT": { "shape": [1, 131, 1024], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "4 * args.qkPairs * args.headDim" }] }
+    },
+    {
+      "name": "hybrid-d64-s131",
+      "preset": "model",
+      "vars": {
+        "dtype": "float32",
+        "batch": 1,
+        "seq": 131,
+        "heads": 32,
+        "kvHeads": 8,
+        "headDim": 64,
+        "qkPairs": 276672,
+        "attendedKeys": 131,
+        "dtypeBytes": 4
+      },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": { "shape": [1, 131, 2048], "dtype": "float32", "dist": "normal", "seed": 9300, "scale": 1 },
+        "keyT": { "shape": [1, 131, 512], "dtype": "float32", "dist": "normal", "seed": 9301, "scale": 1 },
+        "valueT": { "shape": [1, 131, 512], "dtype": "float32", "dist": "normal", "seed": 9302, "scale": 1 },
+        "pastKeyT": { "shape": [1, 8, 1024, 64], "dtype": "float32", "dist": "normal", "seed": 9303, "scale": 1 },
+        "pastValueT": { "shape": [1, 8, 1024, 64], "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": [131] } },
+        "keyTotalSequenceLengthsT": { "shape": [1], "dtype": "int32", "dist": "constant", "value": 131 }
+      },
+      "outputs": { "outputT": { "shape": [1, 131, 2048], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "4 * args.qkPairs * args.headDim" }] }
+    },
+    {
+      "name": "hybrid-d32-s131-narrow",
+      "preset": "model",
+      "vars": {
+        "dtype": "float32",
+        "batch": 1,
+        "seq": 131,
+        "heads": 32,
+        "kvHeads": 8,
+        "headDim": 32,
+        "qkPairs": 276672,
+        "attendedKeys": 131,
+        "dtypeBytes": 4
+      },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": { "shape": [1, 131, 1024], "dtype": "float32", "dist": "normal", "seed": 9300, "scale": 1 },
+        "keyT": { "shape": [1, 131, 256], "dtype": "float32", "dist": "normal", "seed": 9301, "scale": 1 },
+        "valueT": { "shape": [1, 131, 256], "dtype": "float32", "dist": "normal", "seed": 9302, "scale": 1 },
+        "pastKeyT": { "shape": [1, 8, 1024, 32], "dtype": "float32", "dist": "normal", "seed": 9303, "scale": 1 },
+        "pastValueT": { "shape": [1, 8, 1024, 32], "dtype": "float32", "dist": "normal", "seed": 9304, "scale": 1 },
+        "blockRowIndicesT": {
+          "shape": [4, 17],
+          "dtype": "int32",
+          "data": {
+            "kind": "values",
+            "values": { "$ref": "#/fixtureArrays/sparse-prompt-b1-s1024-h32kv8-d128-blk64_input_blockRowIndicesT" }
+          }
+        },
+        "blockColIndicesT": {
+          "shape": [4, 112],
+          "dtype": "int32",
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/block_column_indices_t_pattern" } }
+        },
+        "totalSequenceLengthT": { "shape": [1], "dtype": "int32", "data": { "kind": "values", "values": [131] } },
+        "keyTotalSequenceLengthsT": { "shape": [1], "dtype": "int32", "dist": "constant", "value": 131 }
+      },
+      "outputs": { "outputT": { "shape": [1, 131, 1024], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "4 * args.qkPairs * args.headDim" }] },
+      "tunables": { "NARROW_MIN_WORKGROUPS": 1 }
     }
   ]
 }

build/webgpu/manifest.json CHANGED Viewed

@@ -54,9 +54,10 @@
     "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",
@@ -82,93 +83,98 @@
     "rotaryPairOk": "not doRotary or (present.cosCacheT and present.sinCacheT and ranks.cosCacheT == 2 and ranks.sinCacheT == 2 and rotaryHalf % 8 == 0 and rotaryDim <= headSize and sameShape(shapes.sinCacheT, shapes.cosCacheT) and tensorDtypes.cosCacheT == tensorDtypes.queryT and tensorDtypes.sinCacheT == tensorDtypes.queryT)",
     "blockIndexShapeOk": "ranks.blockRowIndicesT == 2 and ranks.blockColIndicesT == 2 and dim(shapes.blockColIndicesT, 0) == numLayout and maxBlocks >= 1 and maxNnz >= 0 and maxNnz <= maxBlocks * maxBlocks and tensorDtypes.blockRowIndicesT == \"int32\" and tensorDtypes.blockColIndicesT == \"int32\"",
     "scheduleShapeOk": "(ranks.totalSequenceLengthT == 0 or ranks.totalSequenceLengthT == 1) and numel(shapes.totalSequenceLengthT) == 1 and ranks.keyTotalSequenceLengthsT == 1 and dim(shapes.keyTotalSequenceLengthsT, 0) == batchSize and tensorDtypes.totalSequenceLengthT == \"int32\" and tensorDtypes.keyTotalSequenceLengthsT == \"int32\"",
-    "geometryOk": "tunables.WORKGROUP_SIZE >= 1 and floor(tunables.WORKGROUP_SIZE) == tunables.WORKGROUP_SIZE and pow2ceil(tunables.WORKGROUP_SIZE) == tunables.WORKGROUP_SIZE and tunables.WORKGROUP_SIZE <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.WORKGROUP_SIZE <= device.limits.maxComputeWorkgroupSizeX and tunables.APPEND_WORKGROUP_SIZE >= 1 and floor(tunables.APPEND_WORKGROUP_SIZE) == tunables.APPEND_WORKGROUP_SIZE and tunables.APPEND_WORKGROUP_SIZE <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.APPEND_WORKGROUP_SIZE <= device.limits.maxComputeWorkgroupSizeX and sparseQueryTiles <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and batchSize * numHeads <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(ceilDiv(qRotaryElements, tunables.APPEND_WORKGROUP_SIZE), min(device.limits.maxComputeWorkgroupsPerDimension, 65535)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and (2 * sparseQueryTile * headSize + (3 * sparseQueryTile + 1) * sparseAttnWorkgroup) * 4 <= device.limits.maxComputeWorkgroupStorageSize and sparseAttnWorkgroup <= device.limits.maxComputeInvocationsPerWorkgroup and sparseAttnWorkgroup <= device.limits.maxComputeWorkgroupSizeX",
     "contract": "ranks.queryT == 3 and ranks.outputT == 3 and (tensorDtypes.queryT == \"float32\" or tensorDtypes.queryT == \"float16\") and f16Ok(dtypes.T) and tensorDtypes.pastKeyT == tensorDtypes.queryT and tensorDtypes.pastValueT == tensorDtypes.queryT and tensorDtypes.outputT == tensorDtypes.queryT and numHeads >= 1 and kvNumHeads >= 1 and numHeads % kvNumHeads == 0 and headSize >= 8 and headSize % 8 == 0 and (not doRotary or headSize % 16 == 0) and numLayout >= 1 and numHeads % numLayout == 0 and (sparseBlockSize == 16 or sparseBlockSize == 32 or sparseBlockSize == 64 or sparseBlockSize == 128) and cacheShapeOk and queryShapeOk and kvShapeOk and kvPairOk and rotaryPairOk and blockIndexShapeOk and scheduleShapeOk and dim(shapes.outputT, 0) == batchSize and dim(shapes.outputT, 1) == seqLen and dim(shapes.outputT, 2) == qHidden",
     "packedContract": "contract and packedQkv and not useRotary",
     "packedRotaryContract": "contract and packedQkv and useRotary",
     "separateContract": "contract and not packedQkv and not useRotary",
     "separateRotaryContract": "contract and not packedQkv and useRotary",
-    "sparseVStageWorthIt": "sparseQueryTiles * batchSize * numHeads <= tunables.V_STAGE_MAX_WORKGROUPS",
-    "sgmatQueryTiles": "ceilDiv(seqLen, 64)",
-    "sgmatDirectQuery": "seqLen % 64 == 0",
-    "sparseSgmatTileN": "64 if (64 * 32 + 64 * 64 + 64 * 2 + 128 * 2) * 4 <= device.limits.maxComputeWorkgroupStorageSize else 32",
     "sparseSgmatTileK": "sparseSgmatTileN / 2",
-    "sparseSgmatLdsBytes": "(64 * sparseSgmatTileK + 64 * sparseSgmatTileN + 64 * 2 + 128 * 2) * 4",
-    "sparseSgmatGeometryOk": "256 <= device.limits.maxComputeInvocationsPerWorkgroup and 256 <= device.limits.maxComputeWorkgroupSizeX and sgmatQueryTiles <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and batchSize * numHeads <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and sparseSgmatLdsBytes <= device.limits.maxComputeWorkgroupStorageSize",
     "sparseSgmatOk": "tensorDtypes.queryT == \"float32\" and seqLen >= 64 and sparseBlockSize % 64 == 0 and headSize % 32 == 0 and headSize <= 128 and maxCacheSeq % 64 == 0 and device.features.has(\"subgroups\") and wave32Effective and device.features.has(\"chromium-experimental-subgroup-matrix\") and sparseSgmatGeometryOk",
     "scalar": "dtypes.T",
     "cacheVec": "\"vec4<f16>\" if dtypes.T == \"f16\" else \"vec4<f32>\"",
     "attnWorkgroup": "sparseAttnWorkgroup",
     "usesRotary": "useRotary",
-    "appendWorkgroupSize": "tunables.APPEND_WORKGROUP_SIZE"
   },
   "when": ["geometryOk"],
   "bindings": {
-    "new_key": { "arg": "keyT", "buffer": "read-only-storage", "elementType": "$scalar" },
-    "new_value": { "arg": "valueT", "buffer": "read-only-storage", "elementType": "$scalar" },
-    "present_key": { "arg": "pastKeyT", "buffer": "storage", "elementType": "$scalar" },
-    "present_value": { "arg": "pastValueT", "buffer": "storage", "elementType": "$scalar" },
-    "key_total_sequence_lengths": {
-      "arg": "keyTotalSequenceLengthsT",
-      "buffer": "read-only-storage",
-      "elementType": "i32"
-    },
-    "total_sequence_length": { "arg": "totalSequenceLengthT", "buffer": "read-only-storage", "elementType": "i32" },
     "params": {
-      "buffer": "uniform",
       "struct": [
         { "name": "batchSize", "type": "u32", "value": "batchSize" },
         { "name": "seqLen", "type": "u32", "value": "seqLen" }
       ]
     },
-    "cos_cache": { "arg": "cosCacheT", "buffer": "read-only-storage", "elementType": "$scalar" },
-    "sin_cache": { "arg": "sinCacheT", "buffer": "read-only-storage", "elementType": "$scalar" },
-    "packed_qkv": { "arg": "queryT", "buffer": "read-only-storage", "elementType": "$scalar" },
-    "query": { "arg": "queryT", "buffer": "read-only-storage", "elementType": "$scalar" },
-    "present_key_2": {
       "arg": "pastKeyT",
       "name": "present_key",
       "buffer": "read-only-storage",
       "elementType": "$cacheVec"
     },
-    "present_value_2": {
       "arg": "pastValueT",
       "name": "present_value",
       "buffer": "read-only-storage",
       "elementType": "$cacheVec"
     },
-    "block_row_indices": { "arg": "blockRowIndicesT", "buffer": "read-only-storage", "elementType": "i32" },
-    "block_col_indices": { "arg": "blockColIndicesT", "buffer": "read-only-storage", "elementType": "i32" },
-    "output": { "arg": "outputT", "buffer": "storage", "elementType": "$scalar" },
-    "params_2": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
         { "name": "seqLen", "type": "u32", "value": "seqLen" },
         { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" }
       ]
     },
     "q_rotary": { "scratch": "QRotary", "buffer": "read-only-storage", "elementType": "f32" },
-    "present_key_3": {
       "arg": "pastKeyT",
       "name": "present_key",
       "buffer": "read-only-storage",
       "elementType": "$scalar"
     },
-    "present_value_3": {
       "arg": "pastValueT",
       "name": "present_value",
       "buffer": "read-only-storage",
       "elementType": "$scalar"
     },
-    "q_rotary_2": { "scratch": "QRotary", "name": "q_rotary", "buffer": "storage", "elementType": "f32" }
   },
   "variants": [
     {
       "id": "separate",
       "priority": 0,
       "when": ["separateContract"],
-      "derive": { "vStageWorthIt": "sparseVStageWorthIt" },
       "passes": [
         {
           "id": "append",
@@ -186,8 +192,8 @@
           "id": "attention",
           "name": "SparseAttention.Attention",
           "shader": "sparse-attention.wgsl.jinja",
-          "derive": { "qTile": "sparseQueryTile" },
-          "bindings": ["query", "present_key_2", "present_value_2", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_2"],
           "dispatch": { "x": "sparseQueryTiles", "y": "batchSize * numHeads" }
         }
       ]
@@ -217,18 +223,65 @@
           "id": "attention",
           "name": "SparseAttention.AttentionSgmat",
           "shader": "sparse-attention-sgmat.wgsl.jinja",
-          "bindings": ["query", "present_key_3", "present_value_3", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_2"],
           "dispatch": { "x": "sgmatQueryTiles", "y": "batchSize * numHeads" },
-          "subgroupCollectivesWidth": 32
         }
       ],
       "demoteWhen": ["sgmatQueryTiles * batchSize * numHeads * sparseSgmatTileN < tunables.MATRIX_MIN_WORKGROUPS * 64"]
     },
     {
       "id": "separate_rotary",
       "priority": 10,
       "when": ["separateRotaryContract"],
-      "derive": { "vStageWorthIt": "sparseVStageWorthIt" },
       "intermediates": [{ "id": "QRotary", "dtype": "float32", "shape": "[qRotaryElements]" }],
       "passes": [
         {
@@ -247,7 +300,7 @@
           "id": "qrotary",
           "name": "SparseAttention.QueryRotary",
           "shader": "sparse-q-rotary.wgsl.jinja",
-          "bindings": ["query", "cos_cache", "sin_cache", "q_rotary_2", "key_total_sequence_lengths", "total_sequence_length", "params"],
           "dispatch": {
             "x": "min(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
@@ -258,8 +311,8 @@
           "id": "attention",
           "name": "SparseAttention.Attention",
           "shader": "sparse-attention.wgsl.jinja",
-          "derive": { "qTile": "sparseQueryTile" },
-          "bindings": ["q_rotary", "present_key_2", "present_value_2", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_2"],
           "dispatch": { "x": "sparseQueryTiles", "y": "batchSize * numHeads" }
         }
       ]
@@ -290,7 +343,7 @@
           "id": "qrotary",
           "name": "SparseAttention.QueryRotary",
           "shader": "sparse-q-rotary.wgsl.jinja",
-          "bindings": ["query", "cos_cache", "sin_cache", "q_rotary_2", "key_total_sequence_lengths", "total_sequence_length", "params"],
           "dispatch": {
             "x": "min(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
@@ -301,18 +354,77 @@
           "id": "attention",
           "name": "SparseAttention.AttentionSgmat",
           "shader": "sparse-attention-sgmat.wgsl.jinja",
-          "bindings": ["q_rotary", "present_key_3", "present_value_3", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_2"],
           "dispatch": { "x": "sgmatQueryTiles", "y": "batchSize * numHeads" },
-          "subgroupCollectivesWidth": 32
         }
       ],
       "demoteWhen": ["sgmatQueryTiles * batchSize * numHeads * sparseSgmatTileN < tunables.MATRIX_MIN_WORKGROUPS * 64"]
     },
     {
       "id": "packed",
       "priority": 0,
       "when": ["packedContract"],
-      "derive": { "vStageWorthIt": "sparseVStageWorthIt" },
       "passes": [
         {
           "id": "append",
@@ -330,8 +442,8 @@
           "id": "attention",
           "name": "SparseAttention.Attention",
           "shader": "sparse-attention.wgsl.jinja",
-          "derive": { "qTile": "sparseQueryTile" },
-          "bindings": ["query", "present_key_2", "present_value_2", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_2"],
           "dispatch": { "x": "sparseQueryTiles", "y": "batchSize * numHeads" }
         }
       ]
@@ -361,18 +473,65 @@
           "id": "attention",
           "name": "SparseAttention.AttentionSgmat",
           "shader": "sparse-attention-sgmat.wgsl.jinja",
-          "bindings": ["query", "present_key_3", "present_value_3", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_2"],
           "dispatch": { "x": "sgmatQueryTiles", "y": "batchSize * numHeads" },
-          "subgroupCollectivesWidth": 32
         }
       ],
       "demoteWhen": ["sgmatQueryTiles * batchSize * numHeads * sparseSgmatTileN < tunables.MATRIX_MIN_WORKGROUPS * 64"]
     },
     {
       "id": "packed_rotary",
       "priority": 10,
       "when": ["packedRotaryContract"],
-      "derive": { "vStageWorthIt": "sparseVStageWorthIt" },
       "intermediates": [{ "id": "QRotary", "dtype": "float32", "shape": "[qRotaryElements]" }],
       "passes": [
         {
@@ -391,7 +550,7 @@
           "id": "qrotary",
           "name": "SparseAttention.QueryRotary",
           "shader": "sparse-q-rotary.wgsl.jinja",
-          "bindings": ["query", "cos_cache", "sin_cache", "q_rotary_2", "key_total_sequence_lengths", "total_sequence_length", "params"],
           "dispatch": {
             "x": "min(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
@@ -402,8 +561,8 @@
           "id": "attention",
           "name": "SparseAttention.Attention",
           "shader": "sparse-attention.wgsl.jinja",
-          "derive": { "qTile": "sparseQueryTile" },
-          "bindings": ["q_rotary", "present_key_2", "present_value_2", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_2"],
           "dispatch": { "x": "sparseQueryTiles", "y": "batchSize * numHeads" }
         }
       ]
@@ -434,7 +593,7 @@
           "id": "qrotary",
           "name": "SparseAttention.QueryRotary",
           "shader": "sparse-q-rotary.wgsl.jinja",
-          "bindings": ["query", "cos_cache", "sin_cache", "q_rotary_2", "key_total_sequence_lengths", "total_sequence_length", "params"],
           "dispatch": {
             "x": "min(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
@@ -445,12 +604,71 @@
           "id": "attention",
           "name": "SparseAttention.AttentionSgmat",
           "shader": "sparse-attention-sgmat.wgsl.jinja",
-          "bindings": ["q_rotary", "present_key_3", "present_value_3", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_2"],
           "dispatch": { "x": "sgmatQueryTiles", "y": "batchSize * numHeads" },
-          "subgroupCollectivesWidth": 32
         }
       ],
       "demoteWhen": ["sgmatQueryTiles * batchSize * numHeads * sparseSgmatTileN < tunables.MATRIX_MIN_WORKGROUPS * 64"]
     }
   ]
 }

     "sparseBlockSize": "attrs.sparse_block_size",
     "headSize": "dim(shapes.pastKeyT, 3)",
     "headVec": "headSize / 4",
+    "sparseRequestedQueryTile": "max(tunables.QUERY_TILE, 8) if device.features.has(\"subgroups\") and wave32Effective and not device.features.has(\"shader-f16\") and ceilDiv(seqLen, 8) * batchSize * numHeads >= tunables.NARROW_MIN_WORKGROUPS else tunables.QUERY_TILE",
+    "sparseWidthBound": "min(256, max(32, pow2ceil(headVec))) if ceilDiv(seqLen, max(1, min(sparseBlockSize, sparseRequestedQueryTile))) * batchSize * numHeads >= tunables.NARROW_MIN_WORKGROUPS else max(256, tunables.WORKGROUP_SIZE)",
     "sparseQueryTileCap": "max(1, floor((device.limits.maxComputeWorkgroupStorageSize / 4 - sparseWidthBound) / (2 * headSize + 3 * sparseWidthBound)))",
+    "sparseQueryTileWant": "min(sparseRequestedQueryTile, 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",
     "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\"",
+    "sparseAttnBaseLdsBytes": "(2 * sparseQueryTile * headSize + (3 * sparseQueryTile + 1) * sparseAttnWorkgroup) * 4",
+    "geometryOk": "tunables.WORKGROUP_SIZE >= 1 and floor(tunables.WORKGROUP_SIZE) == tunables.WORKGROUP_SIZE and pow2ceil(tunables.WORKGROUP_SIZE) == tunables.WORKGROUP_SIZE and tunables.WORKGROUP_SIZE <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.WORKGROUP_SIZE <= device.limits.maxComputeWorkgroupSizeX and tunables.APPEND_WORKGROUP_SIZE >= 1 and floor(tunables.APPEND_WORKGROUP_SIZE) == tunables.APPEND_WORKGROUP_SIZE and tunables.APPEND_WORKGROUP_SIZE <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.APPEND_WORKGROUP_SIZE <= device.limits.maxComputeWorkgroupSizeX and sparseQueryTiles <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and batchSize * numHeads <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(ceilDiv(qRotaryElements, tunables.APPEND_WORKGROUP_SIZE), min(device.limits.maxComputeWorkgroupsPerDimension, 65535)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and sparseAttnBaseLdsBytes <= 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",
+    "sparseValueParts": "max(1, min(floor(sparseAttnWorkgroup / headVec), floor((device.limits.maxComputeWorkgroupStorageSize - sparseAttnBaseLdsBytes) / (sparseQueryTile * headSize * 4))))",
+    "sparseVStageWorthIt": "sparseAttnWorkgroups <= tunables.V_STAGE_MAX_WORKGROUPS and headVec <= sparseAttnWorkgroup and sparseAttnBaseLdsBytes + 16 * headSize * 4 <= device.limits.maxComputeWorkgroupStorageSize",
+    "sparseSgmatWorkgroup": "256",
+    "sparseSgmatTileM": "64",
+    "sgmatQueryTiles": "ceilDiv(seqLen, sparseSgmatTileM)",
+    "sgmatDirectQuery": "seqLen % sparseSgmatTileM == 0",
+    "sparseSgmatTileN": "64 if (64 * 32 + 64 * 64 + 64 * 3 + 128 * 2) * 4 <= device.limits.maxComputeWorkgroupStorageSize else 32",
     "sparseSgmatTileK": "sparseSgmatTileN / 2",
+    "sparseSgmatLdsBytes": "(64 * sparseSgmatTileK + 64 * sparseSgmatTileN + 64 * 3 + 128 * 2) * 4",
+    "sparseSgmatGeometryOk": "sparseSgmatWorkgroup <= device.limits.maxComputeInvocationsPerWorkgroup and sparseSgmatWorkgroup <= device.limits.maxComputeWorkgroupSizeX and sgmatQueryTiles <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and batchSize * numHeads <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and sparseSgmatLdsBytes <= device.limits.maxComputeWorkgroupStorageSize",
     "sparseSgmatOk": "tensorDtypes.queryT == \"float32\" and seqLen >= 64 and sparseBlockSize % 64 == 0 and headSize % 32 == 0 and headSize <= 128 and maxCacheSeq % 64 == 0 and device.features.has(\"subgroups\") and wave32Effective and device.features.has(\"chromium-experimental-subgroup-matrix\") and sparseSgmatGeometryOk",
     "scalar": "dtypes.T",
     "cacheVec": "\"vec4<f16>\" if dtypes.T == \"f16\" else \"vec4<f32>\"",
     "attnWorkgroup": "sparseAttnWorkgroup",
     "usesRotary": "useRotary",
+    "appendWorkgroupSize": "tunables.APPEND_WORKGROUP_SIZE",
+    "sparseTailRows": "seqLen % sparseSgmatTileM",
+    "sparsePrefixRows": "seqLen - sparseTailRows",
+    "sparseTailWorkgroup": "min(256, max(32, pow2ceil(headVec))) if batchSize * numHeads >= tunables.NARROW_MIN_WORKGROUPS else tunables.WORKGROUP_SIZE",
+    "sparseTailBaseLdsBytes": "(2 * sparseTailRows * headSize + (3 * sparseTailRows + 1) * sparseTailWorkgroup) * 4",
+    "sparseTailValueParts": "max(1, min(floor(sparseTailWorkgroup / headVec), floor((device.limits.maxComputeWorkgroupStorageSize - sparseTailBaseLdsBytes) / (max(1, sparseTailRows) * headSize * 4))))",
+    "sparseTailVStageWorthIt": "batchSize * numHeads <= tunables.V_STAGE_MAX_WORKGROUPS and headVec <= sparseTailWorkgroup and sparseTailBaseLdsBytes + 16 * headSize * 4 <= device.limits.maxComputeWorkgroupStorageSize",
+    "sparseTailOk": "sparseTailRows > 0 and sparseTailRows <= sparseQueryTile and sparsePrefixRows >= sparseSgmatTileM and sparseTailBaseLdsBytes <= device.limits.maxComputeWorkgroupStorageSize and sparseTailWorkgroup <= device.limits.maxComputeInvocationsPerWorkgroup and sparseTailWorkgroup <= device.limits.maxComputeWorkgroupSizeX"
   },
   "when": ["geometryOk"],
   "bindings": {
+    "new_key": { "arg": "keyT", "elementType": "$scalar" },
+    "new_value": { "arg": "valueT", "elementType": "$scalar" },
+    "present_key": { "arg": "pastKeyT", "elementType": "$scalar" },
+    "present_value": { "arg": "pastValueT", "elementType": "$scalar" },
+    "key_total_sequence_lengths": { "arg": "keyTotalSequenceLengthsT", "elementType": "i32" },
+    "total_sequence_length": { "arg": "totalSequenceLengthT", "elementType": "i32" },
     "params": {
       "struct": [
         { "name": "batchSize", "type": "u32", "value": "batchSize" },
         { "name": "seqLen", "type": "u32", "value": "seqLen" }
       ]
     },
+    "cos_cache": { "arg": "cosCacheT", "elementType": "$scalar" },
+    "sin_cache": { "arg": "sinCacheT", "elementType": "$scalar" },
+    "packed_qkv": { "arg": "queryT", "elementType": "$scalar" },
+    "query": { "arg": "queryT", "elementType": "$scalar" },
+    "present_key_packed": {
       "arg": "pastKeyT",
       "name": "present_key",
       "buffer": "read-only-storage",
       "elementType": "$cacheVec"
     },
+    "present_value_packed": {
       "arg": "pastValueT",
       "name": "present_value",
       "buffer": "read-only-storage",
       "elementType": "$cacheVec"
     },
+    "block_row_indices": { "arg": "blockRowIndicesT", "elementType": "i32" },
+    "block_col_indices": { "arg": "blockColIndicesT", "elementType": "i32" },
+    "output": { "arg": "outputT", "elementType": "$scalar" },
+    "params_packed": {
       "name": "params",
       "struct": [
         { "name": "seqLen", "type": "u32", "value": "seqLen" },
         { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" }
       ]
     },
     "q_rotary": { "scratch": "QRotary", "buffer": "read-only-storage", "elementType": "f32" },
+    "present_key_scalar": {
       "arg": "pastKeyT",
       "name": "present_key",
       "buffer": "read-only-storage",
       "elementType": "$scalar"
     },
+    "present_value_scalar": {
       "arg": "pastValueT",
       "name": "present_value",
       "buffer": "read-only-storage",
       "elementType": "$scalar"
     },
+    "q_rotary_f32": { "scratch": "QRotary", "name": "q_rotary", "elementType": "f32" }
   },
   "variants": [
     {
       "id": "separate",
       "priority": 0,
       "when": ["separateContract"],
+      "derive": { "valueParts": "sparseValueParts", "vStageWorthIt": "sparseVStageWorthIt and valueParts == 1" },
       "passes": [
         {
           "id": "append",
           "id": "attention",
           "name": "SparseAttention.Attention",
           "shader": "sparse-attention.wgsl.jinja",
+          "derive": { "qTile": "sparseQueryTile", "queryOffset": "0", "promptTail": "false" },
+          "bindings": ["query", "present_key_packed", "present_value_packed", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_packed"],
           "dispatch": { "x": "sparseQueryTiles", "y": "batchSize * numHeads" }
         }
       ]
           "id": "attention",
           "name": "SparseAttention.AttentionSgmat",
           "shader": "sparse-attention-sgmat.wgsl.jinja",
+          "bindings": ["query", "present_key_scalar", "present_value_scalar", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_packed"],
           "dispatch": { "x": "sgmatQueryTiles", "y": "batchSize * numHeads" },
+          "derive": { "splitQueryTail": "false" }
         }
       ],
       "demoteWhen": ["sgmatQueryTiles * batchSize * numHeads * sparseSgmatTileN < tunables.MATRIX_MIN_WORKGROUPS * 64"]
     },
+    {
+      "id": "separate_sgmat_tail",
+      "priority": 21,
+      "when": ["separateContract", "sparseSgmatOk", "sparseTailOk"],
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
+      },
+      "passes": [
+        {
+          "id": "append",
+          "name": "SparseAttention.Append",
+          "shader": "sparse-kv-append.wgsl.jinja",
+          "derive": { "kvSource": "\"new_key\"", "vSource": "\"new_value\"" },
+          "bindings": ["new_key", "new_value", "present_key", "present_value", "key_total_sequence_lengths", "total_sequence_length", "params"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "z": 1
+          }
+        },
+        {
+          "id": "attention",
+          "name": "SparseAttention.AttentionSgmat",
+          "shader": "sparse-attention-sgmat.wgsl.jinja",
+          "bindings": ["query", "present_key_scalar", "present_value_scalar", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_packed"],
+          "dispatch": { "x": "sgmatQueryTiles", "y": "batchSize * numHeads" },
+          "derive": { "splitQueryTail": "true" }
+        },
+        {
+          "id": "tail",
+          "name": "SparseAttention.Tail",
+          "shader": "sparse-attention.wgsl.jinja",
+          "derive": {
+            "qTile": "sparseTailRows",
+            "queryOffset": "sparsePrefixRows",
+            "attnWorkgroup": "sparseTailWorkgroup",
+            "valueParts": "sparseTailValueParts",
+            "vStageWorthIt": "sparseTailVStageWorthIt and sparseTailValueParts == 1",
+            "promptTail": "true"
+          },
+          "bindings": ["query", "present_key_packed", "present_value_packed", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_packed"],
+          "dispatch": { "x": "1", "y": "batchSize * numHeads" }
+        }
+      ],
+      "demoteWhen": ["sparsePrefixRows / sparseSgmatTileM * batchSize * numHeads * sparseSgmatTileN < tunables.MATRIX_MIN_WORKGROUPS * 64", "headSize < sparseSgmatTileM"]
+    },
     {
       "id": "separate_rotary",
       "priority": 10,
       "when": ["separateRotaryContract"],
+      "derive": { "valueParts": "sparseValueParts", "vStageWorthIt": "sparseVStageWorthIt and valueParts == 1" },
       "intermediates": [{ "id": "QRotary", "dtype": "float32", "shape": "[qRotaryElements]" }],
       "passes": [
         {
           "id": "qrotary",
           "name": "SparseAttention.QueryRotary",
           "shader": "sparse-q-rotary.wgsl.jinja",
+          "bindings": ["query", "cos_cache", "sin_cache", "q_rotary_f32", "key_total_sequence_lengths", "total_sequence_length", "params"],
           "dispatch": {
             "x": "min(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
           "id": "attention",
           "name": "SparseAttention.Attention",
           "shader": "sparse-attention.wgsl.jinja",
+          "derive": { "qTile": "sparseQueryTile", "queryOffset": "0", "promptTail": "false" },
+          "bindings": ["q_rotary", "present_key_packed", "present_value_packed", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_packed"],
           "dispatch": { "x": "sparseQueryTiles", "y": "batchSize * numHeads" }
         }
       ]
           "id": "qrotary",
           "name": "SparseAttention.QueryRotary",
           "shader": "sparse-q-rotary.wgsl.jinja",
+          "bindings": ["query", "cos_cache", "sin_cache", "q_rotary_f32", "key_total_sequence_lengths", "total_sequence_length", "params"],
           "dispatch": {
             "x": "min(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
           "id": "attention",
           "name": "SparseAttention.AttentionSgmat",
           "shader": "sparse-attention-sgmat.wgsl.jinja",
+          "bindings": ["q_rotary", "present_key_scalar", "present_value_scalar", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_packed"],
           "dispatch": { "x": "sgmatQueryTiles", "y": "batchSize * numHeads" },
+          "derive": { "splitQueryTail": "false" }
         }
       ],
       "demoteWhen": ["sgmatQueryTiles * batchSize * numHeads * sparseSgmatTileN < tunables.MATRIX_MIN_WORKGROUPS * 64"]
     },
+    {
+      "id": "separate_rotary_sgmat_tail",
+      "priority": 31,
+      "when": ["separateRotaryContract", "sparseSgmatOk", "sparseTailOk"],
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
+      },
+      "intermediates": [{ "id": "QRotary", "dtype": "float32", "shape": "[qRotaryElements]" }],
+      "passes": [
+        {
+          "id": "append",
+          "name": "SparseAttention.Append",
+          "shader": "sparse-kv-append.wgsl.jinja",
+          "derive": { "kvSource": "\"new_key\"", "vSource": "\"new_value\"" },
+          "bindings": ["new_key", "new_value", "present_key", "present_value", "key_total_sequence_lengths", "total_sequence_length", "cos_cache", "sin_cache", "params"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "z": 1
+          }
+        },
+        {
+          "id": "qrotary",
+          "name": "SparseAttention.QueryRotary",
+          "shader": "sparse-q-rotary.wgsl.jinja",
+          "bindings": ["query", "cos_cache", "sin_cache", "q_rotary_f32", "key_total_sequence_lengths", "total_sequence_length", "params"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "z": 1
+          }
+        },
+        {
+          "id": "attention",
+          "name": "SparseAttention.AttentionSgmat",
+          "shader": "sparse-attention-sgmat.wgsl.jinja",
+          "bindings": ["q_rotary", "present_key_scalar", "present_value_scalar", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_packed"],
+          "dispatch": { "x": "sgmatQueryTiles", "y": "batchSize * numHeads" },
+          "derive": { "splitQueryTail": "true" }
+        },
+        {
+          "id": "tail",
+          "name": "SparseAttention.Tail",
+          "shader": "sparse-attention.wgsl.jinja",
+          "derive": {
+            "qTile": "sparseTailRows",
+            "queryOffset": "sparsePrefixRows",
+            "attnWorkgroup": "sparseTailWorkgroup",
+            "valueParts": "sparseTailValueParts",
+            "vStageWorthIt": "sparseTailVStageWorthIt and sparseTailValueParts == 1",
+            "promptTail": "true"
+          },
+          "bindings": ["q_rotary", "present_key_packed", "present_value_packed", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_packed"],
+          "dispatch": { "x": "1", "y": "batchSize * numHeads" }
+        }
+      ],
+      "demoteWhen": ["sparsePrefixRows / sparseSgmatTileM * batchSize * numHeads * sparseSgmatTileN < tunables.MATRIX_MIN_WORKGROUPS * 64", "headSize < sparseSgmatTileM"]
+    },
     {
       "id": "packed",
       "priority": 0,
       "when": ["packedContract"],
+      "derive": { "valueParts": "sparseValueParts", "vStageWorthIt": "sparseVStageWorthIt and valueParts == 1" },
       "passes": [
         {
           "id": "append",
           "id": "attention",
           "name": "SparseAttention.Attention",
           "shader": "sparse-attention.wgsl.jinja",
+          "derive": { "qTile": "sparseQueryTile", "queryOffset": "0", "promptTail": "false" },
+          "bindings": ["query", "present_key_packed", "present_value_packed", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_packed"],
           "dispatch": { "x": "sparseQueryTiles", "y": "batchSize * numHeads" }
         }
       ]
           "id": "attention",
           "name": "SparseAttention.AttentionSgmat",
           "shader": "sparse-attention-sgmat.wgsl.jinja",
+          "bindings": ["query", "present_key_scalar", "present_value_scalar", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_packed"],
           "dispatch": { "x": "sgmatQueryTiles", "y": "batchSize * numHeads" },
+          "derive": { "splitQueryTail": "false" }
         }
       ],
       "demoteWhen": ["sgmatQueryTiles * batchSize * numHeads * sparseSgmatTileN < tunables.MATRIX_MIN_WORKGROUPS * 64"]
     },
+    {
+      "id": "packed_sgmat_tail",
+      "priority": 21,
+      "when": ["packedContract", "sparseSgmatOk", "sparseTailOk"],
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
+      },
+      "passes": [
+        {
+          "id": "append",
+          "name": "SparseAttention.Append",
+          "shader": "sparse-kv-append.wgsl.jinja",
+          "derive": { "kvSource": "\"packed_qkv\"", "vSource": "\"packed_qkv\"" },
+          "bindings": ["packed_qkv", "present_key", "present_value", "key_total_sequence_lengths", "total_sequence_length", "params"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "z": 1
+          }
+        },
+        {
+          "id": "attention",
+          "name": "SparseAttention.AttentionSgmat",
+          "shader": "sparse-attention-sgmat.wgsl.jinja",
+          "bindings": ["query", "present_key_scalar", "present_value_scalar", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_packed"],
+          "dispatch": { "x": "sgmatQueryTiles", "y": "batchSize * numHeads" },
+          "derive": { "splitQueryTail": "true" }
+        },
+        {
+          "id": "tail",
+          "name": "SparseAttention.Tail",
+          "shader": "sparse-attention.wgsl.jinja",
+          "derive": {
+            "qTile": "sparseTailRows",
+            "queryOffset": "sparsePrefixRows",
+            "attnWorkgroup": "sparseTailWorkgroup",
+            "valueParts": "sparseTailValueParts",
+            "vStageWorthIt": "sparseTailVStageWorthIt and sparseTailValueParts == 1",
+            "promptTail": "true"
+          },
+          "bindings": ["query", "present_key_packed", "present_value_packed", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_packed"],
+          "dispatch": { "x": "1", "y": "batchSize * numHeads" }
+        }
+      ],
+      "demoteWhen": ["sparsePrefixRows / sparseSgmatTileM * batchSize * numHeads * sparseSgmatTileN < tunables.MATRIX_MIN_WORKGROUPS * 64", "headSize < sparseSgmatTileM"]
+    },
     {
       "id": "packed_rotary",
       "priority": 10,
       "when": ["packedRotaryContract"],
+      "derive": { "valueParts": "sparseValueParts", "vStageWorthIt": "sparseVStageWorthIt and valueParts == 1" },
       "intermediates": [{ "id": "QRotary", "dtype": "float32", "shape": "[qRotaryElements]" }],
       "passes": [
         {
           "id": "qrotary",
           "name": "SparseAttention.QueryRotary",
           "shader": "sparse-q-rotary.wgsl.jinja",
+          "bindings": ["query", "cos_cache", "sin_cache", "q_rotary_f32", "key_total_sequence_lengths", "total_sequence_length", "params"],
           "dispatch": {
             "x": "min(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
           "id": "attention",
           "name": "SparseAttention.Attention",
           "shader": "sparse-attention.wgsl.jinja",
+          "derive": { "qTile": "sparseQueryTile", "queryOffset": "0", "promptTail": "false" },
+          "bindings": ["q_rotary", "present_key_packed", "present_value_packed", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_packed"],
           "dispatch": { "x": "sparseQueryTiles", "y": "batchSize * numHeads" }
         }
       ]
           "id": "qrotary",
           "name": "SparseAttention.QueryRotary",
           "shader": "sparse-q-rotary.wgsl.jinja",
+          "bindings": ["query", "cos_cache", "sin_cache", "q_rotary_f32", "key_total_sequence_lengths", "total_sequence_length", "params"],
           "dispatch": {
             "x": "min(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
           "id": "attention",
           "name": "SparseAttention.AttentionSgmat",
           "shader": "sparse-attention-sgmat.wgsl.jinja",
+          "bindings": ["q_rotary", "present_key_scalar", "present_value_scalar", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_packed"],
           "dispatch": { "x": "sgmatQueryTiles", "y": "batchSize * numHeads" },
+          "derive": { "splitQueryTail": "false" }
         }
       ],
       "demoteWhen": ["sgmatQueryTiles * batchSize * numHeads * sparseSgmatTileN < tunables.MATRIX_MIN_WORKGROUPS * 64"]
+    },
+    {
+      "id": "packed_rotary_sgmat_tail",
+      "priority": 31,
+      "when": ["packedRotaryContract", "sparseSgmatOk", "sparseTailOk"],
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
+      },
+      "intermediates": [{ "id": "QRotary", "dtype": "float32", "shape": "[qRotaryElements]" }],
+      "passes": [
+        {
+          "id": "append",
+          "name": "SparseAttention.Append",
+          "shader": "sparse-kv-append.wgsl.jinja",
+          "derive": { "kvSource": "\"packed_qkv\"", "vSource": "\"packed_qkv\"" },
+          "bindings": ["packed_qkv", "present_key", "present_value", "key_total_sequence_lengths", "total_sequence_length", "cos_cache", "sin_cache", "params"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize * kvNumHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "z": 1
+          }
+        },
+        {
+          "id": "qrotary",
+          "name": "SparseAttention.QueryRotary",
+          "shader": "sparse-q-rotary.wgsl.jinja",
+          "bindings": ["query", "cos_cache", "sin_cache", "q_rotary_f32", "key_total_sequence_lengths", "total_sequence_length", "params"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize * numHeads * seqLen * headSize), (appendWorkgroupSize)), 65535)",
+            "z": 1
+          }
+        },
+        {
+          "id": "attention",
+          "name": "SparseAttention.AttentionSgmat",
+          "shader": "sparse-attention-sgmat.wgsl.jinja",
+          "bindings": ["q_rotary", "present_key_scalar", "present_value_scalar", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_packed"],
+          "dispatch": { "x": "sgmatQueryTiles", "y": "batchSize * numHeads" },
+          "derive": { "splitQueryTail": "true" }
+        },
+        {
+          "id": "tail",
+          "name": "SparseAttention.Tail",
+          "shader": "sparse-attention.wgsl.jinja",
+          "derive": {
+            "qTile": "sparseTailRows",
+            "queryOffset": "sparsePrefixRows",
+            "attnWorkgroup": "sparseTailWorkgroup",
+            "valueParts": "sparseTailValueParts",
+            "vStageWorthIt": "sparseTailVStageWorthIt and sparseTailValueParts == 1",
+            "promptTail": "true"
+          },
+          "bindings": ["q_rotary", "present_key_packed", "present_value_packed", "block_row_indices", "block_col_indices", "key_total_sequence_lengths", "total_sequence_length", "output", "params_packed"],
+          "dispatch": { "x": "1", "y": "batchSize * numHeads" }
+        }
+      ],
+      "demoteWhen": ["sparsePrefixRows / sparseSgmatTileM * batchSize * numHeads * sparseSgmatTileN < tunables.MATRIX_MIN_WORKGROUPS * 64", "headSize < sparseSgmatTileM"]
     }
   ]
 }

build/webgpu/metadata.json CHANGED Viewed

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

 {
   "name": "com.microsoft.SparseAttention",
+  "id": "_com_microsoft_sparseattention_webgpu_c2657b2",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "4NNfSy0Pz7J1r5bacbFNwgbyHBbeJ3OYn3iajMdCQ64=",
+      "manifest.json": "8xLb9NHTZ4tCa1RARx5gvMP/WuEU4tsurxnNVELHnPc=",
+      "sparse-attention-sgmat.wgsl.jinja": "bUA2Sjw4qg7s/2KvOkUddx7N6v718rkkQNG7yQko+gI=",
+      "sparse-attention.wgsl.jinja": "rrRJNjo17muFxfSrY5YgMA9zkYih1zSDLYwZqDbViKQ=",
+      "sparse-kv-append.wgsl.jinja": "P1nJPEk4RcBb6GcajLnqUqMRV15GsBkIwCfLWBls6YE=",
+      "sparse-q-rotary.wgsl.jinja": "fTkf6Ht1LF4ygO2DviYl8kYqoGd7LKL3bSD1XBG+1sg=",
+      "test.json": "wyUajoZ/2iaiyzrdWd5+dceFdPQAmEAy/ol6vKTe0os="
     }
   },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
   "webgpu": {
+    "manifestSpec": "2.1",
     "variants": {
       "separate": ["sparse-attention.wgsl.jinja", "sparse-kv-append.wgsl.jinja"],
       "separate_sgmat": ["sparse-attention-sgmat.wgsl.jinja", "sparse-kv-append.wgsl.jinja"],
+      "separate_sgmat_tail": ["sparse-attention-sgmat.wgsl.jinja", "sparse-attention.wgsl.jinja", "sparse-kv-append.wgsl.jinja"],
       "separate_rotary": ["sparse-attention.wgsl.jinja", "sparse-kv-append.wgsl.jinja", "sparse-q-rotary.wgsl.jinja"],
       "separate_rotary_sgmat": ["sparse-attention-sgmat.wgsl.jinja", "sparse-kv-append.wgsl.jinja", "sparse-q-rotary.wgsl.jinja"],
+      "separate_rotary_sgmat_tail": ["sparse-attention-sgmat.wgsl.jinja", "sparse-attention.wgsl.jinja", "sparse-kv-append.wgsl.jinja", "sparse-q-rotary.wgsl.jinja"],
       "packed": ["sparse-attention.wgsl.jinja", "sparse-kv-append.wgsl.jinja"],
       "packed_sgmat": ["sparse-attention-sgmat.wgsl.jinja", "sparse-kv-append.wgsl.jinja"],
+      "packed_sgmat_tail": ["sparse-attention-sgmat.wgsl.jinja", "sparse-attention.wgsl.jinja", "sparse-kv-append.wgsl.jinja"],
       "packed_rotary": ["sparse-attention.wgsl.jinja", "sparse-kv-append.wgsl.jinja", "sparse-q-rotary.wgsl.jinja"],
+      "packed_rotary_sgmat": ["sparse-attention-sgmat.wgsl.jinja", "sparse-kv-append.wgsl.jinja", "sparse-q-rotary.wgsl.jinja"],
+      "packed_rotary_sgmat_tail": ["sparse-attention-sgmat.wgsl.jinja", "sparse-attention.wgsl.jinja", "sparse-kv-append.wgsl.jinja", "sparse-q-rotary.wgsl.jinja"]
     }
   }
 }

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

@@ -10,8 +10,7 @@ fn past_sequence_length(batch: u32) -> u32 {
   let total = u32(key_total_sequence_lengths[batch]);
   return select(0u, total - params.seqLen, total >= params.seqLen);
 }
-{%- endmacro %}
 enable subgroups;
 {% if pinSubgroupSize32 %}
 enable subgroup_size_control;
@@ -19,19 +18,18 @@ enable subgroup_size_control;
 enable chromium_experimental_subgroup_matrix;
 diagnostic(off, chromium.subgroup_matrix_uniformity);
 {{ env.wgsl.resourceDeclarations }}
 // Subgroup-matrix attention over 64-query tiles. Each workgroup processes one
-// `(batch, query tile, query head)` tuple. Selected sparse blocks are traversed
-// as 64-query score tiles; their contiguous K/V cache rows load directly as matrix
-// fragments. Sequences of complete query tiles also load Q directly; sequences
-// with a partial final tile stage Q with zero padding. The manifest bounds tile widths by workgroup storage.
 //
-// The first sweep folds score tiles into per-row `(max, denominator)` softmax
-// statistics. The second recomputes the same scores, applies the completed
-// normalization, and accumulates P.V into result fragments. Causal bounds,
-// duplicate CSR columns, dense rows, and all-masked rows are handled explicitly.
 const Q_HEADS: u32 = {{ numHeads }}u;
 const KV_HEADS: u32 = {{ kvNumHeads }}u;
 const HEAD_DIM: u32 = {{ headSize }}u;
@@ -47,7 +45,7 @@ const Q_STRIDE: u32 = {{ packedStride if packedQkv else numHeads * headSize }}u;
 // Eight 32-lane subgroups form a 4x2 grid. The device storage budget chooses
 // 64 or 32 key columns and a matching head-dimension staging width. Both divide
 // the admitted sparse-block and head dimensions without a key or head tail.
-const TILE_M: u32 = 64u;
 const TILE_N: u32 = {{ sparseSgmatTileN }}u;
 const TILE_K: u32 = {{ sparseSgmatTileK }}u;
 const SUB_TILES: u32 = {{ (sparseBlockSize / sparseSgmatTileN) | int }}u;
@@ -74,49 +72,41 @@ 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));
 }
-// Q staging for the score GEMM; the score epilogues alias it as the
-// fragment-store scratch (8 subgroups x scoreColBlocks banks x 64 elements).
 var<workgroup> tile_q: array<f32, {{ 64 * sparseSgmatTileK }}>;
 // Tile probabilities for the P.V GEMM; the output epilogue aliases it as the
 // result-fragment scratch once the last key tile's readers are done.
 var<workgroup> prob_tile: array<f32, {{ 64 * sparseSgmatTileN }}>;
 var<workgroup> row_m: array<f32, 64>;
 var<workgroup> row_d: array<f32, 64>;
 // Per-key-tile row partials, one slot per (row, subgroup column group).
 var<workgroup> part_m: array<f32, 128>;
 var<workgroup> part_d: array<f32, 128>;
-{% 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() }}
-{% set queryMatrixSource = "tile_q" if not sgmatDirectQuery else ("q_rotary" if usesRotary else "query") %}
-{% set queryMatrixStride = "TILE_K" if not sgmatDirectQuery else ("HEAD_DIM" if usesRotary else "Q_STRIDE") %}
-{% macro query_matrix_offset(rb) -%}
-{% if not sgmatDirectQuery -%}
-(base_a + {{ rb * 8 }}u) * TILE_K + step
-{%- elif usesRotary -%}
-((batch * Q_HEADS + head) * params.seqLen + tile0 + base_a + {{ rb * 8 }}u) * HEAD_DIM + k_base + step
-{%- else -%}
-(batch * params.seqLen + tile0 + base_a + {{ rb * 8 }}u) * Q_STRIDE + head * HEAD_DIM + k_base + step
-{%- endif %}
-{%- endmacro %}
-{% macro score_tile() %}
         // S = Q.K^T, retaining the same sequence of 8-wide matrix operations.
         // Fully populated query tiles need neither staging nor K-loop barriers.
-        // A sequence with a partial query tile retains zero-padded staging.
         for (var k_base = 0u; k_base < HEAD_DIM; k_base += TILE_K) {
-{% if not sgmatDirectQuery %}
           {
             let a_row = li / 4u;
             let a_col = (li % 4u) * {{ (sparseSgmatTileK / 4) | int }}u;
@@ -140,7 +130,7 @@ fn scale_value() -> f32 {
           for (var step = 0u; step < TILE_K; step += 8u) {
 {% for rb in range(2) %}
             let mat_a{{ rb }} = subgroupMatrixLoad<subgroup_matrix_left<f32, 8, 8>, row_major>(
-              &{{ queryMatrixSource }}, {{ query_matrix_offset(rb) }}, {{ queryMatrixStride }}
             );
 {% endfor %}
 {% for cb in range(scoreColBlocks) %}
@@ -156,13 +146,13 @@ fn scale_value() -> f32 {
 {% endfor %}
 {% endfor %}
           }
-{% if not sgmatDirectQuery %}
           workgroupBarrier();
 {% endif %}
         }
 {% endmacro %}
-{% macro sweep(phase) %}
   // Consecutive queries 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 its own row's tiles.
@@ -203,20 +193,28 @@ fn scale_value() -> f32 {
         var mat_s{{ rb }}{{ cb }}: subgroup_matrix_result<f32, 8, 8>;
 {% endfor %}
 {% endfor %}
-{{ score_tile() }}
 {% for rb in range(2) %}
 {% if rb > 0 %}
         // The banks alias the Q staging tile; the previous row block's readers
         // must finish before this one overwrites them.
         workgroupBarrier();
 {% endif %}
-{% if phase == "stats" %}
         // All four lanes of a quad carry the same score row (row_in_block is
         // lane / 4), so the accumulator below is a partial over one row and
         // the butterfly merging it is quad-uniform.
         var tile_stat_m{{ rb }} = -FLT_MAX;
         var tile_stat_d{{ rb }} = 0.0;
-{% endif %}
 {% for cb in range(scoreColBlocks) %}
         subgroupMatrixStore<row_major>(
           &tile_q, (subgroup * {{ scoreColBlocks }}u + {{ cb }}u) * 64u, mat_s{{ rb }}{{ cb }}, 8u
@@ -229,28 +227,21 @@ fn scale_value() -> f32 {
           let key = key_base + base_b + {{ cb * 8 }}u + col_in_block + pair;
           let q_abs = q_abs0 + r;
           let allowed = r < rows_live && q_abs / SPARSE_BLOCK == mask_row && key <= q_abs;
-{% if phase == "stats" %}
           if (allowed) {
-            let scored = tile_q[
-              (subgroup * {{ scoreColBlocks }}u + {{ cb }}u) * 64u + row_in_block * 8u + col_in_block + pair
-            ] * scale;
             let prev_m = tile_stat_m{{ rb }};
             tile_stat_m{{ rb }} = max(tile_stat_m{{ rb }}, scored);
             tile_stat_d{{ rb }} = tile_stat_d{{ rb }} * exp_shift(prev_m, tile_stat_m{{ rb }})
               + exp_shift(scored, tile_stat_m{{ rb }});
           }
-{% else %}
-          var prob = 0.0;
-          if (allowed) {
-            prob = exp_shift(tile_q[
-              (subgroup * {{ scoreColBlocks }}u + {{ cb }}u) * 64u + row_in_block * 8u + col_in_block + pair
-            ] * scale, row_m[r]);
-          }
-          prob_tile[r * TILE_N + base_b + {{ cb * 8 }}u + col_in_block + pair] = prob;
-{% endif %}
         }
 {% endfor %}
-{% if phase == "stats" %}
         // Butterfly the quad unconditionally: a lane whose row ran past the
         // query tail carries the exact identity (-FLT_MAX, 0), which merges to
         // a no-op, and a subgroup shuffle under a partial guard would not be
@@ -270,9 +261,9 @@ fn scale_value() -> f32 {
           part_m[stat_row * 2u + subtile_idx] = tile_stat_m{{ rb }};
           part_d[stat_row * 2u + subtile_idx] = tile_stat_d{{ rb }};
         }
-{% endif %}
 {% endfor %}
-{% if phase == "stats" %}
         workgroupBarrier();
         // Fold both column groups' partials into the running row statistics,
         // in the same (max, rescale, add) form as the per-lane walk: an empty
@@ -288,11 +279,44 @@ fn scale_value() -> f32 {
             merged_d = merged_d * exp_shift(merged_m, new_m) + d2 * exp_shift(m2, new_m);
             merged_m = new_m;
           }
           row_m[li] = merged_m;
           row_d[li] = merged_d;
         }
         workgroupBarrier();
-{% else %}
         workgroupBarrier();
         // O += P.V. P streams from shared memory; V rows are contiguous cache
         // rows, loaded directly as right-hand fragments. A masked or padded
@@ -319,19 +343,22 @@ fn scale_value() -> f32 {
         }
         // Orders this tile's prob_tile reads before the next tile rewrites it.
         workgroupBarrier();
-{% endif %}
       }
     }
   }
 {% endmacro %}
-@compute @workgroup_size(256, 1, 1){{ " @subgroup_size(32)" if pinSubgroupSize32 else "" }}
 fn main(
   @builtin(workgroup_id) wg: vec3<u32>,
   @builtin(local_invocation_index) li: u32,
   @builtin(subgroup_invocation_id) lane: u32
 ) {
   let tile0 = wg.x * TILE_M;
   let head = wg.y % Q_HEADS;
   let batch = wg.y / Q_HEADS;
@@ -367,14 +394,13 @@ fn main(
 {% endfor %}
 {% endfor %}
-  for (var r = li; r < TILE_M; r += 256u) {
     row_m[r] = -FLT_MAX;
     row_d[r] = 0.0;
   }
   workgroupBarrier();
-{{ sweep("stats") }}
-{{ sweep("apply") }}
   // Normalize by the final denominators and store. Publish only as many
   // fragment columns per batch as fit prob_tile, keeping the smaller key tile's
@@ -413,7 +439,7 @@ fn main(
     if (!(row_d[r] > 0.0)) {
       let key_bound = q_abs0 + r + 1u;
       let cache_base = (batch * KV_HEADS + kv_head) * MAX_CACHE_SEQ * HEAD_DIM;
-      for (var d = li; d < HEAD_DIM; d += 256u) {
         var total = 0.0;
         for (var key = 0u; key < key_bound; key++) {
           total += f32(present_value[cache_base + key * HEAD_DIM + d]);

   let total = u32(key_total_sequence_lengths[batch]);
   return select(0u, total - params.seqLen, total >= params.seqLen);
 }
+{% endmacro %}
 enable subgroups;
 {% if pinSubgroupSize32 %}
 enable subgroup_size_control;
 enable chromium_experimental_subgroup_matrix;
 diagnostic(off, chromium.subgroup_matrix_uniformity);
 {{ env.wgsl.resourceDeclarations }}
 // Subgroup-matrix attention over 64-query tiles. Each workgroup processes one
+// (batch, query tile, query head). K/V rows and complete query tiles load
+// directly as matrix fragments; only the final partial query tile stages Q.
+// The device workgroup-storage budget bounds the key and staging tile widths.
 //
+// One score sweep updates each row's online softmax statistics and rescales its
+// accumulated P.V fragments before adding the current tile. The score scratch
+// doubles as a bounded bank for fragment rescaling, avoiding a second QK sweep
+// or a full output staging array. Causal bounds, duplicate CSR columns, dense
+// rows, and all-masked rows retain the shared sparse-attention semantics.
 const Q_HEADS: u32 = {{ numHeads }}u;
 const KV_HEADS: u32 = {{ kvNumHeads }}u;
 const HEAD_DIM: u32 = {{ headSize }}u;
 // Eight 32-lane subgroups form a 4x2 grid. The device storage budget chooses
 // 64 or 32 key columns and a matching head-dimension staging width. Both divide
 // the admitted sparse-block and head dimensions without a key or head tail.
+const TILE_M: u32 = {{ sparseSgmatTileM }}u;
 const TILE_N: u32 = {{ sparseSgmatTileN }}u;
 const TILE_K: u32 = {{ sparseSgmatTileK }}u;
 const SUB_TILES: u32 = {{ (sparseBlockSize / sparseSgmatTileN) | int }}u;
   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));
 }
+// Q staging for the score GEMM; score epilogues and output rescaling alias
+// these banks (8 subgroups x scoreColBlocks banks x 64 elements).
 var<workgroup> tile_q: array<f32, {{ 64 * sparseSgmatTileK }}>;
 // Tile probabilities for the P.V GEMM; the output epilogue aliases it as the
 // result-fragment scratch once the last key tile's readers are done.
 var<workgroup> prob_tile: array<f32, {{ 64 * sparseSgmatTileN }}>;
 var<workgroup> row_m: array<f32, 64>;
 var<workgroup> row_d: array<f32, 64>;
+var<workgroup> row_c: array<f32, 64>;
 // Per-key-tile row partials, one slot per (row, subgroup column group).
 var<workgroup> part_m: array<f32, 128>;
 var<workgroup> part_d: array<f32, 128>;
+{% set ATTN_SCALE_DIM = "HEAD_DIM" %}fn scale_value() -> f32 {
   if (params.scale != 0.0) { return params.scale; }
   return inverseSqrt(f32({{ ATTN_SCALE_DIM }}));
 }
 {{ sparse_schedule() }}
+{% macro score_tile(direct) %}
+{% set queryMatrixSource = "tile_q" if not direct else ("q_rotary" if usesRotary else "query") %}
+{% set queryMatrixStride = "TILE_K" if not direct else ("HEAD_DIM" if usesRotary else "Q_STRIDE") %}
+{% set queryRowOrigin = "" if not direct else ("(batch * Q_HEADS + head) * params.seqLen + tile0 + " if usesRotary else "batch * params.seqLen + tile0 + ") %}
+{% set queryColOrigin = "" if not direct else ("k_base + " if usesRotary else "head * HEAD_DIM + k_base + ") %}
         // S = Q.K^T, retaining the same sequence of 8-wide matrix operations.
         // Fully populated query tiles need neither staging nor K-loop barriers.
+        // Only a partial query tile needs zero-padded staging.
         for (var k_base = 0u; k_base < HEAD_DIM; k_base += TILE_K) {
+{% if not direct %}
           {
             let a_row = li / 4u;
             let a_col = (li % 4u) * {{ (sparseSgmatTileK / 4) | int }}u;
           for (var step = 0u; step < TILE_K; step += 8u) {
 {% for rb in range(2) %}
             let mat_a{{ rb }} = subgroupMatrixLoad<subgroup_matrix_left<f32, 8, 8>, row_major>(
+              &{{ queryMatrixSource }}, ({{ queryRowOrigin }}base_a + {{ rb * 8 }}u) * {{ queryMatrixStride }} + {{ queryColOrigin }}step, {{ queryMatrixStride }}
             );
 {% endfor %}
 {% for cb in range(scoreColBlocks) %}
 {% endfor %}
 {% endfor %}
           }
+{% if not direct %}
           workgroupBarrier();
 {% endif %}
         }
 {% endmacro %}
+{% macro sweep() %}
   // Consecutive queries 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 its own row's tiles.
         var mat_s{{ rb }}{{ cb }}: subgroup_matrix_result<f32, 8, 8>;
 {% endfor %}
 {% endfor %}
+{% if sgmatDirectQuery %}
+{{ score_tile(true) }}
+{% else %}
+        if (rows_live == TILE_M) {
+{{ score_tile(true) }}
+        } else {
+{{ score_tile(false) }}
+        }
+{% endif %}
 {% for rb in range(2) %}
 {% if rb > 0 %}
         // The banks alias the Q staging tile; the previous row block's readers
         // must finish before this one overwrites them.
         workgroupBarrier();
 {% endif %}
         // All four lanes of a quad carry the same score row (row_in_block is
         // lane / 4), so the accumulator below is a partial over one row and
         // the butterfly merging it is quad-uniform.
         var tile_stat_m{{ rb }} = -FLT_MAX;
         var tile_stat_d{{ rb }} = 0.0;
 {% for cb in range(scoreColBlocks) %}
         subgroupMatrixStore<row_major>(
           &tile_q, (subgroup * {{ scoreColBlocks }}u + {{ cb }}u) * 64u, mat_s{{ rb }}{{ cb }}, 8u
           let key = key_base + base_b + {{ cb * 8 }}u + col_in_block + pair;
           let q_abs = q_abs0 + r;
           let allowed = r < rows_live && q_abs / SPARSE_BLOCK == mask_row && key <= q_abs;
+          let scored = tile_q[
+            (subgroup * {{ scoreColBlocks }}u + {{ cb }}u) * 64u + row_in_block * 8u + col_in_block + pair
+          ] * scale;
+          prob_tile[r * TILE_N + base_b + {{ cb * 8 }}u + col_in_block + pair] = scored;
           if (allowed) {
             let prev_m = tile_stat_m{{ rb }};
             tile_stat_m{{ rb }} = max(tile_stat_m{{ rb }}, scored);
             tile_stat_d{{ rb }} = tile_stat_d{{ rb }} * exp_shift(prev_m, tile_stat_m{{ rb }})
               + exp_shift(scored, tile_stat_m{{ rb }});
           }
         }
 {% endfor %}
         // Butterfly the quad unconditionally: a lane whose row ran past the
         // query tail carries the exact identity (-FLT_MAX, 0), which merges to
         // a no-op, and a subgroup shuffle under a partial guard would not be
           part_m[stat_row * 2u + subtile_idx] = tile_stat_m{{ rb }};
           part_d[stat_row * 2u + subtile_idx] = tile_stat_d{{ rb }};
         }
 {% endfor %}
         workgroupBarrier();
         // Fold both column groups' partials into the running row statistics,
         // in the same (max, rescale, add) form as the per-lane walk: an empty
             merged_d = merged_d * exp_shift(merged_m, new_m) + d2 * exp_shift(m2, new_m);
             merged_m = new_m;
           }
+          row_c[li] = exp_shift(row_m[li], merged_m);
           row_m[li] = merged_m;
           row_d[li] = merged_d;
         }
         workgroupBarrier();
+        // Rescale the persistent output fragments through the existing score
+        // scratch, using only as many banks as this device's key tile provides.
+{% for rb in range(2) %}
+{% for cbBase in range(0, pvColBlocks, scoreColBlocks) %}
+{% set cbEnd = pvColBlocks if pvColBlocks < cbBase + scoreColBlocks else cbBase + scoreColBlocks %}
+{% for cb in range(cbBase, cbEnd) %}
+        subgroupMatrixStore<row_major>(&tile_q,
+          (subgroup * {{ scoreColBlocks }}u + {{ cb - cbBase }}u) * 64u, mat_o{{ rb }}{{ cb }}, 8u);
+{% endfor %}
+        workgroupBarrier();
+{% for cb in range(cbBase, cbEnd) %}
+        for (var pair = 0u; pair < 2u; pair++) {
+          let index = (subgroup * {{ scoreColBlocks }}u + {{ cb - cbBase }}u) * 64u + row_in_block * 8u + col_in_block + pair;
+          tile_q[index] *= row_c[base_a + {{ rb * 8 }}u + row_in_block];
+        }
+{% endfor %}
+        workgroupBarrier();
+{% for cb in range(cbBase, cbEnd) %}
+        mat_o{{ rb }}{{ cb }} = subgroupMatrixLoad<subgroup_matrix_result<f32, 8, 8>, row_major>(
+          &tile_q, (subgroup * {{ scoreColBlocks }}u + {{ cb - cbBase }}u) * 64u, 8u);
+{% endfor %}
+        workgroupBarrier();
+{% endfor %}
+{% endfor %}
+        // The single score sweep supplied each row's new normalizer above.
+        for (var i = li; i < TILE_M * TILE_N; i += {{ sparseSgmatWorkgroup }}u) {
+          let r = i / TILE_N;
+          let key = key_base + i % TILE_N;
+          let q_abs = q_abs0 + r;
+          let allowed = r < rows_live && q_abs / SPARSE_BLOCK == mask_row && key <= q_abs;
+          prob_tile[i] = select(0.0, exp_shift(prob_tile[i], row_m[r]), allowed);
+        }
         workgroupBarrier();
         // O += P.V. P streams from shared memory; V rows are contiguous cache
         // rows, loaded directly as right-hand fragments. A masked or padded
         }
         // Orders this tile's prob_tile reads before the next tile rewrites it.
         workgroupBarrier();
       }
     }
   }
 {% endmacro %}
+@compute @workgroup_size({{ sparseSgmatWorkgroup }}, 1, 1){{ " @subgroup_size(32)" if pinSubgroupSize32 else "" }}
 fn main(
   @builtin(workgroup_id) wg: vec3<u32>,
   @builtin(local_invocation_index) li: u32,
   @builtin(subgroup_invocation_id) lane: u32
 ) {
   let tile0 = wg.x * TILE_M;
+{% if splitQueryTail %}
+  // Only a prompt delegates its partial query tile to the portable tail pass.
+  if (tile0 >= {{ sparsePrefixRows }}u && u32(total_sequence_length[0]) == params.seqLen) { return; }
+{% endif %}
   let head = wg.y % Q_HEADS;
   let batch = wg.y / Q_HEADS;
 {% endfor %}
 {% endfor %}
+  for (var r = li; r < TILE_M; r += {{ sparseSgmatWorkgroup }}u) {
     row_m[r] = -FLT_MAX;
     row_d[r] = 0.0;
   }
   workgroupBarrier();
+{{ sweep() }}
   // Normalize by the final denominators and store. Publish only as many
   // fragment columns per batch as fit prob_tile, keeping the smaller key tile's
     if (!(row_d[r] > 0.0)) {
       let key_bound = q_abs0 + r + 1u;
       let cache_base = (batch * KV_HEADS + kv_head) * MAX_CACHE_SEQ * HEAD_DIM;
+      for (var d = li; d < HEAD_DIM; d += {{ sparseSgmatWorkgroup }}u) {
         var total = 0.0;
         for (var key = 0u; key < key_bound; key++) {
           total += f32(present_value[cache_base + key * HEAD_DIM + d]);

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

@@ -9,8 +9,7 @@ fn past_sequence_length(batch: u32) -> u32 {
   let total = u32(key_total_sequence_lengths[batch]);
   return select(0u, total - params.seqLen, total >= params.seqLen);
 }
-{%- endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.SparseAttention, attention pass.
@@ -66,6 +65,7 @@ fn shifted_value(value: f32, maxValue: f32) -> f32 {
   let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
   return select(value - maxValue, 0.0, equalFiniteMax);
 }
 fn exp_shift(value: f32, maxValue: f32) -> f32 {
   return exp(shifted_value(value, maxValue));
 }
@@ -79,6 +79,10 @@ var<workgroup> probs: array<f32, Q_TILE * WG>;
 const V_STAGE_KEYS: u32 = 16u;
 var<workgroup> v_stage: array<vec4<f32>, V_STAGE_KEYS * HEAD_VEC>;
 {% endif %}
 // 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>;
@@ -90,59 +94,6 @@ var<workgroup> key_rows: array<u32, WG>;
 // Both the subgroup and portable barrier-tree engines return the same merged
 // pair to every invocation. Repeated merges require a workgroup barrier between
 // calls before their shared partial storage is reused.
-{% set combineSubgroups = combineSubgroups is defined and combineSubgroups %}
-{% if combineSubgroups %}
-// Cross-subgroup merge that assumes nothing about which invocations share a
-// subgroup or how many subgroups there are: each subgroup's elected lane
-// publishes the subgroup pair in the slot at its OWN invocation index and sets
-// that index's bit in a workgroup bitmask; thread 0 then folds exactly the
-// published slots, in ascending index order (the online (m, d) merge is not
-// float-associative, so the order is fixed), and clears the mask for the next
-// call as it reads it. Workgroup memory starts zeroed, so the mask needs no
-// setup. Same three collectives as a single-subgroup reduce, two barriers.
-var<workgroup> partialM: array<f32, WG>;
-var<workgroup> partialD: array<f32, WG>;
-var<workgroup> leaderMask: array<atomic<u32>, (WG + 31u) / 32u>;
-var<workgroup> combinedMD: vec2<f32>;
-// When the whole workgroup is one subgroup the subgroup reduce already covers
-// it (no barriers, no shared state). `subgroup_size` is the size of the current
-// subgroup and uniform, so the test is exact and may guard the barriers below.
-fn combine_partials(m: f32, d: f32, lidx: u32, sgSize: u32) -> vec2<f32> {
-  let sgM = subgroupMax(m);
-  // A lane with no elements contributes d == 0 (exact identity). A +inf
-  // element made exp(inf - inf) = NaN stick in that lane's d; a NaN element
-  // landed in d via exp(NaN); both survive the merge and are detected by the
-  // code after the reduction.
-  let sgD = subgroupAdd(d * exp_shift(m, sgM));
-  if (sgSize == WG) {
-    return vec2<f32>(sgM, sgD);
-  }
-  if (subgroupElect()) {
-    partialM[lidx] = sgM;
-    partialD[lidx] = sgD;
-    atomicOr(&leaderMask[lidx / 32u], 1u << (lidx % 32u));
-  }
-  workgroupBarrier();
-  if (lidx == 0u) {
-    var accM = -FLT_MAX;
-    var accD = 0.0;
-    for (var w = 0u; w < (WG + 31u) / 32u; w = w + 1u) {
-      var bits = atomicExchange(&leaderMask[w], 0u);
-      while (bits != 0u) {
-        let slot = w * 32u + firstTrailingBit(bits);
-        bits = bits & (bits - 1u);
-        let mNew = max(accM, partialM[slot]);
-        accD = accD * exp_shift(accM, mNew) + partialD[slot] * exp_shift(partialM[slot], mNew);
-        accM = mNew;
-      }
-    }
-    combinedMD = vec2<f32>(accM, accD);
-  }
-  workgroupBarrier();
-  return combinedMD;
-}
-{% else %}
 {% set mdStreamed = mdStreams is defined %}
 {% set mdStreams = mdStreams if mdStreams is defined else 1 %}
 {% set mdExtent = "WG" if mdStreams == 1 else "WG * " ~ mdStreams ~ "u" %}
@@ -207,21 +158,20 @@ fn combine_partials(m: f32, d: f32, lidx: u32) -> vec2<f32> {
   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;
@@ -276,7 +226,12 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
   // 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 %}
   {
@@ -364,11 +319,42 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
 {% 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.
-{% if vStageWorthIt %}
       let tileCount = min(WG, slot_count - tileBase);
 {% for j in range(qTile) %}
       var vSum_{{ j }} = vec4<f32>(0.0);

   let total = u32(key_total_sequence_lengths[batch]);
   return select(0u, total - params.seqLen, total >= params.seqLen);
 }
+{% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.SparseAttention, attention pass.
   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));
 }
 const V_STAGE_KEYS: u32 = 16u;
 var<workgroup> v_stage: array<vec4<f32>, V_STAGE_KEYS * HEAD_VEC>;
 {% endif %}
+{% if valueParts > 1 %}
+// Each lane owns one (key partition, output vec4); query streams share V loads.
+var<workgroup> value_partials: array<vec4<f32>, {{ valueParts * headVec * qTile | int }}>;
+{% endif %}
 // 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>;
 // 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 mdStreamed = mdStreams is defined %}
 {% set mdStreams = mdStreams if mdStreams is defined else 1 %}
 {% set mdExtent = "WG" if mdStreams == 1 else "WG * " ~ mdStreams ~ "u" %}
   return merged;
 }
 {% endif %}
+{% set ATTN_SCALE_DIM = "HEAD_DIM" %}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>) {
+{% if promptTail %}
+  // The matrix pass retains every query when this call includes past history.
+  if (u32(total_sequence_length[0]) != params.seqLen) { return; }
+{% endif %}
+  let tile0 = wg.x * Q_TILE{% if queryOffset > 0 %} + {{ queryOffset }}u{% endif %};
   let head = wg.y % Q_HEADS;
   let batch = wg.y / Q_HEADS;
   let tid = lid.x;
   // 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;
+{% if (seqLen - queryOffset) % qTile != 0 %}
+  // Inactive query lanes must not extend traversal past the final CSR row.
+  let row_last = (past + min(tile0 + Q_TILE, params.seqLen) - 1u) / SPARSE_BLOCK;
+{% else %}
   let row_last = mask_row_{{ qTile - 1 }};
+{% endif %}
   for (var mask_row = row_first; mask_row <= row_last; mask_row = mask_row + 1u) {
 {% else %}
   {
 {% endfor %}
       workgroupBarrier();
+      // One value vector serves every query. When head columns underfill the
+      // workgroup, spare lanes sum disjoint contiguous key ranges and publish
+      // partials; each output owner then folds them in key-range order. The
+      // manifest limits this storage and keeps a single-owner fallback.
+{% if valueParts > 1 %}
+      let tileCount = min(WG, slot_count - tileBase);
+      let valueDim = tid % HEAD_VEC;
+      let valuePart = tid / HEAD_VEC;
+      if (valuePart < {{ valueParts }}u) {
+        let first = tileCount * valuePart / {{ valueParts }}u;
+        let last = tileCount * (valuePart + 1u) / {{ valueParts }}u;
+{% for j in range(qTile) %}
+        var valueSum_{{ j }} = vec4<f32>(0.0);
+{% endfor %}
+        for (var i = first; i < last; i++) {
+          let vv = vec4<f32>(present_value[key_rows[i] / 4u + valueDim]);
+{% for j in range(qTile) %}
+          valueSum_{{ j }} += probs[{{ j }}u * WG + i] * vv;
+{% endfor %}
+        }
+{% for j in range(qTile) %}
+        value_partials[{{ j * valueParts * headVec | int }}u + tid] = valueSum_{{ j }};
+{% endfor %}
+      }
+      workgroupBarrier();
+      if (tid < HEAD_VEC) {
+{% for j in range(qTile) %}
+        var valueSum_{{ j }} = value_partials[{{ j * valueParts * headVec | int }}u + tid];
+{% for p in range(1, valueParts) %}
+        valueSum_{{ j }} += value_partials[{{ (j * valueParts + p) * headVec | int }}u + tid];
+{% endfor %}
+        running_out[{{ j }}u * HEAD_VEC + tid] = running_out[{{ j }}u * HEAD_VEC + tid] * correction_{{ j }} + valueSum_{{ j }};
+{% endfor %}
+      }
+      workgroupBarrier();
+{% elif vStageWorthIt %}
       let tileCount = min(WG, slot_count - tileBase);
 {% for j in range(qTile) %}
       var vSum_{{ j }} = vec4<f32>(0.0);

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

@@ -9,7 +9,7 @@ fn past_sequence_length(batch: u32) -> u32 {
   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
@@ -44,8 +44,12 @@ fn rotary_is_first(d: u32) -> bool {
 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 %}
 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.SparseAttention, KV append pass.
@@ -72,15 +76,11 @@ const ROTARY_DIM: u32 = {{ rotaryDim }}u;
 {{ sparse_schedule() }}
 {% if usesRotary %}
 {{ sparse_rotary(rotaryInterleaved) }}
 {% endif %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-  // 2D-folded flat index: gid.y carries the high bits past the
-  // per-axis dispatch fold width. Reduces to gid.x when the dispatch does not fold.
-  let index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
   let count = params.batchSize * KV_HEADS * params.seqLen * HEAD_DIM;
   if (index >= count) {
     return;

   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
 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 %}
+{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
+{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};{% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.SparseAttention, KV append pass.
 {{ sparse_schedule() }}
 {% if usesRotary %}
 {{ sparse_rotary(rotaryInterleaved) }}
 {% endif %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d("WG", "index", "") }}
   let count = params.batchSize * KV_HEADS * params.seqLen * HEAD_DIM;
   if (index >= count) {
     return;

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

@@ -9,7 +9,7 @@ fn past_sequence_length(batch: u32) -> u32 {
   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
@@ -44,8 +44,12 @@ fn rotary_is_first(d: u32) -> bool {
 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 %}
 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.SparseAttention, query rotary pass.
@@ -61,14 +65,10 @@ 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>) {
-  // 2D-folded flat index: gid.y carries the high bits past the
-  // per-axis dispatch fold width. Reduces to gid.x when the dispatch does not fold.
-  let index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
   let count = params.batchSize * Q_HEADS * params.seqLen * HEAD_DIM;
   if (index >= count) {
     return;

   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
 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 %}
+{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
+{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};{% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.SparseAttention, query rotary pass.
 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>) {
+{{ flat_index_2d("WG", "index", "") }}
   let count = params.batchSize * Q_HEADS * params.seqLen * HEAD_DIM;
   if (index >= count) {
     return;

build/webgpu/test.json CHANGED Viewed

The diff for this file is too large to render. See raw diff