sync c928d21e6cc1

Browse files

Files changed (6) hide show

README.md +5 -0
build/webgpu/manifest.json +293 -2
build/webgpu/metadata.json +6 -5
build/webgpu/sparse-attention-sgmat.wgsl.jinja +420 -0
build/webgpu/sparse-attention.wgsl.jinja +43 -0
build/webgpu/test.json +316 -0

README.md CHANGED Viewed

@@ -60,12 +60,17 @@ Attributes and default values (overridable per request):
 | `T` | `float32`, `float16` |
 | `M` | `int32` |
 ## Files
 - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
 - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 - [`sparse-attention.wgsl.jinja`](build/webgpu/sparse-attention.wgsl.jinja)
 - [`sparse-kv-append.wgsl.jinja`](build/webgpu/sparse-kv-append.wgsl.jinja)
 - [`sparse-q-rotary.wgsl.jinja`](build/webgpu/sparse-q-rotary.wgsl.jinja)

 | `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.
 ## Files
 - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
 - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
+- [`sparse-attention-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)
 - [`sparse-q-rotary.wgsl.jinja`](build/webgpu/sparse-q-rotary.wgsl.jinja)

build/webgpu/manifest.json CHANGED Viewed

@@ -131,8 +131,18 @@
     "sinCacheT": { "kind": "tensor", "semantic": "sin_cache", "role": "input", "required": false },
     "outputT": { "kind": "tensor", "semantic": "output", "role": "output" }
   },
-  "tunables": { "WORKGROUP_SIZE": 128, "APPEND_WORKGROUP_SIZE": 256, "NARROW_MIN_WORKGROUPS": 1024, "QUERY_TILE": 4 },
   "derive": {
     "batchSize": "dim(shapes.queryT, 0)",
     "seqLen": "dim(shapes.queryT, 1)",
     "numHeads": "attrs.num_heads",
@@ -173,7 +183,12 @@
     "packedContract": "contract and packedQkv and not useRotary",
     "packedRotaryContract": "contract and packedQkv and useRotary",
     "separateContract": "contract and not packedQkv and not useRotary",
-    "separateRotaryContract": "contract and not packedQkv and useRotary"
   },
   "constants": {
     "scalar": "dtypes.T",
@@ -556,6 +571,76 @@
         }
       }
     ],
     "attentionRotary": [
       { "name": "q_rotary", "semantic": "QRotary", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
       {
@@ -619,6 +704,70 @@
           ]
         }
       }
     ]
   },
   "variants": [
@@ -646,6 +795,36 @@
           "dispatch": { "x": "sparseQueryTiles", "y": "batchSize * numHeads" },
           "constants": { "qTile": "sparseQueryTile" }
         }
       ]
     },
     {
@@ -683,6 +862,47 @@
           "dispatch": { "x": "sparseQueryTiles", "y": "batchSize * numHeads" },
           "constants": { "qTile": "sparseQueryTile" }
         }
       ]
     },
     {
@@ -709,6 +929,36 @@
           "dispatch": { "x": "sparseQueryTiles", "y": "batchSize * numHeads" },
           "constants": { "qTile": "sparseQueryTile" }
         }
       ]
     },
     {
@@ -746,6 +996,47 @@
           "dispatch": { "x": "sparseQueryTiles", "y": "batchSize * numHeads" },
           "constants": { "qTile": "sparseQueryTile" }
         }
       ]
     }
   ]

     "sinCacheT": { "kind": "tensor", "semantic": "sin_cache", "role": "input", "required": false },
     "outputT": { "kind": "tensor", "semantic": "output", "role": "output" }
   },
+  "tunables": {
+    "WORKGROUP_SIZE": 128,
+    "APPEND_WORKGROUP_SIZE": 256,
+    "NARROW_MIN_WORKGROUPS": 1024,
+    "QUERY_TILE": 4,
+    "V_STAGE_MAX_WORKGROUPS": 512
+  },
   "derive": {
+    "wave32Adapter": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 32 and device.adapterInfo.subgroupMaxSize == 32",
+    "canPinSubgroupSize32": "device.features.has(\"subgroups\") and device.features.has(\"subgroup-size-control\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize <= 32 and device.adapterInfo.subgroupMaxSize >= 32",
+    "pinSubgroupSize32": "canPinSubgroupSize32 and not wave32Adapter",
+    "wave32Effective": "wave32Adapter or pinSubgroupSize32",
     "batchSize": "dim(shapes.queryT, 0)",
     "seqLen": "dim(shapes.queryT, 1)",
     "numHeads": "attrs.num_heads",
     "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)",
+    "sparseSgmatLdsBytes": "(64 * 32 + 64 * 64 + 64 * 2 + 128 * 2) * 4",
+    "sparseSgmatGeometryOk": "256 <= device.limits.maxComputeInvocationsPerWorkgroup and 256 <= device.limits.maxComputeWorkgroupSizeX and sgmatQueryTiles <= device.limits.maxComputeWorkgroupsPerDimension and batchSize * numHeads <= device.limits.maxComputeWorkgroupsPerDimension and sparseSgmatLdsBytes <= device.limits.maxComputeWorkgroupStorageSize",
+    "sparseSgmatOk": "tensorDtypes.queryT == \"float32\" and seqLen >= 64 and sparseBlockSize % 64 == 0 and headSize % 32 == 0 and headSize <= 128 and maxCacheSeq % 64 == 0 and device.features.has(\"subgroups\") and wave32Effective and device.features.has(\"chromium-experimental-subgroup-matrix\") and sparseSgmatGeometryOk"
   },
   "constants": {
     "scalar": "dtypes.T",
         }
       }
     ],
+    "attentionSgmat": [
+      {
+        "name": "query",
+        "arg": "queryT",
+        "semantic": "query",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$scalar"
+      },
+      {
+        "name": "present_key",
+        "arg": "pastKeyT",
+        "semantic": "past_key",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$scalar"
+      },
+      {
+        "name": "present_value",
+        "arg": "pastValueT",
+        "semantic": "past_value",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$scalar"
+      },
+      {
+        "name": "block_row_indices",
+        "arg": "blockRowIndicesT",
+        "semantic": "block_row_indices",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "i32"
+      },
+      {
+        "name": "block_col_indices",
+        "arg": "blockColIndicesT",
+        "semantic": "block_col_indices",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "i32"
+      },
+      {
+        "name": "key_total_sequence_lengths",
+        "arg": "keyTotalSequenceLengthsT",
+        "semantic": "key_total_sequence_lengths",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "i32"
+      },
+      {
+        "name": "total_sequence_length",
+        "arg": "totalSequenceLengthT",
+        "semantic": "total_sequence_length",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "i32"
+      },
+      {
+        "name": "output",
+        "arg": "outputT",
+        "semantic": "output",
+        "buffer": { "type": "storage" },
+        "elementType": "$scalar"
+      },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "seqLen", "type": "u32", "value": "seqLen" },
+            { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" }
+          ]
+        }
+      }
+    ],
     "attentionRotary": [
       { "name": "q_rotary", "semantic": "QRotary", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
       {
           ]
         }
       }
+    ],
+    "attentionSgmatRotary": [
+      { "name": "q_rotary", "semantic": "QRotary", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
+      {
+        "name": "present_key",
+        "arg": "pastKeyT",
+        "semantic": "past_key",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$scalar"
+      },
+      {
+        "name": "present_value",
+        "arg": "pastValueT",
+        "semantic": "past_value",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$scalar"
+      },
+      {
+        "name": "block_row_indices",
+        "arg": "blockRowIndicesT",
+        "semantic": "block_row_indices",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "i32"
+      },
+      {
+        "name": "block_col_indices",
+        "arg": "blockColIndicesT",
+        "semantic": "block_col_indices",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "i32"
+      },
+      {
+        "name": "key_total_sequence_lengths",
+        "arg": "keyTotalSequenceLengthsT",
+        "semantic": "key_total_sequence_lengths",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "i32"
+      },
+      {
+        "name": "total_sequence_length",
+        "arg": "totalSequenceLengthT",
+        "semantic": "total_sequence_length",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "i32"
+      },
+      {
+        "name": "output",
+        "arg": "outputT",
+        "semantic": "output",
+        "buffer": { "type": "storage" },
+        "elementType": "$scalar"
+      },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "seqLen", "type": "u32", "value": "seqLen" },
+            { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" }
+          ]
+        }
+      }
     ]
   },
   "variants": [
           "dispatch": { "x": "sparseQueryTiles", "y": "batchSize * numHeads" },
           "constants": { "qTile": "sparseQueryTile" }
         }
+      ],
+      "constants": { "vStageWorthIt": "sparseVStageWorthIt" }
+    },
+    {
+      "id": "separate_sgmat",
+      "priority": 20,
+      "when": ["separateContract", "geometryOk", "sparseSgmatOk"],
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
+      },
+      "passes": [
+        {
+          "id": "append",
+          "name": "SparseAttention.Append",
+          "shader": "sparse-kv-append.wgsl.jinja",
+          "bindings": "appendSeparate",
+          "constants": { "kvSource": "\"new_key\"", "vSource": "\"new_value\"" },
+          "dispatch": {
+            "threads": "batchSize * kvNumHeads * seqLen * headSize",
+            "workgroupSize": "constants.appendWorkgroupSize"
+          }
+        },
+        {
+          "id": "attention",
+          "name": "SparseAttention.AttentionSgmat",
+          "shader": "sparse-attention-sgmat.wgsl.jinja",
+          "bindings": "attentionSgmat",
+          "dispatch": { "x": "sgmatQueryTiles", "y": "batchSize * numHeads" }
+        }
       ]
     },
     {
           "dispatch": { "x": "sparseQueryTiles", "y": "batchSize * numHeads" },
           "constants": { "qTile": "sparseQueryTile" }
         }
+      ],
+      "constants": { "vStageWorthIt": "sparseVStageWorthIt" }
+    },
+    {
+      "id": "separate_rotary_sgmat",
+      "priority": 30,
+      "when": ["separateRotaryContract", "geometryOk", "sparseSgmatOk"],
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
+      },
+      "intermediates": [{ "id": "QRotary", "dtype": "float32", "shape": "[qRotaryElements]" }],
+      "passes": [
+        {
+          "id": "append",
+          "name": "SparseAttention.Append",
+          "shader": "sparse-kv-append.wgsl.jinja",
+          "bindings": "appendSeparateRotary",
+          "constants": { "kvSource": "\"new_key\"", "vSource": "\"new_value\"" },
+          "dispatch": {
+            "threads": "batchSize * kvNumHeads * seqLen * headSize",
+            "workgroupSize": "constants.appendWorkgroupSize"
+          }
+        },
+        {
+          "id": "qrotary",
+          "name": "SparseAttention.QueryRotary",
+          "shader": "sparse-q-rotary.wgsl.jinja",
+          "bindings": "qRotary",
+          "dispatch": {
+            "threads": "batchSize * numHeads * seqLen * headSize",
+            "workgroupSize": "constants.appendWorkgroupSize"
+          }
+        },
+        {
+          "id": "attention",
+          "name": "SparseAttention.AttentionSgmat",
+          "shader": "sparse-attention-sgmat.wgsl.jinja",
+          "bindings": "attentionSgmatRotary",
+          "dispatch": { "x": "sgmatQueryTiles", "y": "batchSize * numHeads" }
+        }
       ]
     },
     {
           "dispatch": { "x": "sparseQueryTiles", "y": "batchSize * numHeads" },
           "constants": { "qTile": "sparseQueryTile" }
         }
+      ],
+      "constants": { "vStageWorthIt": "sparseVStageWorthIt" }
+    },
+    {
+      "id": "packed_sgmat",
+      "priority": 20,
+      "when": ["packedContract", "geometryOk", "sparseSgmatOk"],
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
+      },
+      "passes": [
+        {
+          "id": "append",
+          "name": "SparseAttention.Append",
+          "shader": "sparse-kv-append.wgsl.jinja",
+          "bindings": "appendPacked",
+          "constants": { "kvSource": "\"packed_qkv\"", "vSource": "\"packed_qkv\"" },
+          "dispatch": {
+            "threads": "batchSize * kvNumHeads * seqLen * headSize",
+            "workgroupSize": "constants.appendWorkgroupSize"
+          }
+        },
+        {
+          "id": "attention",
+          "name": "SparseAttention.AttentionSgmat",
+          "shader": "sparse-attention-sgmat.wgsl.jinja",
+          "bindings": "attentionSgmat",
+          "dispatch": { "x": "sgmatQueryTiles", "y": "batchSize * numHeads" }
+        }
       ]
     },
     {
           "dispatch": { "x": "sparseQueryTiles", "y": "batchSize * numHeads" },
           "constants": { "qTile": "sparseQueryTile" }
         }
+      ],
+      "constants": { "vStageWorthIt": "sparseVStageWorthIt" }
+    },
+    {
+      "id": "packed_rotary_sgmat",
+      "priority": 30,
+      "when": ["packedRotaryContract", "geometryOk", "sparseSgmatOk"],
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
+      },
+      "intermediates": [{ "id": "QRotary", "dtype": "float32", "shape": "[qRotaryElements]" }],
+      "passes": [
+        {
+          "id": "append",
+          "name": "SparseAttention.Append",
+          "shader": "sparse-kv-append.wgsl.jinja",
+          "bindings": "appendPackedRotary",
+          "constants": { "kvSource": "\"packed_qkv\"", "vSource": "\"packed_qkv\"" },
+          "dispatch": {
+            "threads": "batchSize * kvNumHeads * seqLen * headSize",
+            "workgroupSize": "constants.appendWorkgroupSize"
+          }
+        },
+        {
+          "id": "qrotary",
+          "name": "SparseAttention.QueryRotary",
+          "shader": "sparse-q-rotary.wgsl.jinja",
+          "bindings": "qRotary",
+          "dispatch": {
+            "threads": "batchSize * numHeads * seqLen * headSize",
+            "workgroupSize": "constants.appendWorkgroupSize"
+          }
+        },
+        {
+          "id": "attention",
+          "name": "SparseAttention.AttentionSgmat",
+          "shader": "sparse-attention-sgmat.wgsl.jinja",
+          "bindings": "attentionSgmatRotary",
+          "dispatch": { "x": "sgmatQueryTiles", "y": "batchSize * numHeads" }
+        }
       ]
     }
   ]

build/webgpu/metadata.json CHANGED Viewed

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

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

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

	@@ -0,0 +1,420 @@

+{% macro sparse_schedule() %}
+// How much history precedes this call's tokens. Prompt mode is determined from
+// scalar total_sequence_length rather than per-batch values and forces the past
+// length to zero; a padded prompt row must not be read as if it had history.
+fn past_sequence_length(batch: u32) -> u32 {
+  if (u32(total_sequence_length[0]) == params.seqLen) {
+    return 0u;
+  }
+  let total = u32(key_total_sequence_lengths[batch]);
+  return select(0u, total - params.seqLen, total >= params.seqLen);
+}
+{%- endmacro %}
+enable subgroups;
+{% if pinSubgroupSize32 %}
+enable subgroup_size_control;
+{% endif %}
+enable chromium_experimental_subgroup_matrix;
+diagnostic(off, chromium.subgroup_matrix_uniformity);
+{{ env.wgsl.resourceDeclarations }}
+// com.microsoft.SparseAttention, subgroup-matrix attention pass (prefill band).
+// One workgroup per (batch, 64-query tile, query head). The block-sparse layout
+// already lists dense SPARSE_BLOCK-wide key tiles, so each selected block is
+// walked as 64x64 score tiles on 8x8 subgroup matrices with no gather step: the
+// K and V rows of a block are contiguous rows of the head-major cache and load
+// directly as matrix fragments. Only Q goes through shared memory, for the
+// query-tail zero padding.
+//
+// The pass runs the selected blocks twice. Sweep one folds each score tile
+// into running per-row (max, denominator) softmax statistics; sweep two
+// recomputes the same tiles (bit-identical: same operands, same order),
+// exponentiates against the now-final row max, and accumulates P.V into
+// subgroup-matrix result fragments. Recomputing QK costs half the ideal FLOPs
+// again, but it is what lets both GEMMs stay on the matrix unit: a running
+// online rescale of the output fragments would need an elementwise multiply on
+// opaque fragments, which only exists as a shared-memory round trip per key
+// tile, and the output accumulator alone (64 x head_size f32) is the whole
+// workgroup storage budget.
+//
+// Causality, CSR duplicate suppression, the dense-row shortcut, and the
+// all-masked-row uniform average reproduce the online kernel's semantics; a
+// key past a query's bound contributes probability zero and multiplies away
+// against whatever the cache holds there, exactly as the online kernel's
+// value walk does.
+const Q_HEADS: u32 = {{ numHeads }}u;
+const KV_HEADS: u32 = {{ kvNumHeads }}u;
+const HEAD_DIM: u32 = {{ headSize }}u;
+const MAX_CACHE_SEQ: u32 = {{ maxCacheSeq }}u;
+const SPARSE_BLOCK: u32 = {{ sparseBlockSize }}u;
+const NUM_LAYOUT: u32 = {{ numLayout }}u;
+const ROW_STRIDE: u32 = {{ maxBlocks + 1 }}u;
+const COL_STRIDE: u32 = {{ maxNnz }}u;
+const Q_HIDDEN: u32 = {{ numHeads * headSize }}u;
+{% if not usesRotary %}
+const Q_STRIDE: u32 = {{ packedStride if packedQkv else numHeads * headSize }}u;
+{% endif %}
+// 64 query rows x 64 key columns per score tile, staged over the head
+// dimension in 32-wide chunks; 8 subgroups of 32 lanes as a 4x2 grid, each
+// owning a 16x32 score sub-tile. sparse_block_size here is a multiple of 64
+// and head_size a multiple of 32, so every key tile is a whole aligned run of
+// cache rows and the head loop has no tail.
+const TILE_M: u32 = 64u;
+const TILE_N: u32 = 64u;
+const TILE_K: u32 = 32u;
+const SUB_TILES: u32 = {{ (sparseBlockSize / 64) | int }}u;
+{% set pvColBlocks = (headSize / 16) | int %}
+// Each of the two subgroup column groups owns half the output head dimension.
+const PV_COL_SPAN: u32 = {{ (headSize / 2) | int }}u;
+// FLT_MAX, not -inf, as the online (m, d) accumulator init: merges must keep
+// `m - m` finite so an empty lane / all--inf row contributes the exact
+// accumulator identity (m, d) = (-FLT_MAX, 0). Operator epilogues interpret
+// a zero final denominator according to their public semantics. Using -inf
+// here changes +inf-row behavior.
+const FLT_MAX: f32 = 3.4028234663852886e38;
+fn is_finite_f32(value: f32) -> bool {
+  return select(false, value <= FLT_MAX, value >= -FLT_MAX);
+}
+// x - m that is exactly 0 when x equals a finite m, so exp(shifted) == 1
+// exactly at the row max. `x - x` on an infinite max is a legal fast-math
+// fold to 0, which would silently turn +inf rows finite — the explicit
+// equality test keeps the NaN propagation of the serial kernels.
+fn shifted_value(value: f32, maxValue: f32) -> f32 {
+  let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
+  return select(value - maxValue, 0.0, equalFiniteMax);
+}
+fn exp_shift(value: f32, maxValue: f32) -> f32 {
+  return exp(shifted_value(value, maxValue));
+}
+// Q staging for the score GEMM; the score epilogues alias it as the
+// fragment-store scratch (8 subgroups x 4 banks x 64 elements).
+var<workgroup> tile_q: array<f32, 2048>;
+// Tile probabilities for the P.V GEMM; the output epilogue aliases it as the
+// result-fragment scratch once the last key tile's readers are done.
+var<workgroup> prob_tile: array<f32, 4096>;
+var<workgroup> row_m: array<f32, 64>;
+var<workgroup> row_d: array<f32, 64>;
+// Per-key-tile row partials, one slot per (row, subgroup column group).
+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() }}
+{% macro score_tile() %}
+        // S = Q.K^T for this key tile, accumulated over the head dimension in
+        // TILE_K chunks. Q is staged with query tails zero-padded; K loads
+        // straight from the cache as column-major right fragments.
+        for (var k_base = 0u; k_base < HEAD_DIM; k_base += TILE_K) {
+          {
+            let a_row = li / 4u;
+            let a_col = (li % 4u) * 8u;
+            for (var i = 0u; i < 8u; i++) {
+              let k = k_base + a_col + i;
+{% if usesRotary %}
+              let q_index = ((batch * Q_HEADS + head) * params.seqLen
+                + min(tile0 + a_row, params.seqLen - 1u)) * HEAD_DIM + k;
+              tile_q[a_row * TILE_K + a_col + i] =
+                select(0.0, q_rotary[q_index], a_row < rows_live);
+{% else %}
+              let q_index = (batch * params.seqLen + min(tile0 + a_row, params.seqLen - 1u))
+                * Q_STRIDE + head * HEAD_DIM + k;
+              tile_q[a_row * TILE_K + a_col + i] =
+                select(0.0, f32(query[q_index]), a_row < rows_live);
+{% endif %}
+            }
+          }
+          workgroupBarrier();
+          for (var step = 0u; step < TILE_K; step += 8u) {
+{% for rb in range(2) %}
+            let mat_a{{ rb }} = subgroupMatrixLoad<subgroup_matrix_left<f32, 8, 8>, row_major>(
+              &tile_q, (base_a + {{ rb * 8 }}u) * TILE_K + step, TILE_K
+            );
+{% endfor %}
+{% for cb in range(4) %}
+            let mat_b{{ cb }} = subgroupMatrixLoad<subgroup_matrix_right<f32, 8, 8>, col_major>(
+              &present_key, key_row_elem + (base_b + {{ cb * 8 }}u) * HEAD_DIM + k_base + step, HEAD_DIM
+            );
+{% endfor %}
+{% for rb in range(2) %}
+{% for cb in range(4) %}
+            mat_s{{ rb }}{{ cb }} = subgroupMatrixMultiplyAccumulate(
+              mat_a{{ rb }}, mat_b{{ cb }}, mat_s{{ rb }}{{ cb }}
+            );
+{% endfor %}
+{% endfor %}
+          }
+          workgroupBarrier();
+        }
+{% endmacro %}
+{% 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.
+  for (var mask_row = row_first; mask_row <= row_last; mask_row++) {
+    let start = u32(block_row_indices[row_base + mask_row]);
+    let end = u32(block_row_indices[row_base + mask_row + 1u]);
+    // A row listing every causal block is also dense, so its column values are
+    // ignored, matching the online kernel and the pinned dense-layout fixture.
+    let dense_row = layout_is_dense || (end - start == mask_row + 1u);
+    // Furthest causal bound (exclusive) of the tile's queries on this row.
+    let bound_max = min(q_abs0 + rows_live, (mask_row + 1u) * SPARSE_BLOCK);
+    let tile_count = select(
+      (end - start) * SUB_TILES, (bound_max + TILE_N - 1u) / TILE_N, dense_row
+    );
+    for (var t = 0u; t < tile_count; t++) {
+      var key_base = t * TILE_N;
+      var tile_ok = true;
+      if (!dense_row) {
+        let entry = t / SUB_TILES;
+        let block = u32(block_col_indices[layout_id * COL_STRIDE + start + entry]);
+        // CSR is a boolean mask: a repeated column must not repeat its
+        // probability mass, so a duplicate entry's tiles are skipped whole.
+        var duplicate = false;
+        for (var previous = 0u; previous < entry; previous++) {
+          duplicate = duplicate ||
+            u32(block_col_indices[layout_id * COL_STRIDE + start + previous]) == block;
+        }
+        key_base = block * SPARSE_BLOCK + (t % SUB_TILES) * TILE_N;
+        tile_ok = !duplicate;
+      }
+      // A tile at or past every member query's bound is fully masked; skipping
+      // it is exact. All of this tile's control state is derived from
+      // workgroup-uniform reads, so the barriers below stay in uniform flow.
+      if (tile_ok && key_base < bound_max) {
+        let key_row_elem = ((batch * KV_HEADS + kv_head) * MAX_CACHE_SEQ + key_base) * HEAD_DIM;
+{% for rb in range(2) %}
+{% for cb in range(4) %}
+        var mat_s{{ rb }}{{ cb }}: subgroup_matrix_result<f32, 8, 8>;
+{% endfor %}
+{% endfor %}
+{{ 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(4) %}
+        subgroupMatrixStore<row_major>(
+          &tile_q, (subgroup * 4u + {{ cb }}u) * 64u, mat_s{{ rb }}{{ cb }}, 8u
+        );
+{% endfor %}
+        workgroupBarrier();
+{% for cb in range(4) %}
+        for (var pair = 0u; pair < 2u; pair++) {
+          let r = base_a + {{ rb * 8 }}u + row_in_block;
+          let key = key_base + base_b + {{ cb * 8 }}u + col_in_block + pair;
+          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 * 4u + {{ cb }}u) * 64u + row_in_block * 8u + col_in_block + pair
+            ] * scale;
+            let prev_m = tile_stat_m{{ rb }};
+            tile_stat_m{{ rb }} = max(tile_stat_m{{ rb }}, scored);
+            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 * 4u + {{ cb }}u) * 64u + row_in_block * 8u + col_in_block + pair
+            ] * scale, row_m[r]);
+          }
+          prob_tile[r * TILE_N + base_b + {{ cb * 8 }}u + col_in_block + pair] = prob;
+{% 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
+        // uniform.
+{% for xor in [1, 2] %}
+        {
+          let other_m = subgroupShuffleXor(tile_stat_m{{ rb }}, {{ xor }}u);
+          let other_d = subgroupShuffleXor(tile_stat_d{{ rb }}, {{ xor }}u);
+          let merged_m = max(tile_stat_m{{ rb }}, other_m);
+          tile_stat_d{{ rb }} = tile_stat_d{{ rb }} * exp_shift(tile_stat_m{{ rb }}, merged_m)
+            + other_d * exp_shift(other_m, merged_m);
+          tile_stat_m{{ rb }} = merged_m;
+        }
+{% endfor %}
+        if (lane % 4u == 0u) {
+          let stat_row = base_a + {{ rb * 8 }}u + row_in_block;
+          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
+        // partial (-FLT_MAX, 0) is the identity, a NaN partial poisons the
+        // denominator the way a NaN score does.
+        if (li < TILE_M) {
+          var merged_m = row_m[li];
+          var merged_d = row_d[li];
+          for (var g = 0u; g < 2u; g++) {
+            let m2 = part_m[li * 2u + g];
+            let d2 = part_d[li * 2u + g];
+            let new_m = max(merged_m, m2);
+            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
+        // slot carries probability zero and multiplies away.
+        for (var kk = 0u; kk < TILE_N; kk += 8u) {
+{% for rb in range(2) %}
+          let mat_p{{ rb }} = subgroupMatrixLoad<subgroup_matrix_left<f32, 8, 8>, row_major>(
+            &prob_tile, (base_a + {{ rb * 8 }}u) * TILE_N + kk, TILE_N
+          );
+{% endfor %}
+{% for cb in range(pvColBlocks) %}
+          let mat_v{{ cb }} = subgroupMatrixLoad<subgroup_matrix_right<f32, 8, 8>, row_major>(
+            &present_value,
+            key_row_elem + kk * HEAD_DIM + subtile_idx * PV_COL_SPAN + {{ cb * 8 }}u, HEAD_DIM
+          );
+{% endfor %}
+{% for rb in range(2) %}
+{% for cb in range(pvColBlocks) %}
+          mat_o{{ rb }}{{ cb }} = subgroupMatrixMultiplyAccumulate(
+            mat_p{{ rb }}, mat_v{{ cb }}, mat_o{{ rb }}{{ cb }}
+          );
+{% endfor %}
+{% endfor %}
+        }
+        // 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;
+  let past = past_sequence_length(batch);
+  let kv_head = head / (Q_HEADS / KV_HEADS);
+  let scale = scale_value();
+  // Layouts cycle over the heads: four layouts and eight heads give S0 S1 S2 S3 S0 ...
+  let layout_id = head % NUM_LAYOUT;
+  let row_base = layout_id * ROW_STRIDE;
+  // A layout at or above the triangular dense-block count is treated as dense.
+  let dense_nnz = ROW_STRIDE * (ROW_STRIDE - 1u) / 2u;
+  let layout_is_dense = u32(block_row_indices[row_base + ROW_STRIDE - 1u]) >= dense_nnz;
+  let subgroup = li / 32u;
+  let subtile_idy = subgroup % 4u;
+  let subtile_idx = subgroup / 4u;
+  let base_a = subtile_idy * 16u;
+  let base_b = subtile_idx * 32u;
+  let row_in_block = lane / 4u;
+  let col_in_block = (lane % 4u) * 2u;
+  // The dispatch never launches a tile past the last token, so at least one
+  // row of every tile is live.
+  let rows_live = min(TILE_M, params.seqLen - tile0);
+  let q_abs0 = past + tile0;
+  let row_first = q_abs0 / SPARSE_BLOCK;
+  let row_last = (q_abs0 + rows_live - 1u) / SPARSE_BLOCK;
+{% for rb in range(2) %}
+{% for cb in range(pvColBlocks) %}
+  var mat_o{{ rb }}{{ cb }}: subgroup_matrix_result<f32, 8, 8>;
+{% 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; the result fragments
+  // round-trip through prob_tile, which no reader needs any more.
+{% for rb in range(2) %}
+{% if rb > 0 %}
+  workgroupBarrier();
+{% endif %}
+{% for cb in range(pvColBlocks) %}
+  subgroupMatrixStore<row_major>(
+    &prob_tile, (subgroup * {{ pvColBlocks }}u + {{ cb }}u) * 64u, mat_o{{ rb }}{{ cb }}, 8u
+  );
+{% endfor %}
+  workgroupBarrier();
+{% for cb in range(pvColBlocks) %}
+  for (var pair = 0u; pair < 2u; pair++) {
+    let r = base_a + {{ rb * 8 }}u + row_in_block;
+    if (r < rows_live && row_d[r] > 0.0) {
+      let dcol = subtile_idx * PV_COL_SPAN + {{ cb * 8 }}u + col_in_block + pair;
+      let out_index = (batch * params.seqLen + tile0 + r) * Q_HIDDEN + head * HEAD_DIM + dcol;
+      output[out_index] = {{ scalar }}(prob_tile[
+        (subgroup * {{ pvColBlocks }}u + {{ cb }}u) * 64u + row_in_block * 8u + col_in_block + pair
+      ] / row_d[r]);
+    }
+  }
+{% endfor %}
+{% endfor %}
+  // An all-masked row becomes a uniform average over causal history rather
+  // than zero — the online kernel's public semantics, NaN denominators
+  // included.
+  for (var r = 0u; r < rows_live; r++) {
+    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 out_index = (batch * params.seqLen + tile0 + r) * Q_HIDDEN + head * HEAD_DIM + d;
+        output[out_index] = {{ scalar }}(total / f32(key_bound));
+      }
+    }
+  }
+}

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

@@ -79,6 +79,16 @@ fn exp_shift(value: f32, maxValue: f32) -> f32 {
 var<workgroup> q_shared: array<vec4<f32>, Q_TILE * HEAD_VEC>;
 var<workgroup> running_out: array<vec4<f32>, Q_TILE * HEAD_VEC>;
 var<workgroup> probs: array<f32, Q_TILE * WG>;
 // One resolved cache row base per key of the current tile, so the value accumulation
 // re-reads a base instead of re-walking the column list per head dimension.
 var<workgroup> key_rows: array<u32, WG>;
@@ -374,6 +384,38 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
       // so this rescale-and-accumulate needs no further synchronization. One value vector
       // serves every query, which is the other half of the tile's reuse; a key outside a
       // query's causal bound carries prob 0 and is multiplied away.
       let tileCount = min(WG, slot_count - tileBase);
       for (var dv: u32 = tid; dv < HEAD_VEC; dv = dv + WG) {
 {% for j in range(qTile) %}
@@ -391,6 +433,7 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
       }
       // Orders this tile's probs and key_rows reads before the next tile writes them.
       workgroupBarrier();
       tileBase = tileBase + WG;
     }

 var<workgroup> q_shared: array<vec4<f32>, Q_TILE * HEAD_VEC>;
 var<workgroup> running_out: array<vec4<f32>, Q_TILE * HEAD_VEC>;
 var<workgroup> probs: array<f32, Q_TILE * WG>;
+{% if vStageWorthIt %}
+// Staged value rows for the running-output walk (small-dispatch shapes only):
+// the walk itself occupies just HEAD_VEC lanes, so the whole workgroup
+// prefetches each chunk of value rows at full width and the walkers read
+// shared memory instead of chaining dependent global loads. Large dispatches
+// already hide that latency across resident workgroups, and the chunk
+// barriers only cost them, so they keep the direct walk.
+const V_STAGE_KEYS: u32 = 16u;
+var<workgroup> v_stage: array<vec4<f32>, V_STAGE_KEYS * HEAD_VEC>;
+{% endif %}
 // 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>;
       // 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);
+{% endfor %}
+      for (var chunkBase: u32 = 0u; chunkBase < tileCount; chunkBase = chunkBase + V_STAGE_KEYS) {
+        let chunkCount = min(V_STAGE_KEYS, tileCount - chunkBase);
+        for (var slot = tid; slot < chunkCount * HEAD_VEC; slot = slot + WG) {
+          let ki = slot / HEAD_VEC;
+          let dvs = slot % HEAD_VEC;
+          v_stage[slot] = vec4<f32>(present_value[key_rows[chunkBase + ki] / 4u + dvs]);
+        }
+        workgroupBarrier();
+        for (var dv: u32 = tid; dv < HEAD_VEC; dv = dv + WG) {
+          for (var i: u32 = 0u; i < chunkCount; i = i + 1u) {
+            let vv = v_stage[i * HEAD_VEC + dv];
+{% for j in range(qTile) %}
+            vSum_{{ j }} = vSum_{{ j }} + probs[{{ j }}u * WG + chunkBase + i] * vv;
+{% endfor %}
+          }
+        }
+        // Orders this chunk's v_stage reads before the next chunk overwrites it.
+        workgroupBarrier();
+      }
+      for (var dv: u32 = tid; dv < HEAD_VEC; dv = dv + WG) {
+{% for j in range(qTile) %}
+        running_out[{{ j }}u * HEAD_VEC + dv] = running_out[{{ j }}u * HEAD_VEC + dv] * correction_{{ j }} + vSum_{{ j }};
+{% endfor %}
+      }
+      // Orders this tile's probs and key_rows reads before the next tile writes them.
+      workgroupBarrier();
+{% else %}
       let tileCount = min(WG, slot_count - tileBase);
       for (var dv: u32 = tid; dv < HEAD_VEC; dv = dv + WG) {
 {% for j in range(qTile) %}
       }
       // Orders this tile's probs and key_rows reads before the next tile writes them.
       workgroupBarrier();
+{% endif %}
       tileBase = tileBase + WG;
     }

build/webgpu/test.json CHANGED Viewed

@@ -740,6 +740,322 @@
         "pastKeyT": { "dtype": "float32", "shape": [1, 1, 16, 16], "tolerance": 0.000001, "relTolerance": 0.000001 },
         "pastValueT": { "dtype": "float32", "shape": [1, 1, 16, 16], "tolerance": 0.000001, "relTolerance": 0.000001 }
       }
     }
   ]
 }

         "pastKeyT": { "dtype": "float32", "shape": [1, 1, 16, 16], "tolerance": 0.000001, "relTolerance": 0.000001 },
         "pastValueT": { "dtype": "float32", "shape": [1, 1, 16, 16], "tolerance": 0.000001, "relTolerance": 0.000001 }
       }
+    },
+    {
+      "name": "sgmat_prompt_separate_blk64_duplicate_csr",
+      "provenance": {
+        "notes": "Prefill shape on the subgroup-matrix route: 64-token blocks, a duplicate CSR column whose repeated block must not repeat its mass, a dense-by-count row, and an off-diagonal-only row."
+      },
+      "attrs": { "num_heads": 4, "kv_num_heads": 2, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 192, 256],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 192, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 192, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 2, 192, 64], "data": { "kind": "constant", "value": 0.0 } },
+        "pastValueT": { "dtype": "float32", "shape": [1, 2, 192, 64], "data": { "kind": "constant", "value": 0.0 } },
+        "blockRowIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 4],
+          "data": { "kind": "values", "values": [0, 1, 3, 5, 0, 1, 2, 3] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [0, 0, 1, 1, 1, 0, 1, 2, 0, 0] }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [192] } },
+        "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [192] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 192, 256], "tolerance": 0.0001, "relTolerance": 0.0005 },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 2, 192, 64], "tolerance": 0.00002, "relTolerance": 0.0002 },
+        "pastValueT": { "dtype": "float32", "shape": [1, 2, 192, 64], "tolerance": 0.00002, "relTolerance": 0.0002 }
+      }
+    },
+    {
+      "name": "sgmat_prompt_packed_blk64",
+      "provenance": {
+        "notes": "Packed [Q|K|V] prefill on the subgroup-matrix route; layout 0 crosses the triangular dense-count threshold so its column values are ignored, layout 1 keeps a gapped sparse row."
+      },
+      "attrs": { "num_heads": 4, "kv_num_heads": 2, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 128, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 2, 128, 64], "data": { "kind": "constant", "value": 0.0 } },
+        "pastValueT": { "dtype": "float32", "shape": [1, 2, 128, 64], "data": { "kind": "constant", "value": 0.0 } },
+        "blockRowIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 3],
+          "data": { "kind": "values", "values": [0, 1, 3, 0, 1, 2] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 3],
+          "data": { "kind": "values", "values": [0, 0, 1, 0, 1, 0] }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [128] } },
+        "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [128] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 128, 256], "tolerance": 0.0001, "relTolerance": 0.0005 },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 2, 128, 64], "tolerance": 0.00002, "relTolerance": 0.0002 },
+        "pastValueT": { "dtype": "float32", "shape": [1, 2, 128, 64], "tolerance": 0.00002, "relTolerance": 0.0002 }
+      }
+    },
+    {
+      "name": "sgmat_prompt_separate_rotary_neox_blk64",
+      "provenance": {
+        "notes": "Grouped-query prefill with partial NeoX rotary on the subgroup-matrix route; the attention pass reads the rotated-Q intermediate."
+      },
+      "attrs": { "num_heads": 2, "kv_num_heads": 1, "sparse_block_size": 64, "do_rotary": 1, "rotary_interleaved": 0 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 128, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 128, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 1, 128, 64], "data": { "kind": "constant", "value": 0.0 } },
+        "pastValueT": { "dtype": "float32", "shape": [1, 1, 128, 64], "data": { "kind": "constant", "value": 0.0 } },
+        "blockRowIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 3],
+          "data": { "kind": "values", "values": [0, 1, 3, 0, 1, 2] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 3],
+          "data": { "kind": "values", "values": [0, 0, 1, 0, 1, 0] }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [128] } },
+        "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [128] } },
+        "cosCacheT": { "dtype": "float32", "shape": [128, 16], "data": { "kind": "constant", "value": 0.8 } },
+        "sinCacheT": { "dtype": "float32", "shape": [128, 16], "data": { "kind": "constant", "value": 0.6 } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 128, 128], "tolerance": 0.0001, "relTolerance": 0.0005 },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 1, 128, 64], "tolerance": 0.00002, "relTolerance": 0.0002 },
+        "pastValueT": { "dtype": "float32", "shape": [1, 1, 128, 64], "tolerance": 0.00002, "relTolerance": 0.0002 }
+      }
+    },
+    {
+      "name": "sgmat_prompt_packed_rotary_interleaved_blk64",
+      "provenance": {
+        "notes": "Packed [Q|K|V] prefill with full interleaved rotary at head_size 32, the narrowest head the subgroup-matrix route admits (a single reduction chunk)."
+      },
+      "attrs": { "num_heads": 2, "kv_num_heads": 2, "sparse_block_size": 64, "do_rotary": 1, "rotary_interleaved": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 128, 192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 2, 128, 32], "data": { "kind": "constant", "value": 0.0 } },
+        "pastValueT": { "dtype": "float32", "shape": [1, 2, 128, 32], "data": { "kind": "constant", "value": 0.0 } },
+        "blockRowIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 3],
+          "data": { "kind": "values", "values": [0, 1, 3, 0, 1, 2] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 3],
+          "data": { "kind": "values", "values": [0, 0, 1, 0, 1, 0] }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [128] } },
+        "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [128] } },
+        "cosCacheT": { "dtype": "float32", "shape": [128, 16], "data": { "kind": "constant", "value": 0.8 } },
+        "sinCacheT": { "dtype": "float32", "shape": [128, 16], "data": { "kind": "constant", "value": 0.6 } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 128, 64], "tolerance": 0.0001, "relTolerance": 0.0005 },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 2, 128, 32], "tolerance": 0.00002, "relTolerance": 0.0002 },
+        "pastValueT": { "dtype": "float32", "shape": [1, 2, 128, 32], "tolerance": 0.00002, "relTolerance": 0.0002 }
+      }
+    },
+    {
+      "name": "sgmat_prompt_tail_s100_explicit_scale",
+      "provenance": {
+        "notes": "A 100-token prefill leaves the second query tile 36 rows live, exercising the staged-Q zero padding and guarded stores; the explicit scale attribute overrides the head-size default."
+      },
+      "attrs": { "num_heads": 2, "kv_num_heads": 1, "sparse_block_size": 64, "scale": 0.25 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 100, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 100, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 100, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 1, 128, 32], "data": { "kind": "constant", "value": 0.0 } },
+        "pastValueT": { "dtype": "float32", "shape": [1, 1, 128, 32], "data": { "kind": "constant", "value": 0.0 } },
+        "blockRowIndicesT": { "dtype": "int32", "shape": [1, 3], "data": { "kind": "values", "values": [0, 1, 3] } },
+        "blockColIndicesT": { "dtype": "int32", "shape": [1, 3], "data": { "kind": "values", "values": [0, 0, 1] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [100] } },
+        "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [100] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 100, 64], "tolerance": 0.0001, "relTolerance": 0.0005 },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 1, 128, 32], "tolerance": 0.00002, "relTolerance": 0.0002 },
+        "pastValueT": { "dtype": "float32", "shape": [1, 1, 128, 32], "tolerance": 0.00002, "relTolerance": 0.0002 }
+      }
+    },
+    {
+      "name": "sgmat_prompt_blk128_subtiles",
+      "provenance": {
+        "notes": "128-token sparse blocks decompose into two 64-key tiles per block at head_size 128, the flagship geometry; the second mask row skips block 0 so its diagonal block carries the ragged causal edge."
+      },
+      "attrs": { "num_heads": 2, "kv_num_heads": 1, "sparse_block_size": 128 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 256, 256],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 256, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 256, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 1, 256, 128], "data": { "kind": "constant", "value": 0.0 } },
+        "pastValueT": { "dtype": "float32", "shape": [1, 1, 256, 128], "data": { "kind": "constant", "value": 0.0 } },
+        "blockRowIndicesT": { "dtype": "int32", "shape": [1, 3], "data": { "kind": "values", "values": [0, 1, 2] } },
+        "blockColIndicesT": { "dtype": "int32", "shape": [1, 2], "data": { "kind": "values", "values": [0, 1] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } },
+        "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 256, 256], "tolerance": 0.0001, "relTolerance": 0.0005 },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 1, 256, 128], "tolerance": 0.00002, "relTolerance": 0.0002 },
+        "pastValueT": { "dtype": "float32", "shape": [1, 1, 256, 128], "tolerance": 0.00002, "relTolerance": 0.0002 }
+      }
+    },
+    {
+      "name": "sgmat_chunked_two_mask_rows",
+      "provenance": {
+        "notes": "Chunked prefill with per-batch histories of 100 and 68 tokens: each 64-query tile straddles two mask rows at different split points, the misaligned path a runtime past length forces."
+      },
+      "attrs": { "num_heads": 2, "kv_num_heads": 2, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 64, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [2, 64, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [2, 64, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.017, "scale": 0.25 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 2, 256, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.013, "scale": 0.25 }
+        },
+        "blockRowIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [0, 1, 3, 5, 7, 0, 1, 2, 4, 5] }
+        },
+        "blockColIndicesT": {
+          "dtype": "int32",
+          "shape": [2, 7],
+          "data": { "kind": "values", "values": [0, 0, 1, 0, 2, 1, 3, 0, 1, 0, 2, 3, 0, 0] }
+        },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } },
+        "keyTotalSequenceLengthsT": {
+          "dtype": "int32",
+          "shape": [2],
+          "data": { "kind": "values", "values": [164, 132] }
+        }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [2, 64, 128], "tolerance": 0.0001, "relTolerance": 0.0005 },
+        "pastKeyT": { "dtype": "float32", "shape": [2, 2, 256, 64], "tolerance": 0.00002, "relTolerance": 0.0002 },
+        "pastValueT": { "dtype": "float32", "shape": [2, 2, 256, 64], "tolerance": 0.00002, "relTolerance": 0.0002 }
+      }
+    },
+    {
+      "name": "sgmat_empty_row_uniform_fallback",
+      "provenance": {
+        "notes": "The second mask row selects no blocks, so all 64 of its queries carry zero probability mass and take the uniform average over causal history on the subgroup-matrix route."
+      },
+      "attrs": { "num_heads": 1, "kv_num_heads": 1, "sparse_block_size": 64 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 128, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 128, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 128, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 1, 128, 64], "data": { "kind": "constant", "value": 0.0 } },
+        "pastValueT": { "dtype": "float32", "shape": [1, 1, 128, 64], "data": { "kind": "constant", "value": 0.0 } },
+        "blockRowIndicesT": { "dtype": "int32", "shape": [1, 3], "data": { "kind": "values", "values": [0, 1, 1] } },
+        "blockColIndicesT": { "dtype": "int32", "shape": [1, 1], "data": { "kind": "values", "values": [0] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [128] } },
+        "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [128] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 128, 64], "tolerance": 0.0001, "relTolerance": 0.0005 },
+        "pastKeyT": { "dtype": "float32", "shape": [1, 1, 128, 64], "tolerance": 0.00002, "relTolerance": 0.0002 },
+        "pastValueT": { "dtype": "float32", "shape": [1, 1, 128, 64], "tolerance": 0.00002, "relTolerance": 0.0002 }
+      }
     }
   ]
 }