sync 91d990483a17

Browse files

Files changed (8) hide show

README.md +31 -18
build/webgpu/attn-flash-decode-splitk-merge.wgsl.jinja +2 -13
build/webgpu/bench.json +1 -2
build/webgpu/manifest.json +285 -543
build/webgpu/metadata.json +21 -9
build/webgpu/paged-attention.wgsl.jinja +33 -46
build/webgpu/paged-scatter-kv.wgsl.jinja +17 -18
build/webgpu/test.json +51 -7

README.md CHANGED Viewed

@@ -18,25 +18,25 @@ See the [ONNX Runtime `PagedAttention` contrib-operator spec](https://github.com
 ## Inputs
-| Name | Bind key | Logical dtype | WebGPU storage | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- | --- |
-| `query` | `queryT` | `T` | same as logical dtype | `2` | — | Packed queries of shape `(num_tokens, num_heads * head_size)`, or `(num_tokens, (num_heads + 2 * kv_num_heads) * head_size)` when `key` and `value` are absent and Q, K and V share one row. | required |
-| `key` | `keyT` | `T` | same as logical dtype | `2` | — | Keys of shape `(num_tokens, kv_num_heads * head_size)`. Absent means `query` carries packed `[Q\|K\|V]`. | optional |
-| `value` | `valueT` | `T` | same as logical dtype | `2` | — | Values of shape `(num_tokens, kv_num_heads * head_size)`. Present exactly when `key` is. | optional |
-| `key_cache` | `keyCacheT` | `T` | same as logical dtype | `4` | — | Block-based key cache of shape `(num_blocks, block_size, kv_num_heads, head_size)`, updated in place. | required |
-| `value_cache` | `valueCacheT` | `T` | same as logical dtype | `4` | — | Block-based value cache with the same shape as `key_cache`, updated in place. | required |
-| `cumulative_sequence_length` | `cumulativeSequenceLengthT` | `S` | `int32` | `1` | — | Exclusive prefix sums of the per-sequence token counts, shape `(batch_size + 1)`; sequence `b` owns packed tokens `[cum[b], cum[b+1])`. | required |
-| `past_seqlens` | `pastSeqlensT` | `S` | `int32` | `1` | — | Cached history length per sequence, shape `(batch_size)`. | required |
-| `block_table` | `blockTableT` | `S` | `int32` | `2` | — | Physical block index per sequence and logical block, shape `(batch_size, max_blocks_per_sequence)`. | required |
-| `slot_mapping` | `slotMappingT` | `S` | `int32` | `1` | — | Flat destination slot, `block_id * block_size + offset`, for each token; `-1` suppresses that token's cache write. When omitted, the slot is derived from `past_seqlens`. `block_table` remains required because it defines the read path. | optional |
 ## Outputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `output` | `outputT` | `T` | `2` | derived; see description | Attention output of shape `(num_tokens, num_heads * head_size)`. | required |
-| `key_cache` | `keyCacheT` | `T` | `4` | same as `key_cache` | Optional return alias for the updated in-place key cache. The runtime updates both caches together even when only this alias is requested. | optional |
-| `value_cache` | `valueCacheT` | `T` | `4` | same as `value_cache` | Optional return alias for the updated in-place value cache. The runtime updates both caches together even when only this alias is requested. | optional |
 ## Attributes
@@ -44,7 +44,7 @@ Attributes and default values (overridable per request):
 | Attribute | Default | Description |
 | --- | --- | --- |
-| `is_causal` | `1` | Whether to apply causal masking. This package supports only value 1. Older ORT schema revisions omit this attribute and are always causal. |
 | `kv_num_heads` | — | Number of key/value heads. |
 | `num_heads` | — | Number of query heads. |
 | `scale` | — | Scale applied to query-key products; zero or omission selects `1 / sqrt(head_size)`. |
@@ -56,13 +56,22 @@ Attributes and default values (overridable per request):
 | `T` | `float16` |
 | `S` | `int32` |
 ## Device requirements
 Every implementation variant requires `shader-f16`; the package has no variant-level fallback without that capability.
 ## 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
@@ -72,10 +81,14 @@ Every implementation variant requires `shader-f16`; the package has no variant-l
 ## Use with `@huggingface/kernels`
-The loader derives every required output's shape and logical dtype from the manifest contract and this call.
-It then allocates the result tensors automatically.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

 ## Inputs
+| Name | Upstream name | Logical dtype | WebGPU storage | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- | --- |
+| `queryT` | `query` | `T` | same as logical dtype | `2` | — | Packed queries of shape `(num_tokens, num_heads * head_size)`, or `(num_tokens, (num_heads + 2 * kv_num_heads) * head_size)` when `key` and `value` are absent and Q, K and V share one row. | required |
+| `keyT` | `key` | `T` | same as logical dtype | `2` | — | Keys of shape `(num_tokens, kv_num_heads * head_size)`. Absent means `query` carries packed `[Q\|K\|V]`. | optional |
+| `valueT` | `value` | `T` | same as logical dtype | `2` | — | Values of shape `(num_tokens, kv_num_heads * head_size)`. Present exactly when `key` is. | optional |
+| `keyCacheT` | `key_cache` | `T` | same as logical dtype | `4` | — | Block-based key cache of shape `(num_blocks, block_size, kv_num_heads, head_size)`, updated in place. | required |
+| `valueCacheT` | `value_cache` | `T` | same as logical dtype | `4` | — | Block-based value cache with the same shape as `key_cache`, updated in place. | required |
+| `cumulativeSequenceLengthT` | `cumulative_sequence_length` | `S` | `int32` | `1` | — | Exclusive prefix sums of the per-sequence token counts, shape `(batch_size + 1)`; sequence `b` owns packed tokens `[cum[b], cum[b+1])`. | required |
+| `pastSeqlensT` | `past_seqlens` | `S` | `int32` | `1` | — | Cached history length per sequence, shape `(batch_size)`. | required |
+| `blockTableT` | `block_table` | `S` | `int32` | `2` | — | Physical block index per sequence and logical block, shape `(batch_size, max_blocks_per_sequence)`. | required |
+| `slotMappingT` | `slot_mapping` | `S` | `int32` | `1` | — | Flat destination slot, `block_id * block_size + offset`, for each token; `-1` suppresses that token's cache write. When omitted, the slot is derived from `past_seqlens`. `block_table` remains required because it defines the read path. | optional |
 ## Outputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `outputT` | `output` | `T` | `2` | derived | Attention output of shape `(num_tokens, num_heads * head_size)`. | required |
+| `keyCacheT` | `key_cache` | `T` | `4` | same as `keyCacheT` | Optional return alias for the updated in-place key cache. Both caches are updated even when only this result is requested. | optional |
+| `valueCacheT` | `value_cache` | `T` | `4` | same as `valueCacheT` | Optional return alias for the updated in-place value cache. Both caches are updated even when only this result is requested. | optional |
 ## Attributes
 | Attribute | Default | Description |
 | --- | --- | --- |
+| `is_causal` | `1` | Whether to apply causal masking. This package supports only value 1. |
 | `kv_num_heads` | — | Number of key/value heads. |
 | `num_heads` | — | Number of query heads. |
 | `scale` | — | Scale applied to query-key products; zero or omission selects `1 / sqrt(head_size)`. |
 | `T` | `float16` |
 | `S` | `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_derived_splitk` — Splits each token's key history into contiguous ranges, processes one range per workgroup, and merges the resulting online-softmax states. This adds parallel key ranges for grouped-query decode shapes with few `(token, KV head)` tuples.
+- `separate_slot_splitk` — Splits each token's key history into contiguous ranges, processes one range per workgroup, and merges the resulting online-softmax states. This adds parallel key ranges for grouped-query decode shapes with few `(token, KV head)` tuples.
+- `packed_derived_splitk` — Splits each token's key history into contiguous ranges, processes one range per workgroup, and merges the resulting online-softmax states. This adds parallel key ranges for grouped-query decode shapes with few `(token, KV head)` tuples.
+- `packed_slot_splitk` — Splits each token's key history into contiguous ranges, processes one range per workgroup, and merges the resulting online-softmax states. This adds parallel key ranges for grouped-query decode shapes with few `(token, KV head)` tuples.
 ## Device requirements
 Every implementation variant requires `shader-f16`; the package has no variant-level fallback without that capability.
 ## Files
+- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
 - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 ## Use with `@huggingface/kernels`
+```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.2
+```
+Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

build/webgpu/attn-flash-decode-splitk-merge.wgsl.jinja CHANGED Viewed

@@ -3,6 +3,7 @@ enable f16;
 {% endif %}
 {% if splitQueries is not defined %}{% set splitQueries = false %}{% endif %}
 {% if hasGate is not defined %}{% set hasGate = false %}{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 // Split-K flash decode, pass 2 of 2. Combines the per-split un-normalized online
@@ -22,13 +23,6 @@ const NUM_SPLITS: u32 = {{ numSplits }}u;
 const Q_SEQ: u32 = {{ qSeq }}u;
 {% endif %}
 const Q_HIDDEN_V4: u32 = {{ qHiddenV4 }}u;
-{% if hasBias %}
-// Packed [Q; K; V] bias rows. The V bias is token-independent, so the split
-// pass leaves it off and it is added here once after the final normalize,
-// matching the non-split flash-decode semantics exactly.
-const HEAD_DIM: u32 = {{ headDim }}u;
-const Q_HIDDEN: u32 = {{ qHidden }}u;
-{% endif %}
 // 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
@@ -99,13 +93,8 @@ fn main(
   // and globalDenom lands on exactly 0. Emit 0 rather than dividing so an all-empty
   // row has a defined zero result.
   var outValue = select(vec4<f32>(0.0), outV / globalDenom, globalDenom > 0.0);
-{% if hasBias %}
-  // V bias row base: 2 * qHidden (skip the packed Q and K bias blocks) + this head.
-  let vBiasBase = 2u * Q_HIDDEN + h * HEAD_DIM + d4 * 4u;
-  outValue = outValue + vec4<f32>(bias[vBiasBase], bias[vBiasBase + 1u], bias[vBiasBase + 2u], bias[vBiasBase + 3u]);
-{% endif %}
 {% if hasGate %}
-  // Qwen3.5 full attention gates the normalized attention output elementwise.
   let gateV = vec4<f32>(gate[qBaseV4 + d4]);
   outValue = outValue * (vec4<f32>(1.0) / (vec4<f32>(1.0) + exp(-gateV)));
 {% endif %}

 {% endif %}
 {% if splitQueries is not defined %}{% set splitQueries = false %}{% endif %}
 {% if hasGate is not defined %}{% set hasGate = false %}{% endif %}
+{% set qSeq = qSeq | default(0) %}
 {{ env.wgsl.resourceDeclarations }}
 // Split-K flash decode, pass 2 of 2. Combines the per-split un-normalized online
 const Q_SEQ: u32 = {{ qSeq }}u;
 {% endif %}
 const Q_HIDDEN_V4: u32 = {{ qHiddenV4 }}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
   // and globalDenom lands on exactly 0. Emit 0 rather than dividing so an all-empty
   // row has a defined zero result.
   var outValue = select(vec4<f32>(0.0), outV / globalDenom, globalDenom > 0.0);
 {% if hasGate %}
+  // The gated route multiplies the normalized attention output elementwise by its gate.
   let gateV = vec4<f32>(gate[qBaseV4 + d4]);
   outValue = outValue * (vec4<f32>(1.0) / (vec4<f32>(1.0) + exp(-gateV)));
 {% endif %}

build/webgpu/bench.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "com.microsoft.PagedAttention",
   "tunableSpace": {
     "WORKGROUP_SIZE": [32, 64, 128],
     "SCATTER_WORKGROUP_SIZE": [64, 128, 256],
@@ -32,7 +31,7 @@
         "metrics": [{ "type": "bandwidth", "value": "2 * args.batch * args.kvHeads * args.past * args.headDim * 2" }]
       },
       "provenance": {
-        "notes": "The KV bytes are counted off named vars like the model cases, not hardcoded: the literal this replaced used the QUERY head count for a grouped-query shape and four bytes for a float16 cache, so it over-stated the traffic eightfold and the case reported an impossible rate."
       }
     },
     {

 {
   "tunableSpace": {
     "WORKGROUP_SIZE": [32, 64, 128],
     "SCATTER_WORKGROUP_SIZE": [64, 128, 256],
         "metrics": [{ "type": "bandwidth", "value": "2 * args.batch * args.kvHeads * args.past * args.headDim * 2" }]
       },
       "provenance": {
+        "notes": "The bandwidth metric derives KV traffic from the key/value head count and the two-byte float16 cache element size."
       }
     },
     {

build/webgpu/manifest.json CHANGED Viewed

@@ -2,133 +2,40 @@
   "domain": "com.microsoft",
   "name": "PagedAttention",
   "sinceVersion": 1,
-  "description": "Attention over a block-based (paged) KV cache: `cumulative_sequence_length` marks the sequence boundaries and `block_table` maps a sequence's history onto scattered blocks. This step's K/V are scattered into the cache, then attended with that history. Grouped-query heads, `scale`, packed `[Q|K|V]`, `slot_mapping`, and float16 cache storage are supported; the cache outputs alias the input caches and are updated in place. Rotary embeddings, softcap, local windows, LATENT layout, narrower value heads, quantized KV, head sinks, q/k normalization, scales, and attention metadata are not implemented.",
-  "inputs": [
-    {
-      "role": "query",
-      "dtype": "T",
-      "rank": 2,
-      "description": "Packed queries of shape `(num_tokens, num_heads * head_size)`, or `(num_tokens, (num_heads + 2 * kv_num_heads) * head_size)` when `key` and `value` are absent and Q, K and V share one row."
-    },
-    {
-      "role": "key",
-      "dtype": "T",
-      "rank": 2,
-      "optional": true,
-      "description": "Keys of shape `(num_tokens, kv_num_heads * head_size)`. Absent means `query` carries packed `[Q|K|V]`."
-    },
-    {
-      "role": "value",
-      "dtype": "T",
-      "rank": 2,
-      "optional": true,
-      "description": "Values of shape `(num_tokens, kv_num_heads * head_size)`. Present exactly when `key` is."
-    },
-    {
-      "role": "key_cache",
-      "dtype": "T",
-      "rank": 4,
-      "description": "Block-based key cache of shape `(num_blocks, block_size, kv_num_heads, head_size)`, updated in place."
-    },
-    {
-      "role": "value_cache",
-      "dtype": "T",
-      "rank": 4,
-      "description": "Block-based value cache with the same shape as `key_cache`, updated in place."
-    },
-    {
-      "role": "cumulative_sequence_length",
-      "dtype": "S",
-      "rank": 1,
-      "description": "Exclusive prefix sums of the per-sequence token counts, shape `(batch_size + 1)`; sequence `b` owns packed tokens `[cum[b], cum[b+1])`."
-    },
-    {
-      "role": "past_seqlens",
-      "dtype": "S",
-      "rank": 1,
-      "description": "Cached history length per sequence, shape `(batch_size)`."
-    },
-    {
-      "role": "block_table",
-      "dtype": "S",
-      "rank": 2,
-      "description": "Physical block index per sequence and logical block, shape `(batch_size, max_blocks_per_sequence)`."
-    },
-    {
-      "role": "slot_mapping",
-      "dtype": "S",
-      "rank": 1,
-      "optional": true,
-      "description": "Flat destination slot, `block_id * block_size + offset`, for each token; `-1` suppresses that token's cache write. When omitted, the slot is derived from `past_seqlens`. `block_table` remains required because it defines the read path."
-    }
-  ],
-  "outputs": [
-    {
-      "role": "output",
       "dtype": "T",
       "rank": 2,
-      "shape": "[dim(shapes.queryT, 0), attrs.num_heads * headSize]",
-      "description": "Attention output of shape `(num_tokens, num_heads * head_size)`."
     },
-    {
-      "role": "key_cache",
-      "dtype": "T",
-      "rank": 4,
-      "shape": "shapes.keyCacheT",
-      "optional": true,
-      "description": "Optional return alias for the updated in-place key cache. The runtime updates both caches together even when only this alias is requested."
-    },
-    {
-      "role": "value_cache",
-      "dtype": "T",
-      "rank": 4,
-      "shape": "shapes.valueCacheT",
-      "optional": true,
-      "description": "Optional return alias for the updated in-place value cache. The runtime updates both caches together even when only this alias is requested."
-    }
-  ],
-  "attributes": { "is_causal": 1 },
-  "attributeDescriptions": {
-    "is_causal": "Whether to apply causal masking. This package supports only value 1. Older ORT schema revisions omit this attribute and are always causal.",
-    "kv_num_heads": "Number of key/value heads.",
-    "num_heads": "Number of query heads.",
-    "scale": "Scale applied to query-key products; zero or omission selects `1 / sqrt(head_size)`."
   },
   "attributeConstraints": {
     "is_causal": { "values": [1] },
     "kv_num_heads": { "required": true },
     "num_heads": { "required": true }
   },
   "typeConstraints": { "T": ["float16"], "S": ["int32"] },
-  "args": {
-    "queryT": { "kind": "tensor", "semantic": "query", "role": "input" },
-    "keyT": { "kind": "tensor", "semantic": "key", "role": "input", "required": false },
-    "valueT": { "kind": "tensor", "semantic": "value", "role": "input", "required": false },
-    "keyCacheT": { "kind": "tensor", "semantic": "key_cache", "role": "inout" },
-    "valueCacheT": { "kind": "tensor", "semantic": "value_cache", "role": "inout" },
-    "cumulativeSequenceLengthT": {
-      "kind": "tensor",
-      "semantic": "cumulative_sequence_length",
-      "role": "input",
-      "dtype": "int32"
-    },
-    "pastSeqlensT": { "kind": "tensor", "semantic": "past_seqlens", "role": "input", "dtype": "int32" },
-    "blockTableT": { "kind": "tensor", "semantic": "block_table", "role": "input", "dtype": "int32" },
-    "slotMappingT": {
-      "kind": "tensor",
-      "semantic": "slot_mapping",
-      "role": "input",
-      "dtype": "int32",
-      "required": false
-    },
-    "outputT": { "kind": "tensor", "semantic": "output", "role": "output" }
-  },
   "tunables": {
-    "WORKGROUP_SIZE": 64,
-    "SCATTER_WORKGROUP_SIZE": 64,
-    "MAX_SPLITS": 16,
-    "SPLIT_MIN_KEYS": 128,
-    "SPLIT_TARGET_WORKGROUPS": 1024
   },
   "derive": {
     "tokenCount": "dim(shapes.queryT, 0)",
@@ -150,35 +57,31 @@
     "packedKv": "not present.keyT and not present.valueT and dim(shapes.queryT, 1) == packedStride",
     "slotShapeOk": "dim(shapes.slotMappingT, 0) == tokenCount and ranks.slotMappingT == 1 if present.slotMappingT else true",
     "pagedContractOk": "cacheShapeOk and headLayoutOk and scheduleShapeOk and ioShapeOk and slotShapeOk",
-    "dispatchFits": "tokenCount <= device.limits.maxComputeWorkgroupsPerDimension and attrs.num_heads <= device.limits.maxComputeWorkgroupsPerDimension and tunables.WORKGROUP_SIZE <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.SCATTER_WORKGROUP_SIZE <= device.limits.maxComputeInvocationsPerWorkgroup",
     "pagedBaseWorkgroups": "tokenCount * attrs.kv_num_heads",
     "pagedKeyCeiling": "maxBlocks * blockSize",
     "pagedSplitWg": "min(256, max(32, pow2ceil(headVec)))",
     "pagedSplitCap": "max(1, min(tunables.MAX_SPLITS, pagedKeyCeiling / tunables.SPLIT_MIN_KEYS))",
     "numSplits": "max(1, min(pagedSplitCap, ceilDiv(tunables.SPLIT_TARGET_WORKGROUPS, max(1, pagedBaseWorkgroups))))",
     "pagedSplitScratchBytes": "tokenCount * attrs.num_heads * numSplits * headSize * 4",
     "pagedSplitStatsBytes": "2 * tokenCount * attrs.num_heads * numSplits * 4",
-    "pagedSplitFits": "numSplits >= 2 and cacheVec4Ok and headVec >= 1 and pagedSplitScratchBytes <= device.limits.maxStorageBufferBindingSize and pagedSplitScratchBytes <= device.limits.maxBufferSize and pagedSplitStatsBytes <= device.limits.maxStorageBufferBindingSize and pagedSplitStatsBytes <= device.limits.maxBufferSize and numSplits <= device.limits.maxComputeWorkgroupsPerDimension and headVec <= device.limits.maxComputeInvocationsPerWorkgroup and headVec <= device.limits.maxComputeWorkgroupSizeX and pagedSplitWg <= device.limits.maxComputeInvocationsPerWorkgroup and pagedSplitWg <= device.limits.maxComputeWorkgroupSizeX"
-  },
-  "constants": {
     "aScalar": "dtypes.T",
     "scalar": "dtypes.T",
     "numHeads": "attrs.num_heads",
     "kvNumHeads": "attrs.kv_num_heads",
-    "headSize": "headSize",
-    "blockSize": "blockSize",
-    "maxBlocks": "maxBlocks",
-    "qHidden": "qHidden",
-    "kvHidden": "kvHidden",
-    "packedStride": "packedStride",
     "packedQkv": "not present.keyT",
     "hasSlotMapping": "present.slotMappingT",
-    "qPerKv": "qPerKv",
-    "attnWorkgroup": "tunables.WORKGROUP_SIZE",
-    "headVec": "headVec",
     "cacheVec": "(\"vec4<f16>\" if dtypes.T == \"f16\" else \"vec4<f32>\") if cacheVec4Ok else dtypes.T",
-    "cacheVec4Ok": "cacheVec4Ok",
-    "numSplits": "numSplits",
     "headDimV4": "headVec",
     "headDim": "headSize",
     "qNumHeads": "attrs.num_heads",
@@ -186,367 +89,61 @@
     "hasBias": false,
     "outVec": "\"vec4<f16>\" if dtypes.T == \"f16\" else \"vec4<f32>\""
   },
-  "bindingSets": {
-    "scatterSeparateDerived": [
-      {
-        "name": "key",
-        "arg": "keyT",
-        "semantic": "key",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "value",
-        "arg": "valueT",
-        "semantic": "value",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "key_cache",
-        "arg": "keyCacheT",
-        "semantic": "key_cache",
-        "buffer": { "type": "storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "value_cache",
-        "arg": "valueCacheT",
-        "semantic": "value_cache",
-        "buffer": { "type": "storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "cumulative_sequence_length",
-        "arg": "cumulativeSequenceLengthT",
-        "semantic": "cumulative_sequence_length",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "past_seqlens",
-        "arg": "pastSeqlensT",
-        "semantic": "past_seqlens",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "block_table",
-        "arg": "blockTableT",
-        "semantic": "block_table",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "batchSize", "type": "u32", "value": "batchSize" },
-            { "name": "scatterCount", "type": "u32", "value": "tokenCount * attrs.kv_num_heads * headSize" }
-          ]
-        }
-      }
-    ],
-    "scatterSeparateSlot": [
-      {
-        "name": "key",
-        "arg": "keyT",
-        "semantic": "key",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "value",
-        "arg": "valueT",
-        "semantic": "value",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "key_cache",
-        "arg": "keyCacheT",
-        "semantic": "key_cache",
-        "buffer": { "type": "storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "value_cache",
-        "arg": "valueCacheT",
-        "semantic": "value_cache",
-        "buffer": { "type": "storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "slot_mapping",
-        "arg": "slotMappingT",
-        "semantic": "slot_mapping",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [{ "name": "scatterCount", "type": "u32", "value": "tokenCount * attrs.kv_num_heads * headSize" }]
-        }
-      }
-    ],
-    "scatterPackedDerived": [
-      {
-        "name": "query",
-        "arg": "queryT",
-        "semantic": "query",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "key_cache",
-        "arg": "keyCacheT",
-        "semantic": "key_cache",
-        "buffer": { "type": "storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "value_cache",
-        "arg": "valueCacheT",
-        "semantic": "value_cache",
-        "buffer": { "type": "storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "cumulative_sequence_length",
-        "arg": "cumulativeSequenceLengthT",
-        "semantic": "cumulative_sequence_length",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "past_seqlens",
-        "arg": "pastSeqlensT",
-        "semantic": "past_seqlens",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "block_table",
-        "arg": "blockTableT",
-        "semantic": "block_table",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "batchSize", "type": "u32", "value": "batchSize" },
-            { "name": "scatterCount", "type": "u32", "value": "tokenCount * attrs.kv_num_heads * headSize" }
-          ]
-        }
-      }
-    ],
-    "scatterPackedSlot": [
-      {
-        "name": "query",
-        "arg": "queryT",
-        "semantic": "query",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "key_cache",
-        "arg": "keyCacheT",
-        "semantic": "key_cache",
-        "buffer": { "type": "storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "value_cache",
-        "arg": "valueCacheT",
-        "semantic": "value_cache",
-        "buffer": { "type": "storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "slot_mapping",
-        "arg": "slotMappingT",
-        "semantic": "slot_mapping",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [{ "name": "scatterCount", "type": "u32", "value": "tokenCount * attrs.kv_num_heads * headSize" }]
-        }
-      }
-    ],
-    "attention": [
-      {
-        "name": "query",
-        "arg": "queryT",
-        "semantic": "query",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "key_cache",
-        "arg": "keyCacheT",
-        "semantic": "key_cache",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$cacheVec"
-      },
-      {
-        "name": "value_cache",
-        "arg": "valueCacheT",
-        "semantic": "value_cache",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$cacheVec"
-      },
-      {
-        "name": "cumulative_sequence_length",
-        "arg": "cumulativeSequenceLengthT",
-        "semantic": "cumulative_sequence_length",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "past_seqlens",
-        "arg": "pastSeqlensT",
-        "semantic": "past_seqlens",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "block_table",
-        "arg": "blockTableT",
-        "semantic": "block_table",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "output",
-        "arg": "outputT",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "batchSize", "type": "u32", "value": "batchSize" },
-            { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" }
-          ]
-        }
-      }
-    ],
-    "splitAttention": [
-      {
-        "name": "query",
-        "arg": "queryT",
-        "semantic": "query",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "key_cache",
-        "arg": "keyCacheT",
-        "semantic": "key_cache",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$cacheVec"
-      },
-      {
-        "name": "value_cache",
-        "arg": "valueCacheT",
-        "semantic": "value_cache",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$cacheVec"
-      },
-      {
-        "name": "cumulative_sequence_length",
-        "arg": "cumulativeSequenceLengthT",
-        "semantic": "cumulative_sequence_length",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "past_seqlens",
-        "arg": "pastSeqlensT",
-        "semantic": "past_seqlens",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "block_table",
-        "arg": "blockTableT",
-        "semantic": "block_table",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      { "name": "partial_out", "semantic": "partialOut", "buffer": { "type": "storage" }, "elementType": "vec4<f32>" },
-      {
-        "name": "partial_stats",
-        "semantic": "partialStats",
-        "buffer": { "type": "storage" },
-        "elementType": "vec2<f32>"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "batchSize", "type": "u32", "value": "batchSize" },
-            { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" }
-          ]
-        }
-      }
-    ],
-    "splitMerge": [
-      {
-        "name": "partial_out",
-        "semantic": "partialOut",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "vec4<f32>"
-      },
-      {
-        "name": "partial_stats",
-        "semantic": "partialStats",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "vec2<f32>"
-      },
-      {
-        "name": "output",
-        "arg": "outputT",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$outVec"
-      }
-    ]
   },
   "variants": [
     {
       "id": "separate_derived_splitk",
-      "description": "Splits each token's key history into contiguous ranges, one workgroup per range, and merges the per-range online-softmax states. It serves grouped-query decode shapes, where one workgroup per `(token, KV head)` leaves too few workgroups to fill the device.",
       "priority": 10,
       "requires": { "features": ["shader-f16"] },
-      "when": ["pagedContractOk", "dispatchFits", "pagedSplitFits", "separateKv", "not present.slotMappingT"],
-      "constants": { "attnWorkgroup": "pagedSplitWg" },
       "intermediates": [
         { "id": "partialOut", "dtype": "float32", "shape": "[tokenCount * attrs.num_heads * numSplits * headSize]" },
         { "id": "partialStats", "dtype": "float32", "shape": "[2 * tokenCount * attrs.num_heads * numSplits]" }
@@ -556,59 +153,95 @@
           "id": "scatter",
           "name": "PagedAttention.ScatterKV",
           "shader": "paged-scatter-kv.wgsl.jinja",
-          "bindings": "scatterSeparateDerived",
           "dispatch": {
-            "threads": "tokenCount * attrs.kv_num_heads * headSize",
-            "workgroupSize": "tunables.SCATTER_WORKGROUP_SIZE"
           }
         },
         {
           "id": "split_attention",
           "name": "PagedAttention.AttendSplitK",
-          "source": { "shader": "paged-attention.wgsl.jinja", "inputs": { "splitK": true } },
-          "bindings": "splitAttention",
           "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads", "z": "numSplits" }
         },
         {
           "id": "merge",
           "name": "PagedAttention.AttendSplitKMerge",
-          "source": { "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } },
-          "bindings": "splitMerge",
           "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "tokenCount" }
         }
       ]
     },
     {
       "id": "separate_derived",
       "requires": { "features": ["shader-f16"] },
-      "when": ["pagedContractOk", "dispatchFits", "separateKv", "not present.slotMappingT"],
       "passes": [
         {
           "id": "scatter",
           "name": "PagedAttention.ScatterKV",
           "shader": "paged-scatter-kv.wgsl.jinja",
-          "bindings": "scatterSeparateDerived",
           "dispatch": {
-            "threads": "tokenCount * attrs.kv_num_heads * headSize",
-            "workgroupSize": "tunables.SCATTER_WORKGROUP_SIZE"
           }
         },
         {
           "id": "main",
           "name": "PagedAttention.Attend",
           "shader": "paged-attention.wgsl.jinja",
-          "bindings": "attention",
           "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads" }
         }
       ]
     },
     {
       "id": "separate_slot_splitk",
-      "description": "Splits each token's key history into contiguous ranges, one workgroup per range, and merges the per-range online-softmax states. It serves grouped-query decode shapes, where one workgroup per `(token, KV head)` leaves too few workgroups to fill the device.",
       "priority": 10,
       "requires": { "features": ["shader-f16"] },
-      "when": ["pagedContractOk", "dispatchFits", "pagedSplitFits", "separateKv", "present.slotMappingT"],
-      "constants": { "attnWorkgroup": "pagedSplitWg" },
       "intermediates": [
         { "id": "partialOut", "dtype": "float32", "shape": "[tokenCount * attrs.num_heads * numSplits * headSize]" },
         { "id": "partialStats", "dtype": "float32", "shape": "[2 * tokenCount * attrs.num_heads * numSplits]" }
@@ -618,59 +251,95 @@
           "id": "scatter",
           "name": "PagedAttention.ScatterKV",
           "shader": "paged-scatter-kv.wgsl.jinja",
-          "bindings": "scatterSeparateSlot",
           "dispatch": {
-            "threads": "tokenCount * attrs.kv_num_heads * headSize",
-            "workgroupSize": "tunables.SCATTER_WORKGROUP_SIZE"
           }
         },
         {
           "id": "split_attention",
           "name": "PagedAttention.AttendSplitK",
-          "source": { "shader": "paged-attention.wgsl.jinja", "inputs": { "splitK": true } },
-          "bindings": "splitAttention",
           "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads", "z": "numSplits" }
         },
         {
           "id": "merge",
           "name": "PagedAttention.AttendSplitKMerge",
-          "source": { "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } },
-          "bindings": "splitMerge",
           "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "tokenCount" }
         }
       ]
     },
     {
       "id": "separate_slot",
       "requires": { "features": ["shader-f16"] },
-      "when": ["pagedContractOk", "dispatchFits", "separateKv", "present.slotMappingT"],
       "passes": [
         {
           "id": "scatter",
           "name": "PagedAttention.ScatterKV",
           "shader": "paged-scatter-kv.wgsl.jinja",
-          "bindings": "scatterSeparateSlot",
           "dispatch": {
-            "threads": "tokenCount * attrs.kv_num_heads * headSize",
-            "workgroupSize": "tunables.SCATTER_WORKGROUP_SIZE"
           }
         },
         {
           "id": "main",
           "name": "PagedAttention.Attend",
           "shader": "paged-attention.wgsl.jinja",
-          "bindings": "attention",
           "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads" }
         }
       ]
     },
     {
       "id": "packed_derived_splitk",
-      "description": "Splits each token's key history into contiguous ranges, one workgroup per range, and merges the per-range online-softmax states. It serves grouped-query decode shapes, where one workgroup per `(token, KV head)` leaves too few workgroups to fill the device.",
       "priority": 10,
       "requires": { "features": ["shader-f16"] },
-      "when": ["pagedContractOk", "dispatchFits", "pagedSplitFits", "packedKv", "not present.slotMappingT"],
-      "constants": { "attnWorkgroup": "pagedSplitWg" },
       "intermediates": [
         { "id": "partialOut", "dtype": "float32", "shape": "[tokenCount * attrs.num_heads * numSplits * headSize]" },
         { "id": "partialStats", "dtype": "float32", "shape": "[2 * tokenCount * attrs.num_heads * numSplits]" }
@@ -680,59 +349,95 @@
           "id": "scatter",
           "name": "PagedAttention.ScatterKV",
           "shader": "paged-scatter-kv.wgsl.jinja",
-          "bindings": "scatterPackedDerived",
           "dispatch": {
-            "threads": "tokenCount * attrs.kv_num_heads * headSize",
-            "workgroupSize": "tunables.SCATTER_WORKGROUP_SIZE"
           }
         },
         {
           "id": "split_attention",
           "name": "PagedAttention.AttendSplitK",
-          "source": { "shader": "paged-attention.wgsl.jinja", "inputs": { "splitK": true } },
-          "bindings": "splitAttention",
           "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads", "z": "numSplits" }
         },
         {
           "id": "merge",
           "name": "PagedAttention.AttendSplitKMerge",
-          "source": { "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } },
-          "bindings": "splitMerge",
           "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "tokenCount" }
         }
       ]
     },
     {
       "id": "packed_derived",
       "requires": { "features": ["shader-f16"] },
-      "when": ["pagedContractOk", "dispatchFits", "packedKv", "not present.slotMappingT"],
       "passes": [
         {
           "id": "scatter",
           "name": "PagedAttention.ScatterKV",
           "shader": "paged-scatter-kv.wgsl.jinja",
-          "bindings": "scatterPackedDerived",
           "dispatch": {
-            "threads": "tokenCount * attrs.kv_num_heads * headSize",
-            "workgroupSize": "tunables.SCATTER_WORKGROUP_SIZE"
           }
         },
         {
           "id": "main",
           "name": "PagedAttention.Attend",
           "shader": "paged-attention.wgsl.jinja",
-          "bindings": "attention",
           "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads" }
         }
       ]
     },
     {
       "id": "packed_slot_splitk",
-      "description": "Splits each token's key history into contiguous ranges, one workgroup per range, and merges the per-range online-softmax states. It serves grouped-query decode shapes, where one workgroup per `(token, KV head)` leaves too few workgroups to fill the device.",
       "priority": 10,
       "requires": { "features": ["shader-f16"] },
-      "when": ["pagedContractOk", "dispatchFits", "pagedSplitFits", "packedKv", "present.slotMappingT"],
-      "constants": { "attnWorkgroup": "pagedSplitWg" },
       "intermediates": [
         { "id": "partialOut", "dtype": "float32", "shape": "[tokenCount * attrs.num_heads * numSplits * headSize]" },
         { "id": "partialStats", "dtype": "float32", "shape": "[2 * tokenCount * attrs.num_heads * numSplits]" }
@@ -742,48 +447,85 @@
           "id": "scatter",
           "name": "PagedAttention.ScatterKV",
           "shader": "paged-scatter-kv.wgsl.jinja",
-          "bindings": "scatterPackedSlot",
           "dispatch": {
-            "threads": "tokenCount * attrs.kv_num_heads * headSize",
-            "workgroupSize": "tunables.SCATTER_WORKGROUP_SIZE"
           }
         },
         {
           "id": "split_attention",
           "name": "PagedAttention.AttendSplitK",
-          "source": { "shader": "paged-attention.wgsl.jinja", "inputs": { "splitK": true } },
-          "bindings": "splitAttention",
           "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads", "z": "numSplits" }
         },
         {
           "id": "merge",
           "name": "PagedAttention.AttendSplitKMerge",
-          "source": { "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } },
-          "bindings": "splitMerge",
           "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "tokenCount" }
         }
       ]
     },
     {
       "id": "packed_slot",
       "requires": { "features": ["shader-f16"] },
-      "when": ["pagedContractOk", "dispatchFits", "packedKv", "present.slotMappingT"],
       "passes": [
         {
           "id": "scatter",
           "name": "PagedAttention.ScatterKV",
           "shader": "paged-scatter-kv.wgsl.jinja",
-          "bindings": "scatterPackedSlot",
           "dispatch": {
-            "threads": "tokenCount * attrs.kv_num_heads * headSize",
-            "workgroupSize": "tunables.SCATTER_WORKGROUP_SIZE"
           }
         },
         {
           "id": "main",
           "name": "PagedAttention.Attend",
           "shader": "paged-attention.wgsl.jinja",
-          "bindings": "attention",
           "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads" }
         }
       ]

   "domain": "com.microsoft",
   "name": "PagedAttention",
   "sinceVersion": 1,
+  "inputs": {
+    "queryT": { "onnx": "query", "dtype": "T", "rank": 2 },
+    "keyT": { "onnx": "key", "dtype": "T", "rank": 2, "optional": true },
+    "valueT": { "onnx": "value", "dtype": "T", "rank": 2, "optional": true },
+    "keyCacheT": { "onnx": "key_cache", "dtype": "T", "rank": 4 },
+    "valueCacheT": { "onnx": "value_cache", "dtype": "T", "rank": 4 },
+    "cumulativeSequenceLengthT": { "onnx": "cumulative_sequence_length", "dtype": "S", "rank": 1, "storage": "int32" },
+    "pastSeqlensT": { "onnx": "past_seqlens", "dtype": "S", "rank": 1, "storage": "int32" },
+    "blockTableT": { "onnx": "block_table", "dtype": "S", "rank": 2, "storage": "int32" },
+    "slotMappingT": { "onnx": "slot_mapping", "dtype": "S", "rank": 1, "optional": true, "storage": "int32" }
+  },
+  "outputs": {
+    "outputT": {
+      "onnx": "output",
       "dtype": "T",
       "rank": 2,
+      "shape": "[dim(shapes.queryT, 0), attrs.num_heads * headSize]"
     },
+    "keyCacheT": { "onnx": "key_cache", "dtype": "T", "rank": 4, "optional": true, "shape": "shapes.keyCacheT" },
+    "valueCacheT": { "onnx": "value_cache", "dtype": "T", "rank": 4, "optional": true, "shape": "shapes.valueCacheT" }
   },
+  "attributes": { "is_causal": { "default": 1 }, "kv_num_heads": {}, "num_heads": {}, "scale": {} },
   "attributeConstraints": {
     "is_causal": { "values": [1] },
     "kv_num_heads": { "required": true },
     "num_heads": { "required": true }
   },
   "typeConstraints": { "T": ["float16"], "S": ["int32"] },
   "tunables": {
+    "WORKGROUP_SIZE": { "default": 64 },
+    "SCATTER_WORKGROUP_SIZE": { "default": 64 },
+    "MAX_SPLITS": { "default": 16 },
+    "SPLIT_MIN_KEYS": { "default": 128 },
+    "SPLIT_TARGET_WORKGROUPS": { "default": 1024 }
   },
   "derive": {
     "tokenCount": "dim(shapes.queryT, 0)",
     "packedKv": "not present.keyT and not present.valueT and dim(shapes.queryT, 1) == packedStride",
     "slotShapeOk": "dim(shapes.slotMappingT, 0) == tokenCount and ranks.slotMappingT == 1 if present.slotMappingT else true",
     "pagedContractOk": "cacheShapeOk and headLayoutOk and scheduleShapeOk and ioShapeOk and slotShapeOk",
+    "dispatchFits": "tokenCount <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and attrs.num_heads <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and tunables.WORKGROUP_SIZE <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.SCATTER_WORKGROUP_SIZE <= device.limits.maxComputeInvocationsPerWorkgroup",
     "pagedBaseWorkgroups": "tokenCount * attrs.kv_num_heads",
+    "pagedAttnQueryFloats": "2 * headSize * qPerKv",
+    "pagedAttnLaneFloats": "3 * qPerKv + 1",
+    "pagedAttnWgBudget": "floor(device.limits.maxComputeWorkgroupStorageSize / 4 - pagedAttnQueryFloats) / pagedAttnLaneFloats",
+    "pagedAttnWg": "max(32, min(tunables.WORKGROUP_SIZE, pow2ceil(max(1, floor(pagedAttnWgBudget)) + 1) / 2))",
+    "pagedAttnStorageBytes": "(pagedAttnQueryFloats + pagedAttnLaneFloats * pagedAttnWg) * 4",
+    "pagedAttnStorageFits": "pagedAttnStorageBytes <= device.limits.maxComputeWorkgroupStorageSize and pagedAttnWg <= device.limits.maxComputeInvocationsPerWorkgroup and pagedAttnWg <= device.limits.maxComputeWorkgroupSizeX",
     "pagedKeyCeiling": "maxBlocks * blockSize",
     "pagedSplitWg": "min(256, max(32, pow2ceil(headVec)))",
+    "pagedSplitStorageBytes": "(2 * headSize * qPerKv + (3 * qPerKv + 1) * pagedSplitWg) * 4",
+    "pagedSplitStorageFits": "pagedSplitStorageBytes <= device.limits.maxComputeWorkgroupStorageSize",
     "pagedSplitCap": "max(1, min(tunables.MAX_SPLITS, pagedKeyCeiling / tunables.SPLIT_MIN_KEYS))",
     "numSplits": "max(1, min(pagedSplitCap, ceilDiv(tunables.SPLIT_TARGET_WORKGROUPS, max(1, pagedBaseWorkgroups))))",
     "pagedSplitScratchBytes": "tokenCount * attrs.num_heads * numSplits * headSize * 4",
     "pagedSplitStatsBytes": "2 * tokenCount * attrs.num_heads * numSplits * 4",
+    "pagedSplitFits": "numSplits >= 2 and cacheVec4Ok and headVec >= 1 and pagedSplitScratchBytes <= device.limits.maxStorageBufferBindingSize and pagedSplitScratchBytes <= device.limits.maxBufferSize and pagedSplitStatsBytes <= device.limits.maxStorageBufferBindingSize and pagedSplitStatsBytes <= device.limits.maxBufferSize and numSplits <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and headVec <= device.limits.maxComputeInvocationsPerWorkgroup and headVec <= device.limits.maxComputeWorkgroupSizeX and pagedSplitWg <= device.limits.maxComputeInvocationsPerWorkgroup and pagedSplitWg <= device.limits.maxComputeWorkgroupSizeX",
     "aScalar": "dtypes.T",
     "scalar": "dtypes.T",
     "numHeads": "attrs.num_heads",
     "kvNumHeads": "attrs.kv_num_heads",
     "packedQkv": "not present.keyT",
     "hasSlotMapping": "present.slotMappingT",
+    "attnWorkgroup": "pagedAttnWg",
     "cacheVec": "(\"vec4<f16>\" if dtypes.T == \"f16\" else \"vec4<f32>\") if cacheVec4Ok else dtypes.T",
     "headDimV4": "headVec",
     "headDim": "headSize",
     "qNumHeads": "attrs.num_heads",
     "hasBias": false,
     "outVec": "\"vec4<f16>\" if dtypes.T == \"f16\" else \"vec4<f32>\""
   },
+  "when": ["pagedContractOk", "dispatchFits"],
+  "bindings": {
+    "key": { "arg": "keyT", "buffer": "read-only-storage", "elementType": "$aScalar" },
+    "value": { "arg": "valueT", "buffer": "read-only-storage", "elementType": "$aScalar" },
+    "key_cache": { "arg": "keyCacheT", "buffer": "storage", "elementType": "$aScalar" },
+    "value_cache": { "arg": "valueCacheT", "buffer": "storage", "elementType": "$aScalar" },
+    "cumulative_sequence_length": {
+      "arg": "cumulativeSequenceLengthT",
+      "buffer": "read-only-storage",
+      "elementType": "i32"
+    },
+    "past_seqlens": { "arg": "pastSeqlensT", "buffer": "read-only-storage", "elementType": "i32" },
+    "block_table": { "arg": "blockTableT", "buffer": "read-only-storage", "elementType": "i32" },
+    "params": {
+      "buffer": "uniform",
+      "struct": [
+        { "name": "batchSize", "type": "u32", "value": "batchSize" },
+        { "name": "scatterCount", "type": "u32", "value": "tokenCount * attrs.kv_num_heads * headSize" }
+      ]
+    },
+    "slot_mapping": { "arg": "slotMappingT", "buffer": "read-only-storage", "elementType": "i32" },
+    "params_2": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [{ "name": "scatterCount", "type": "u32", "value": "tokenCount * attrs.kv_num_heads * headSize" }]
+    },
+    "query": { "arg": "queryT", "buffer": "read-only-storage", "elementType": "$aScalar" },
+    "key_cache_2": {
+      "arg": "keyCacheT",
+      "name": "key_cache",
+      "buffer": "read-only-storage",
+      "elementType": "$cacheVec"
+    },
+    "value_cache_2": {
+      "arg": "valueCacheT",
+      "name": "value_cache",
+      "buffer": "read-only-storage",
+      "elementType": "$cacheVec"
+    },
+    "params_3": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "batchSize", "type": "u32", "value": "batchSize" },
+        { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" }
+      ]
+    }
   },
   "variants": [
     {
       "id": "separate_derived_splitk",
       "priority": 10,
+      "when": ["pagedSplitFits", "pagedSplitStorageFits", "separateKv", "not present.slotMappingT"],
       "requires": { "features": ["shader-f16"] },
+      "derive": { "attnWorkgroup": "pagedSplitWg" },
       "intermediates": [
         { "id": "partialOut", "dtype": "float32", "shape": "[tokenCount * attrs.num_heads * numSplits * headSize]" },
         { "id": "partialStats", "dtype": "float32", "shape": "[2 * tokenCount * attrs.num_heads * numSplits]" }
           "id": "scatter",
           "name": "PagedAttention.ScatterKV",
           "shader": "paged-scatter-kv.wgsl.jinja",
+          "bindings": ["key", "value", "key_cache", "value_cache", "cumulative_sequence_length", "past_seqlens", "block_table", "params"],
           "dispatch": {
+            "x": "min(ceilDiv((tokenCount * attrs.kv_num_heads * headSize), (tunables.SCATTER_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((tokenCount * attrs.kv_num_heads * headSize), (tunables.SCATTER_WORKGROUP_SIZE)), 65535)",
+            "z": 1
           }
         },
         {
           "id": "split_attention",
           "name": "PagedAttention.AttendSplitK",
+          "shader": "paged-attention.wgsl.jinja",
+          "derive": { "splitK": true },
+          "bindings": [
+            "query",
+            "key_cache_2",
+            "value_cache_2",
+            "cumulative_sequence_length",
+            "past_seqlens",
+            "block_table",
+            { "scratch": "partialOut", "name": "partial_out", "elementType": "vec4<f32>" },
+            { "scratch": "partialStats", "name": "partial_stats", "elementType": "vec2<f32>" },
+            "params_3"
+          ],
           "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads", "z": "numSplits" }
         },
         {
           "id": "merge",
           "name": "PagedAttention.AttendSplitKMerge",
+          "shader": "attn-flash-decode-splitk-merge.wgsl.jinja",
+          "derive": { "layout": "\"bsh\"" },
+          "bindings": [
+            {
+              "scratch": "partialOut",
+              "name": "partial_out",
+              "buffer": "read-only-storage",
+              "elementType": "vec4<f32>"
+            },
+            {
+              "scratch": "partialStats",
+              "name": "partial_stats",
+              "buffer": "read-only-storage",
+              "elementType": "vec2<f32>"
+            },
+            { "arg": "outputT", "name": "output", "elementType": "$outVec" }
+          ],
           "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "tokenCount" }
         }
       ]
     },
     {
       "id": "separate_derived",
+      "when": ["pagedAttnStorageFits", "separateKv", "not present.slotMappingT"],
       "requires": { "features": ["shader-f16"] },
       "passes": [
         {
           "id": "scatter",
           "name": "PagedAttention.ScatterKV",
           "shader": "paged-scatter-kv.wgsl.jinja",
+          "bindings": ["key", "value", "key_cache", "value_cache", "cumulative_sequence_length", "past_seqlens", "block_table", "params"],
           "dispatch": {
+            "x": "min(ceilDiv((tokenCount * attrs.kv_num_heads * headSize), (tunables.SCATTER_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((tokenCount * attrs.kv_num_heads * headSize), (tunables.SCATTER_WORKGROUP_SIZE)), 65535)",
+            "z": 1
           }
         },
         {
           "id": "main",
           "name": "PagedAttention.Attend",
           "shader": "paged-attention.wgsl.jinja",
+          "bindings": [
+            "query",
+            "key_cache_2",
+            "value_cache_2",
+            "cumulative_sequence_length",
+            "past_seqlens",
+            "block_table",
+            { "arg": "outputT", "name": "output", "elementType": "$aScalar" },
+            "params_3"
+          ],
           "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads" }
         }
       ]
     },
     {
       "id": "separate_slot_splitk",
       "priority": 10,
+      "when": ["pagedSplitFits", "pagedSplitStorageFits", "separateKv", "present.slotMappingT"],
       "requires": { "features": ["shader-f16"] },
+      "derive": { "attnWorkgroup": "pagedSplitWg" },
       "intermediates": [
         { "id": "partialOut", "dtype": "float32", "shape": "[tokenCount * attrs.num_heads * numSplits * headSize]" },
         { "id": "partialStats", "dtype": "float32", "shape": "[2 * tokenCount * attrs.num_heads * numSplits]" }
           "id": "scatter",
           "name": "PagedAttention.ScatterKV",
           "shader": "paged-scatter-kv.wgsl.jinja",
+          "bindings": ["key", "value", "key_cache", "value_cache", "slot_mapping", "params_2"],
           "dispatch": {
+            "x": "min(ceilDiv((tokenCount * attrs.kv_num_heads * headSize), (tunables.SCATTER_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((tokenCount * attrs.kv_num_heads * headSize), (tunables.SCATTER_WORKGROUP_SIZE)), 65535)",
+            "z": 1
           }
         },
         {
           "id": "split_attention",
           "name": "PagedAttention.AttendSplitK",
+          "shader": "paged-attention.wgsl.jinja",
+          "derive": { "splitK": true },
+          "bindings": [
+            "query",
+            "key_cache_2",
+            "value_cache_2",
+            "cumulative_sequence_length",
+            "past_seqlens",
+            "block_table",
+            { "scratch": "partialOut", "name": "partial_out", "elementType": "vec4<f32>" },
+            { "scratch": "partialStats", "name": "partial_stats", "elementType": "vec2<f32>" },
+            "params_3"
+          ],
           "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads", "z": "numSplits" }
         },
         {
           "id": "merge",
           "name": "PagedAttention.AttendSplitKMerge",
+          "shader": "attn-flash-decode-splitk-merge.wgsl.jinja",
+          "derive": { "layout": "\"bsh\"" },
+          "bindings": [
+            {
+              "scratch": "partialOut",
+              "name": "partial_out",
+              "buffer": "read-only-storage",
+              "elementType": "vec4<f32>"
+            },
+            {
+              "scratch": "partialStats",
+              "name": "partial_stats",
+              "buffer": "read-only-storage",
+              "elementType": "vec2<f32>"
+            },
+            { "arg": "outputT", "name": "output", "elementType": "$outVec" }
+          ],
           "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "tokenCount" }
         }
       ]
     },
     {
       "id": "separate_slot",
+      "when": ["pagedAttnStorageFits", "separateKv", "present.slotMappingT"],
       "requires": { "features": ["shader-f16"] },
       "passes": [
         {
           "id": "scatter",
           "name": "PagedAttention.ScatterKV",
           "shader": "paged-scatter-kv.wgsl.jinja",
+          "bindings": ["key", "value", "key_cache", "value_cache", "slot_mapping", "params_2"],
           "dispatch": {
+            "x": "min(ceilDiv((tokenCount * attrs.kv_num_heads * headSize), (tunables.SCATTER_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((tokenCount * attrs.kv_num_heads * headSize), (tunables.SCATTER_WORKGROUP_SIZE)), 65535)",
+            "z": 1
           }
         },
         {
           "id": "main",
           "name": "PagedAttention.Attend",
           "shader": "paged-attention.wgsl.jinja",
+          "bindings": [
+            "query",
+            "key_cache_2",
+            "value_cache_2",
+            "cumulative_sequence_length",
+            "past_seqlens",
+            "block_table",
+            { "arg": "outputT", "name": "output", "elementType": "$aScalar" },
+            "params_3"
+          ],
           "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads" }
         }
       ]
     },
     {
       "id": "packed_derived_splitk",
       "priority": 10,
+      "when": ["pagedSplitFits", "pagedSplitStorageFits", "packedKv", "not present.slotMappingT"],
       "requires": { "features": ["shader-f16"] },
+      "derive": { "attnWorkgroup": "pagedSplitWg" },
       "intermediates": [
         { "id": "partialOut", "dtype": "float32", "shape": "[tokenCount * attrs.num_heads * numSplits * headSize]" },
         { "id": "partialStats", "dtype": "float32", "shape": "[2 * tokenCount * attrs.num_heads * numSplits]" }
           "id": "scatter",
           "name": "PagedAttention.ScatterKV",
           "shader": "paged-scatter-kv.wgsl.jinja",
+          "bindings": ["query", "key_cache", "value_cache", "cumulative_sequence_length", "past_seqlens", "block_table", "params"],
           "dispatch": {
+            "x": "min(ceilDiv((tokenCount * attrs.kv_num_heads * headSize), (tunables.SCATTER_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((tokenCount * attrs.kv_num_heads * headSize), (tunables.SCATTER_WORKGROUP_SIZE)), 65535)",
+            "z": 1
           }
         },
         {
           "id": "split_attention",
           "name": "PagedAttention.AttendSplitK",
+          "shader": "paged-attention.wgsl.jinja",
+          "derive": { "splitK": true },
+          "bindings": [
+            "query",
+            "key_cache_2",
+            "value_cache_2",
+            "cumulative_sequence_length",
+            "past_seqlens",
+            "block_table",
+            { "scratch": "partialOut", "name": "partial_out", "elementType": "vec4<f32>" },
+            { "scratch": "partialStats", "name": "partial_stats", "elementType": "vec2<f32>" },
+            "params_3"
+          ],
           "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads", "z": "numSplits" }
         },
         {
           "id": "merge",
           "name": "PagedAttention.AttendSplitKMerge",
+          "shader": "attn-flash-decode-splitk-merge.wgsl.jinja",
+          "derive": { "layout": "\"bsh\"" },
+          "bindings": [
+            {
+              "scratch": "partialOut",
+              "name": "partial_out",
+              "buffer": "read-only-storage",
+              "elementType": "vec4<f32>"
+            },
+            {
+              "scratch": "partialStats",
+              "name": "partial_stats",
+              "buffer": "read-only-storage",
+              "elementType": "vec2<f32>"
+            },
+            { "arg": "outputT", "name": "output", "elementType": "$outVec" }
+          ],
           "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "tokenCount" }
         }
       ]
     },
     {
       "id": "packed_derived",
+      "when": ["pagedAttnStorageFits", "packedKv", "not present.slotMappingT"],
       "requires": { "features": ["shader-f16"] },
       "passes": [
         {
           "id": "scatter",
           "name": "PagedAttention.ScatterKV",
           "shader": "paged-scatter-kv.wgsl.jinja",
+          "bindings": ["query", "key_cache", "value_cache", "cumulative_sequence_length", "past_seqlens", "block_table", "params"],
           "dispatch": {
+            "x": "min(ceilDiv((tokenCount * attrs.kv_num_heads * headSize), (tunables.SCATTER_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((tokenCount * attrs.kv_num_heads * headSize), (tunables.SCATTER_WORKGROUP_SIZE)), 65535)",
+            "z": 1
           }
         },
         {
           "id": "main",
           "name": "PagedAttention.Attend",
           "shader": "paged-attention.wgsl.jinja",
+          "bindings": [
+            "query",
+            "key_cache_2",
+            "value_cache_2",
+            "cumulative_sequence_length",
+            "past_seqlens",
+            "block_table",
+            { "arg": "outputT", "name": "output", "elementType": "$aScalar" },
+            "params_3"
+          ],
           "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads" }
         }
       ]
     },
     {
       "id": "packed_slot_splitk",
       "priority": 10,
+      "when": ["pagedSplitFits", "pagedSplitStorageFits", "packedKv", "present.slotMappingT"],
       "requires": { "features": ["shader-f16"] },
+      "derive": { "attnWorkgroup": "pagedSplitWg" },
       "intermediates": [
         { "id": "partialOut", "dtype": "float32", "shape": "[tokenCount * attrs.num_heads * numSplits * headSize]" },
         { "id": "partialStats", "dtype": "float32", "shape": "[2 * tokenCount * attrs.num_heads * numSplits]" }
           "id": "scatter",
           "name": "PagedAttention.ScatterKV",
           "shader": "paged-scatter-kv.wgsl.jinja",
+          "bindings": ["query", "key_cache", "value_cache", "slot_mapping", "params_2"],
           "dispatch": {
+            "x": "min(ceilDiv((tokenCount * attrs.kv_num_heads * headSize), (tunables.SCATTER_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((tokenCount * attrs.kv_num_heads * headSize), (tunables.SCATTER_WORKGROUP_SIZE)), 65535)",
+            "z": 1
           }
         },
         {
           "id": "split_attention",
           "name": "PagedAttention.AttendSplitK",
+          "shader": "paged-attention.wgsl.jinja",
+          "derive": { "splitK": true },
+          "bindings": [
+            "query",
+            "key_cache_2",
+            "value_cache_2",
+            "cumulative_sequence_length",
+            "past_seqlens",
+            "block_table",
+            { "scratch": "partialOut", "name": "partial_out", "elementType": "vec4<f32>" },
+            { "scratch": "partialStats", "name": "partial_stats", "elementType": "vec2<f32>" },
+            "params_3"
+          ],
           "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads", "z": "numSplits" }
         },
         {
           "id": "merge",
           "name": "PagedAttention.AttendSplitKMerge",
+          "shader": "attn-flash-decode-splitk-merge.wgsl.jinja",
+          "derive": { "layout": "\"bsh\"" },
+          "bindings": [
+            {
+              "scratch": "partialOut",
+              "name": "partial_out",
+              "buffer": "read-only-storage",
+              "elementType": "vec4<f32>"
+            },
+            {
+              "scratch": "partialStats",
+              "name": "partial_stats",
+              "buffer": "read-only-storage",
+              "elementType": "vec2<f32>"
+            },
+            { "arg": "outputT", "name": "output", "elementType": "$outVec" }
+          ],
           "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "tokenCount" }
         }
       ]
     },
     {
       "id": "packed_slot",
+      "when": ["pagedAttnStorageFits", "packedKv", "present.slotMappingT"],
       "requires": { "features": ["shader-f16"] },
       "passes": [
         {
           "id": "scatter",
           "name": "PagedAttention.ScatterKV",
           "shader": "paged-scatter-kv.wgsl.jinja",
+          "bindings": ["query", "key_cache", "value_cache", "slot_mapping", "params_2"],
           "dispatch": {
+            "x": "min(ceilDiv((tokenCount * attrs.kv_num_heads * headSize), (tunables.SCATTER_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((tokenCount * attrs.kv_num_heads * headSize), (tunables.SCATTER_WORKGROUP_SIZE)), 65535)",
+            "z": 1
           }
         },
         {
           "id": "main",
           "name": "PagedAttention.Attend",
           "shader": "paged-attention.wgsl.jinja",
+          "bindings": [
+            "query",
+            "key_cache_2",
+            "value_cache_2",
+            "cumulative_sequence_length",
+            "past_seqlens",
+            "block_table",
+            { "arg": "outputT", "name": "output", "elementType": "$aScalar" },
+            "params_3"
+          ],
           "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads" }
         }
       ]

build/webgpu/metadata.json CHANGED Viewed

@@ -1,20 +1,32 @@
 {
   "name": "com.microsoft.PagedAttention",
-  "id": "_com_microsoft_pagedattention_webgpu_da91370",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "attn-flash-decode-splitk-merge.wgsl.jinja": "8RI/yP6118n/rLRo0OG9NJ+kZ2xdmJyj093Tw8C07nc=",
-      "bench.json": "UwS/SwZvyxZNimhsfEODD8GUpSN1saYbJFO+2AJpBFE=",
-      "manifest.json": "4OMoAQRJcGm7t8TLViM098p15V0i788gvT/vpkfbt1c=",
-      "paged-attention.wgsl.jinja": "LcAWM3XNSdLIc67tQSUJ6v9Sa4QcRIetF3cnNwkdu9c=",
-      "paged-scatter-kv.wgsl.jinja": "WiuV+h3AOSDU0ENfvZ5iequjvL9TFR3h1gCQgGswPRA=",
-      "test.json": "AZ89NKbsfzyghJ0lIz57f3tm0r9H62tEPhY1uP5yWEM="
     }
   },
-  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
-  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.PagedAttention" }
 }

 {
   "name": "com.microsoft.PagedAttention",
+  "id": "_com_microsoft_pagedattention_webgpu_a50bcf6",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "attn-flash-decode-splitk-merge.wgsl.jinja": "uOgPgm5HioxH0ha2Z7Q0b/8qESsCmwKAS/198ogYX3I=",
+      "bench.json": "n1VW3Kl/whbn6CsIdfs0MHNrvUTskejhRkd5jqAAGmk=",
+      "manifest.json": "T7kLnpUD+K08/2E51GlZEK8v9YsGXliBIVnku1LRgRM=",
+      "paged-attention.wgsl.jinja": "NS4JduPGQ8x2YxgIap9fKuggGtyqse0u503A7rvzSM4=",
+      "paged-scatter-kv.wgsl.jinja": "EEb8eRJICdQOiQPsThu7rGm97W8iYFFH5zL3vArnQ8Q=",
+      "test.json": "UtJ6CP4M4fFlZm0lZdIrcX+8AxB/5gIC88fXZb6Iw5Q="
     }
   },
+  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.0",
+    "variants": {
+      "separate_derived_splitk": ["attn-flash-decode-splitk-merge.wgsl.jinja", "paged-attention.wgsl.jinja", "paged-scatter-kv.wgsl.jinja"],
+      "separate_derived": ["paged-attention.wgsl.jinja", "paged-scatter-kv.wgsl.jinja"],
+      "separate_slot_splitk": ["attn-flash-decode-splitk-merge.wgsl.jinja", "paged-attention.wgsl.jinja", "paged-scatter-kv.wgsl.jinja"],
+      "separate_slot": ["paged-attention.wgsl.jinja", "paged-scatter-kv.wgsl.jinja"],
+      "packed_derived_splitk": ["attn-flash-decode-splitk-merge.wgsl.jinja", "paged-attention.wgsl.jinja", "paged-scatter-kv.wgsl.jinja"],
+      "packed_derived": ["paged-attention.wgsl.jinja", "paged-scatter-kv.wgsl.jinja"],
+      "packed_slot_splitk": ["attn-flash-decode-splitk-merge.wgsl.jinja", "paged-attention.wgsl.jinja", "paged-scatter-kv.wgsl.jinja"],
+      "packed_slot": ["paged-attention.wgsl.jinja", "paged-scatter-kv.wgsl.jinja"]
+    }
+  }
 }

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

@@ -1,24 +1,17 @@
-{% set splitK = source.splitK is defined and source.splitK %}
 {% set keyLo = "splitStart" if splitK else "0u" %}
 {% set keyHi = "splitEnd" if splitK else "key_bound" %}
 enable f16;
 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.PagedAttention, attention pass.
-// One workgroup per (packed token, KV head). The workgroup sweeps that token's
-// KV history in tiles of WG keys with one key per thread, merging each tile into
-// a running online softmax, so the score row is never materialized.
-// A workgroup serves every query head that shares its KV head, reusing each
-// cache read across the group. Q_PER_KV == 1 reduces to ordinary attention.
-// The history is not contiguous: key j of sequence b lives in physical block
-// block_table[b, j / BLOCK_SIZE] at offset j % BLOCK_SIZE, which is the whole
-// point of a paged cache. Each tile resolves that indirection once per key and
-// parks the resulting row bases in workgroup memory, so the value accumulation
-// below re-reads a base instead of re-walking the block table for every head
-// dimension.
-// Masking is causal by construction rather than by comparison: a token's key
-// bound is its own absolute position plus one, so keys beyond it are never
-// visited.
 {% if splitK %}
 // Split-K partitions the key axis into NUM_SPLITS contiguous ranges. Each range
 // writes its unnormalized online state (max, denominator, and weighted V sum),
@@ -87,20 +80,22 @@ var<workgroup> key_rows: array<u32, WG>;
 // calls before their shared partial storage is reused.
 {% set combineSubgroups = combineSubgroups is defined and combineSubgroups %}
 {% if combineSubgroups %}
-// Per-subgroup partials are published into a deterministic slot: the subgroup's
-// ordinal index within the workgroup (lidx / sgSize). The online (m, d) merge
-// is not float-associative, so thread 0 must fold partials in a fixed order.
-// Subgroups partition a workgroup into contiguous ordinal ranges on supported
-// backends, so the ordinal slot is unique per subgroup and every slot in
-// [0, subgroupCount) is written (each subgroup elects one leader).
-// Sized for the worst case of one partial per invocation.
 var<workgroup> partialM: array<f32, WG>;
 var<workgroup> partialD: array<f32, WG>;
 var<workgroup> combinedMD: vec2<f32>;
 // When the whole workgroup is one subgroup the subgroup reduce already covers
-// it (no barriers, no shared state); otherwise subgroup leaders publish
-// partials through shared memory and thread 0 folds them in ordinal order.
 fn combine_partials(m: f32, d: f32, lidx: u32, sgSize: u32) -> vec2<f32> {
   let sgM = subgroupMax(m);
   // A lane with no elements contributes d == 0 (exact identity). A +inf
@@ -111,31 +106,24 @@ fn combine_partials(m: f32, d: f32, lidx: u32, sgSize: u32) -> vec2<f32> {
   if (sgSize == WG) {
     return vec2<f32>(sgM, sgD);
   }
-  let subgroupCount = (WG + sgSize - 1u) / sgSize;
-  // Pre-seed every fold slot with the (max, denom) identity. The fold below reads a
-  // fixed subgroupCount slots in ordinal order (for determinism), but a slot whose
-  // subgroup elects no leader this call — e.g. a fully out-of-window key tile in the
-  // flash-attention loop that re-uses this shared memory each iteration — would
-  // otherwise read stale shared memory. Identity makes such a slot a no-op.
-  // (max identity = -FLT_MAX, denom identity = 0.)
-  if (lidx < subgroupCount) {
-    partialM[lidx] = -FLT_MAX;
-    partialD[lidx] = 0.0;
-  }
-  workgroupBarrier();
   if (subgroupElect()) {
-    let slot = lidx / sgSize;
-    partialM[slot] = sgM;
-    partialD[slot] = sgD;
   }
   workgroupBarrier();
   if (lidx == 0u) {
     var accM = -FLT_MAX;
     var accD = 0.0;
-    for (var i = 0u; i < subgroupCount; i = i + 1u) {
-      let mNew = max(accM, partialM[i]);
-      accD = accD * exp_shift(accM, mNew) + partialD[i] * exp_shift(partialM[i], mNew);
-      accM = mNew;
     }
     combinedMD = vec2<f32>(accM, accD);
   }
@@ -150,9 +138,8 @@ var<workgroup> partialM: array<f32, {{ mdExtent }}>;
 var<workgroup> partialD: array<f32, {{ mdExtent }}>;
 {% if mdStreamed %}
-// In-place fold of {{ mdStreams }} streams. The caller stores its per-thread
-// partials into partialM/partialD first and reads the merged pair of stream s
-// from slot s * WG afterwards.
 fn combine_partials_streams(lidx: u32) {
   workgroupBarrier();
   var stride = WG / 2u;

+{% set splitK = splitK is defined and splitK %}
 {% set keyLo = "splitStart" if splitK else "0u" %}
 {% set keyHi = "splitEnd" if splitK else "key_bound" %}
 enable f16;
 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.PagedAttention, attention pass.
+// One workgroup per packed token and KV head sweeps the token's history in
+// WG-key tiles and merges them into an online softmax without materializing the
+// score row. Query heads sharing that KV head reuse each resolved cache row.
+// Key j of sequence b resides in block_table[b, j / BLOCK_SIZE] at offset
+// j % BLOCK_SIZE; each tile resolves this mapping once per key and stores the
+// row bases in workgroup memory. The exclusive key bound is the token's absolute
+// position plus one, so later keys are never visited.
 {% if splitK %}
 // Split-K partitions the key axis into NUM_SPLITS contiguous ranges. Each range
 // writes its unnormalized online state (max, denominator, and weighted V sum),
 // calls before their shared partial storage is reused.
 {% set combineSubgroups = combineSubgroups is defined and combineSubgroups %}
 {% if combineSubgroups %}
+// Cross-subgroup merge that assumes nothing about which invocations share a
+// subgroup or how many subgroups there are: each subgroup's elected lane
+// publishes the subgroup pair in the slot at its OWN invocation index and sets
+// that index's bit in a workgroup bitmask; thread 0 then folds exactly the
+// published slots, in ascending index order (the online (m, d) merge is not
+// float-associative, so the order is fixed), and clears the mask for the next
+// call as it reads it. Workgroup memory starts zeroed, so the mask needs no
+// setup. Same three collectives as a single-subgroup reduce, two barriers.
 var<workgroup> partialM: array<f32, WG>;
 var<workgroup> partialD: array<f32, WG>;
+var<workgroup> leaderMask: array<atomic<u32>, (WG + 31u) / 32u>;
 var<workgroup> combinedMD: vec2<f32>;
 // When the whole workgroup is one subgroup the subgroup reduce already covers
+// it (no barriers, no shared state). `subgroup_size` is the size of the current
+// subgroup and uniform, so the test is exact and may guard the barriers below.
 fn combine_partials(m: f32, d: f32, lidx: u32, sgSize: u32) -> vec2<f32> {
   let sgM = subgroupMax(m);
   // A lane with no elements contributes d == 0 (exact identity). A +inf
   if (sgSize == WG) {
     return vec2<f32>(sgM, sgD);
   }
   if (subgroupElect()) {
+    partialM[lidx] = sgM;
+    partialD[lidx] = sgD;
+    atomicOr(&leaderMask[lidx / 32u], 1u << (lidx % 32u));
   }
   workgroupBarrier();
   if (lidx == 0u) {
     var accM = -FLT_MAX;
     var accD = 0.0;
+    for (var w = 0u; w < (WG + 31u) / 32u; w = w + 1u) {
+      var bits = atomicExchange(&leaderMask[w], 0u);
+      while (bits != 0u) {
+        let slot = w * 32u + firstTrailingBit(bits);
+        bits = bits & (bits - 1u);
+        let mNew = max(accM, partialM[slot]);
+        accD = accD * exp_shift(accM, mNew) + partialD[slot] * exp_shift(partialM[slot], mNew);
+        accM = mNew;
+      }
     }
     combinedMD = vec2<f32>(accM, accD);
   }
 var<workgroup> partialD: array<f32, {{ mdExtent }}>;
 {% if mdStreamed %}
+// In-place fold of {{ mdStreams }} streams. Input partials occupy
+// partialM/partialD; stream s returns its merged pair in slot s * WG.
 fn combine_partials_streams(lidx: u32) {
   workgroupBarrier();
   var stride = WG / 2u;

build/webgpu/paged-scatter-kv.wgsl.jinja CHANGED Viewed

@@ -1,30 +1,30 @@
 {% macro flat_index_2d(name="i", bound="params.count", guardInline=false, note="dispatch-limit") %}
 {% if note == "dispatch-limit" %}
-  // 2D-folded flat index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension dispatch limit (outputs > 16.7M elements).
 {% elif note == "limit" %}
-  // 2D-folded flat index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension limit.
 {% elif note == "device-axis" %}
-  // The flat dispatch is folded across x/y at the device's per-axis workgroup
-  // limit; gid.y carries the high portion of the output index.
 {% elif note == "vec4-limit" %}
-  // 2D-folded flat vec4 index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension limit (the dispatch caps x and spills into y).
 {% elif note == "element-limit" %}
   // 2D-folded flat element index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension limit.
 {% elif note == "dispatch" %}
-  // 2D-folded flat index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension dispatch limit.
 {% endif %}
 {% if bound == "" %}
-  let {{ name }} = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
 {%- elif guardInline %}
-  let {{ name }} = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
   if ({{ name }} >= {{ bound }}) { return; }
 {%- else %}
-  let {{ name }} = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
   if ({{ name }} >= {{ bound }}) {
     return;
   }
@@ -80,16 +80,15 @@ fn paged_row_base(block_id: u32, slot_in_block: u32, kv_head: u32) -> u32 {
 {{ paged_sequence_index(withLookup=not hasSlotMapping) }}
 @compute @workgroup_size({{ tunables.SCATTER_WORKGROUP_SIZE }})
-fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
 {{ flat_index_2d(bound="params.scatterCount", note="element-limit") }}
   let dim = i % HEAD_DIM;
   let kv_head = (i / HEAD_DIM) % KV_HEADS;
   let token = i / (KV_HEADS * HEAD_DIM);
 {% if hasSlotMapping %}
-  // slot_mapping gives the flat cache slot directly, and -1 means "do not store
-  // this token" -- the scheduler uses it to suppress writes for a prefix-cache
-  // hit or a rejected speculative token. block_table still defines the read path.
   let slot = slot_mapping[token];
   if (slot < 0) {
     return;

 {% macro flat_index_2d(name="i", bound="params.count", guardInline=false, note="dispatch-limit") %}
 {% if note == "dispatch-limit" %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width (outputs > 16.7M elements).
 {% elif note == "limit" %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
 {% elif note == "device-axis" %}
+  // The flat dispatch is folded across x/y at a fixed per-axis workgroup
+  // width; gid.y carries the high portion of the output index.
 {% elif note == "vec4-limit" %}
+  // 2D-folded flat vec4 index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width (the dispatch caps x and spills into y).
 {% elif note == "element-limit" %}
   // 2D-folded flat element index: gid.y carries the high bits past the
+  // dispatch's per-axis workgroup fold width.
 {% elif note == "dispatch" %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
 {% endif %}
 {% if bound == "" %}
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
 {%- elif guardInline %}
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
   if ({{ name }} >= {{ bound }}) { return; }
 {%- else %}
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
   if ({{ name }} >= {{ bound }}) {
     return;
   }
 {{ paged_sequence_index(withLookup=not hasSlotMapping) }}
 @compute @workgroup_size({{ tunables.SCATTER_WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 {{ flat_index_2d(bound="params.scatterCount", note="element-limit") }}
   let dim = i % HEAD_DIM;
   let kv_head = (i / HEAD_DIM) % KV_HEADS;
   let token = i / (KV_HEADS * HEAD_DIM);
 {% if hasSlotMapping %}
+  // `slot_mapping` gives the flat cache slot directly. A negative slot suppresses
+  // this token's cache write; `block_table` still defines the read path.
   let slot = slot_mapping[token];
   if (slot < 0) {
     return;

build/webgpu/test.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "com.microsoft.PagedAttention",
   "cases": [
     {
       "name": "two_sequences_gqa",
@@ -45,7 +44,9 @@
     },
     {
       "name": "two_sequences_packed_qkv",
-      "provenance": { "notes": "Same schedule with Q, K and V packed on one query row." },
       "attrs": { "num_heads": 4, "kv_num_heads": 2 },
       "inputs": {
         "queryT": {
@@ -76,7 +77,7 @@
     {
       "name": "two_sequences_key_cache_output",
       "provenance": {
-        "notes": "Requests only the updated key-cache alias; the paired schema value-cache output remains internal."
       },
       "attrs": { "num_heads": 4, "kv_num_heads": 2 },
       "inputs": {
@@ -206,7 +207,7 @@
     {
       "name": "slot_mapping_skips_token",
       "provenance": {
-        "notes": "A slot of -1 suppresses that token's cache write, which a scheduler uses for a prefix-cache hit or a rejected speculative token. The cache slot keeps its prior contents and attention still reads it."
       },
       "attrs": { "num_heads": 4, "kv_num_heads": 2 },
       "inputs": {
@@ -520,7 +521,7 @@
     {
       "name": "pinned_paged_history",
       "provenance": {
-        "notes": "Expected values computed by an independent implementation written from the ONNX Runtime schema text alone, with every tensor value rounded to its declared float16 dtype before evaluation, so this case checks the trusted reference as well as the kernels. The block table [2, 0, 1] is deliberately not the identity, so a kernel that ignored it would read the wrong blocks; the two new tokens land in blocks 0 and 1 respectively."
       },
       "attrs": { "num_heads": 2, "kv_num_heads": 1 },
       "inputs": {
@@ -685,7 +686,7 @@
     {
       "name": "splitk_slot_mapping",
       "provenance": {
-        "notes": "The split-K schedule with an explicit slot mapping instead of a derived write slot: same 201-key history over a 256-key block table, so the two split ranges both run and the scatter pass takes the slot-mapping arm."
       },
       "attrs": { "num_heads": 2, "kv_num_heads": 1 },
       "inputs": {
@@ -766,7 +767,7 @@
     {
       "name": "splitk_packed_qkv_slot_mapping",
       "provenance": {
-        "notes": "The split-K schedule with both a packed query row and an explicit slot mapping, the last of the four scatter/query combinations the family covers."
       },
       "attrs": { "num_heads": 2, "kv_num_heads": 1 },
       "inputs": {
@@ -798,6 +799,49 @@
         "slotMappingT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [208] } }
       },
       "outputs": { "outputT": { "dtype": "float16", "shape": [1, 16], "tolerance": 0.00001, "relTolerance": 0.0001 } }
     }
   ]
 }

 {
   "cases": [
     {
       "name": "two_sequences_gqa",
     },
     {
       "name": "two_sequences_packed_qkv",
+      "provenance": {
+        "notes": "Q, K, and V are packed into one query row while preserving the two-sequence paged schedule."
+      },
       "attrs": { "num_heads": 4, "kv_num_heads": 2 },
       "inputs": {
         "queryT": {
     {
       "name": "two_sequences_key_cache_output",
       "provenance": {
+        "notes": "Only the updated key-cache alias is requested; the optional value-cache result is omitted."
       },
       "attrs": { "num_heads": 4, "kv_num_heads": 2 },
       "inputs": {
     {
       "name": "slot_mapping_skips_token",
       "provenance": {
+        "notes": "A negative slot suppresses that token's cache write. The cache slot retains its prior contents and remains available to attention."
       },
       "attrs": { "num_heads": 4, "kv_num_heads": 2 },
       "inputs": {
     {
       "name": "pinned_paged_history",
       "provenance": {
+        "notes": "Expected values are independently derived from the ONNX Runtime schema after rounding every tensor to its declared float16 dtype. A nonidentity block table [2, 0, 1] places the two new tokens in physical blocks 0 and 1."
       },
       "attrs": { "num_heads": 2, "kv_num_heads": 1 },
       "inputs": {
     {
       "name": "splitk_slot_mapping",
       "provenance": {
+        "notes": "A 201-key history over a 256-key table creates two split-K ranges, while explicit slot mapping selects the corresponding cache-scatter path."
       },
       "attrs": { "num_heads": 2, "kv_num_heads": 1 },
       "inputs": {
     {
       "name": "splitk_packed_qkv_slot_mapping",
       "provenance": {
+        "notes": "A packed query row and explicit slot mapping exercise the split-K scatter path with both optional layout modes active."
       },
       "attrs": { "num_heads": 2, "kv_num_heads": 1 },
       "inputs": {
         "slotMappingT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [208] } }
       },
       "outputs": { "outputT": { "dtype": "float16", "shape": [1, 16], "tolerance": 0.00001, "relTolerance": 0.0001 } }
+    },
+    {
+      "name": "mqa_32_query_heads_per_kv_head_workgroup_storage",
+      "provenance": {
+        "notes": "A 32:1 multi-query decode with head size 64 stages 24,832 bytes of query, output, and probability data per workgroup. The case therefore declares a 32 KiB workgroup-storage requirement and must be rejected as inapplicable at the WebGPU 16 KiB minimum."
+      },
+      "requires": { "limits": { "maxComputeWorkgroupStorageSize": 32768 } },
+      "attrs": { "num_heads": 32, "kv_num_heads": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [1, 2048],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.13, "scale": 0.8 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.41, "cosStep": 0.07, "scale": 0.7 }
+        },
+        "keyCacheT": {
+          "dtype": "float16",
+          "shape": [8, 16, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 }
+        },
+        "valueCacheT": {
+          "dtype": "float16",
+          "shape": [8, 16, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.19, "scale": 0.5 }
+        },
+        "cumulativeSequenceLengthT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [0, 1] } },
+        "pastSeqlensT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [40] } },
+        "blockTableT": {
+          "dtype": "int32",
+          "shape": [1, 8],
+          "data": { "kind": "values", "values": [0, 1, 2, 3, 4, 5, 6, 7] }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float16", "shape": [1, 2048], "tolerance": 0.00001, "relTolerance": 0.0001 } }
     }
   ]
 }