sync 2e7068faf55e

Browse files

Files changed (8) hide show

README.md +95 -0
build/webgpu/attn-flash-decode-splitk-merge.wgsl.jinja +113 -0
build/webgpu/bench.json +156 -0
build/webgpu/manifest.json +792 -0
build/webgpu/metadata.json +20 -0
build/webgpu/paged-attention.wgsl.jinja +419 -0
build/webgpu/paged-scatter-kv.wgsl.jinja +121 -0
build/webgpu/test.json +803 -0

README.md CHANGED Viewed

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

 ---
+library_name: kernels
 license: apache-2.0
+tags:
+- kernel
+- webgpu
+- wgsl
 ---
+# com.microsoft.PagedAttention
+`com.microsoft`  ·  ONNX Runtime contrib operator  ·  contrib since_version 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.
+See the [ONNX Runtime `PagedAttention` contrib-operator spec](https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.PagedAttention) for the reference semantics.
+## 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
+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)`. |
+## Type constraints
+| Variable | Allowed dtypes |
+| --- | --- |
+| `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
+- [`attn-flash-decode-splitk-merge.wgsl.jinja`](build/webgpu/attn-flash-decode-splitk-merge.wgsl.jinja)
+- [`paged-attention.wgsl.jinja`](build/webgpu/paged-attention.wgsl.jinja)
+- [`paged-scatter-kv.wgsl.jinja`](build/webgpu/paged-scatter-kv.wgsl.jinja)
+## Use with `@huggingface/kernels`
+The loader derives every required output's shape and logical dtype from the manifest contract and this call.
+It then allocates the result tensors automatically.
+The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+Replace each `*Data` placeholder with a typed array containing the corresponding input data.
+```js
+import { getKernel } from "@huggingface/kernels";
+const kernel = await getKernel("webgpu-kernels/com.microsoft.PagedAttention", { version: 1 });
+const { keyCacheT, valueCacheT, outputT } = await kernel({
+  queryT: { data: queryTData, shape: [2, 4] },
+  keyT: { data: keyTData, shape: [2, 2] },
+  valueT: { data: valueTData, shape: [2, 2] },
+  keyCacheT: { data: keyCacheTData, shape: [3, 2, 1, 2] },
+  valueCacheT: { data: valueCacheTData, shape: [3, 2, 1, 2] },
+  cumulativeSequenceLengthT: { data: cumulativeSequenceLengthTData, shape: [2] },
+  pastSeqlensT: { data: pastSeqlensTData, shape: [1] },
+  blockTableT: { data: blockTableTData, shape: [1, 3] },
+}, {
+  attrs: { num_heads: 2, kv_num_heads: 1 },
+});
+```

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

	@@ -0,0 +1,113 @@

+{% if scalar == "f16" %}
+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
+// states (max, denom, weighted V sum) written by the split pass into the final
+// attention output for each query token, using the online softmax merge:
+//   globalMax  = max_s partialMax[s]
+//   globalDen  = sum_s partialDenom[s] * exp(partialMax[s] - globalMax)
+//   out[d]     = (sum_s partialOut[s, d] * exp(partialMax[s] - globalMax)) / globalDen
+// One workgroup per (batch, head); one thread per output vec4 channel. The
+// per-split (max, denom) scan is tiny (NUM_SPLITS is small), so every thread
+// recomputes it locally and keeps the merge state in registers.
+// Output layout: token-major [batch, seq, heads * headDim], matching the split pass.
+const HEAD_DIM_V4: u32 = {{ headDimV4 }}u;
+const Q_HEADS: u32 = {{ qNumHeads }}u;
+const NUM_SPLITS: u32 = {{ numSplits }}u;
+{% if splitQueries %}
+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
+// 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));
+}
+@compute @workgroup_size(HEAD_DIM_V4, 1, 1)
+fn main(
+  @builtin(workgroup_id) wg: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>
+) {
+{% if splitQueries %}
+  let queryToken = wg.x;
+{% endif %}
+  let h = wg.y;
+  let b = wg.z;
+  let d4 = lid.x;
+  if (h >= Q_HEADS || d4 >= HEAD_DIM_V4{% if splitQueries %} || queryToken >= Q_SEQ{% endif %}) {
+    return;
+  }
+{% if splitQueries %}
+  let mdBase = ((b * Q_SEQ + queryToken) * Q_HEADS + h) * NUM_SPLITS;
+{% else %}
+  let mdBase = (b * Q_HEADS + h) * NUM_SPLITS;
+{% endif %}
+  var globalMax = -FLT_MAX;
+  for (var s: u32 = 0u; s < NUM_SPLITS; s = s + 1u) {
+    globalMax = max(globalMax, partial_stats[mdBase + s].x);
+  }
+  var globalDenom = 0.0;
+  var outV = vec4<f32>(0.0);
+  for (var s: u32 = 0u; s < NUM_SPLITS; s = s + 1u) {
+    let stats = partial_stats[mdBase + s];
+    let w = exp_shift(stats.x, globalMax);
+    globalDenom = globalDenom + stats.y * w;
+{% if splitQueries %}
+    let pBase = (((b * Q_SEQ + queryToken) * Q_HEADS + h) * NUM_SPLITS + s) * HEAD_DIM_V4;
+{% else %}
+    let pBase = ((b * Q_HEADS + h) * NUM_SPLITS + s) * HEAD_DIM_V4;
+{% endif %}
+    outV = outV + partial_out[pBase + d4] * w;
+  }
+{% if splitQueries %}
+  let qBaseV4 = (b * Q_SEQ + queryToken) * Q_HIDDEN_V4 + h * HEAD_DIM_V4;
+{% else %}
+  let qBaseV4 = b * Q_HIDDEN_V4 + h * HEAD_DIM_V4;
+{% endif %}
+  // A split that covers no keys writes vec2(-FLT_MAX, 0), so when every split is empty —
+  // a query whose window admits nothing — globalMax stays -FLT_MAX, exp_shift(x, x) is 1,
+  // 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 %}
+  output[qBaseV4 + d4] = vec4<{{ scalar }}>(outValue);
+}

build/webgpu/bench.json ADDED Viewed

	@@ -0,0 +1,156 @@

+{
+  "op": "com.microsoft.PagedAttention",
+  "tunableSpace": {
+    "WORKGROUP_SIZE": [32, 64, 128],
+    "SCATTER_WORKGROUP_SIZE": [64, 128, 256],
+    "SPLIT_TARGET_WORKGROUPS": [512, 1024, 2048, 4096],
+    "SPLIT_MIN_KEYS": [64, 128, 256],
+    "MAX_SPLITS": [8, 16, 32]
+  },
+  "cases": [
+    {
+      "name": "paged-decode-b8-past1024-h32kv8-d128",
+      "preset": "smoke",
+      "vars": { "dtype": "float16", "batch": 8, "past": 1024, "heads": 32, "kvHeads": 8, "headDim": 128 },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8 },
+      "inputs": {
+        "queryT": { "shape": [8, 4096], "dtype": "float16", "dist": "normal", "seed": 9201, "scale": 1 },
+        "keyT": { "shape": [8, 1024], "dtype": "float16", "dist": "normal", "seed": 9202, "scale": 1 },
+        "valueT": { "shape": [8, 1024], "dtype": "float16", "dist": "normal", "seed": 9203, "scale": 1 },
+        "keyCacheT": { "shape": [512, 16, 8, 128], "dtype": "float16", "dist": "normal", "seed": 9204, "scale": 1 },
+        "valueCacheT": { "shape": [512, 16, 8, 128], "dtype": "float16", "dist": "normal", "seed": 9205, "scale": 1 },
+        "cumulativeSequenceLengthT": {
+          "shape": [9],
+          "dtype": "int32",
+          "data": { "kind": "values", "values": [0, 1, 2, 3, 4, 5, 6, 7, 8] }
+        },
+        "pastSeqlensT": { "shape": [8], "dtype": "int32", "dist": "constant", "value": 1023 },
+        "blockTableT": { "shape": [8, 64], "dtype": "int32", "dist": "linearMod", "seed": 9208, "step": 1, "mod": 512 }
+      },
+      "outputs": { "outputT": { "shape": [8, 4096], "dtype": "float16" } },
+      "bench": {
+        "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."
+      }
+    },
+    {
+      "name": "paged-decode-llama-b16-past2048-h32kv32-d128",
+      "preset": "model",
+      "provenance": {
+        "notes": "Llama class defaults (32 query heads, no GQA, head_dim 128) decoding 16 concurrent sequences with 2048 tokens of paged KV each."
+      },
+      "vars": { "dtype": "float16", "batch": 16, "past": 2048, "heads": 32, "kvHeads": 32, "headDim": 128 },
+      "attrs": { "num_heads": 32, "kv_num_heads": 32 },
+      "inputs": {
+        "queryT": { "shape": [16, 4096], "dtype": "float16", "dist": "normal", "seed": 7600, "scale": 0.3 },
+        "keyT": { "shape": [16, 4096], "dtype": "float16", "dist": "normal", "seed": 7601, "scale": 0.3 },
+        "valueT": { "shape": [16, 4096], "dtype": "float16", "dist": "normal", "seed": 7602, "scale": 0.3 },
+        "keyCacheT": { "shape": [2048, 16, 32, 128], "dtype": "float16", "dist": "normal", "seed": 7603, "scale": 0.3 },
+        "valueCacheT": {
+          "shape": [2048, 16, 32, 128],
+          "dtype": "float16",
+          "dist": "normal",
+          "seed": 7604,
+          "scale": 0.3
+        },
+        "cumulativeSequenceLengthT": {
+          "shape": [17],
+          "dtype": "int32",
+          "dist": "linearMod",
+          "seed": 7605,
+          "step": 1,
+          "mod": 17
+        },
+        "pastSeqlensT": { "shape": [16], "dtype": "int32", "dist": "constant", "value": 2048 },
+        "blockTableT": {
+          "shape": [16, 128],
+          "dtype": "int32",
+          "dist": "linearMod",
+          "seed": 7606,
+          "step": 1,
+          "mod": 2048
+        }
+      },
+      "outputs": { "outputT": { "shape": [16, 4096], "dtype": "float16" } },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "2 * args.batch * args.kvHeads * args.past * args.headDim * 2" }]
+      }
+    },
+    {
+      "name": "paged-decode-mistral-b16-past4096-h32kv8-d128",
+      "preset": "model",
+      "provenance": {
+        "notes": "Mistral class defaults (32 query heads, 8 key/value heads, head_dim 128) at 4096 tokens of history -- the GQA case, where the cache is a quarter the size for the same context."
+      },
+      "vars": { "dtype": "float16", "batch": 16, "past": 4096, "heads": 32, "kvHeads": 8, "headDim": 128 },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8 },
+      "inputs": {
+        "queryT": { "shape": [16, 4096], "dtype": "float16", "dist": "normal", "seed": 7700, "scale": 0.3 },
+        "keyT": { "shape": [16, 1024], "dtype": "float16", "dist": "normal", "seed": 7701, "scale": 0.3 },
+        "valueT": { "shape": [16, 1024], "dtype": "float16", "dist": "normal", "seed": 7702, "scale": 0.3 },
+        "keyCacheT": { "shape": [4096, 16, 8, 128], "dtype": "float16", "dist": "normal", "seed": 7703, "scale": 0.3 },
+        "valueCacheT": { "shape": [4096, 16, 8, 128], "dtype": "float16", "dist": "normal", "seed": 7704, "scale": 0.3 },
+        "cumulativeSequenceLengthT": {
+          "shape": [17],
+          "dtype": "int32",
+          "dist": "linearMod",
+          "seed": 7705,
+          "step": 1,
+          "mod": 17
+        },
+        "pastSeqlensT": { "shape": [16], "dtype": "int32", "dist": "constant", "value": 4096 },
+        "blockTableT": {
+          "shape": [16, 256],
+          "dtype": "int32",
+          "dist": "linearMod",
+          "seed": 7706,
+          "step": 1,
+          "mod": 4096
+        }
+      },
+      "outputs": { "outputT": { "shape": [16, 4096], "dtype": "float16" } },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "2 * args.batch * args.kvHeads * args.past * args.headDim * 2" }]
+      }
+    },
+    {
+      "name": "paged-decode-gemma3-b8-past2048-h8kv4-d256",
+      "preset": "model",
+      "provenance": {
+        "notes": "Gemma3 class defaults (8 query heads, 4 key/value heads, head_dim 256) -- few heads but the widest head dimension in this set."
+      },
+      "vars": { "dtype": "float16", "batch": 8, "past": 2048, "heads": 8, "kvHeads": 4, "headDim": 256 },
+      "attrs": { "num_heads": 8, "kv_num_heads": 4 },
+      "inputs": {
+        "queryT": { "shape": [8, 2048], "dtype": "float16", "dist": "normal", "seed": 7800, "scale": 0.3 },
+        "keyT": { "shape": [8, 1024], "dtype": "float16", "dist": "normal", "seed": 7801, "scale": 0.3 },
+        "valueT": { "shape": [8, 1024], "dtype": "float16", "dist": "normal", "seed": 7802, "scale": 0.3 },
+        "keyCacheT": { "shape": [1024, 16, 4, 256], "dtype": "float16", "dist": "normal", "seed": 7803, "scale": 0.3 },
+        "valueCacheT": { "shape": [1024, 16, 4, 256], "dtype": "float16", "dist": "normal", "seed": 7804, "scale": 0.3 },
+        "cumulativeSequenceLengthT": {
+          "shape": [9],
+          "dtype": "int32",
+          "dist": "linearMod",
+          "seed": 7805,
+          "step": 1,
+          "mod": 9
+        },
+        "pastSeqlensT": { "shape": [8], "dtype": "int32", "dist": "constant", "value": 2048 },
+        "blockTableT": {
+          "shape": [8, 128],
+          "dtype": "int32",
+          "dist": "linearMod",
+          "seed": 7806,
+          "step": 1,
+          "mod": 1024
+        }
+      },
+      "outputs": { "outputT": { "shape": [8, 2048], "dtype": "float16" } },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "2 * args.batch * args.kvHeads * args.past * args.headDim * 2" }]
+      }
+    }
+  ]
+}

build/webgpu/manifest.json ADDED Viewed

	@@ -0,0 +1,792 @@

+{
+  "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)",
+    "blockSize": "dim(shapes.keyCacheT, 1)",
+    "headSize": "dim(shapes.keyCacheT, 3)",
+    "maxBlocks": "dim(shapes.blockTableT, 1)",
+    "batchSize": "dim(shapes.cumulativeSequenceLengthT, 0) - 1",
+    "qHidden": "attrs.num_heads * headSize",
+    "kvHidden": "attrs.kv_num_heads * headSize",
+    "qPerKv": "attrs.num_heads / max(1, attrs.kv_num_heads)",
+    "headVec": "headSize / 4",
+    "cacheVec4Ok": "headSize % 4 == 0",
+    "packedStride": "qHidden + 2 * kvHidden",
+    "cacheShapeOk": "ranks.keyCacheT == 4 and ranks.valueCacheT == 4 and sameShape(shapes.valueCacheT, shapes.keyCacheT) and dim(shapes.keyCacheT, 2) == attrs.kv_num_heads and blockSize > 0 and headSize > 0 and tensorDtypes.keyCacheT == tensorDtypes.queryT and tensorDtypes.valueCacheT == tensorDtypes.queryT",
+    "headLayoutOk": "attrs.num_heads > 0 and attrs.kv_num_heads > 0 and attrs.num_heads % attrs.kv_num_heads == 0",
+    "scheduleShapeOk": "ranks.cumulativeSequenceLengthT == 1 and ranks.pastSeqlensT == 1 and ranks.blockTableT == 2 and batchSize >= 1 and dim(shapes.pastSeqlensT, 0) == batchSize and dim(shapes.blockTableT, 0) == batchSize and maxBlocks >= 1",
+    "ioShapeOk": "ranks.queryT == 2 and ranks.outputT == 2 and dim(shapes.outputT, 0) == tokenCount and dim(shapes.outputT, 1) == qHidden and tensorDtypes.outputT == tensorDtypes.queryT and f16Ok(tensorDtypes.queryT)",
+    "separateKv": "present.keyT and present.valueT and ranks.keyT == 2 and ranks.valueT == 2 and dim(shapes.keyT, 0) == tokenCount and dim(shapes.valueT, 0) == tokenCount and dim(shapes.keyT, 1) == kvHidden and dim(shapes.valueT, 1) == kvHidden and tensorDtypes.keyT == tensorDtypes.queryT and tensorDtypes.valueT == tensorDtypes.queryT and dim(shapes.queryT, 1) == qHidden",
+    "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",
+    "qHiddenV4": "qHidden / 4",
+    "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]" }
+      ],
+      "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": "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]" }
+      ],
+      "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": "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]" }
+      ],
+      "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": "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]" }
+      ],
+      "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": "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" }
+        }
+      ]
+    }
+  ]
+}

build/webgpu/metadata.json ADDED Viewed

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

+{
+  "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" }
+}

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

	@@ -0,0 +1,419 @@

+{% 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),
+// and the merge pass combines those states before the final division. An empty
+// range writes the (-FLT_MAX, 0) identity expected by the merge.
+{% endif %}
+const KV_HEADS: u32 = {{ kvNumHeads }}u;
+const HEAD_DIM: u32 = {{ headSize }}u;
+// Where the head dimension divides into four, the cache is read four elements at
+// a time, reducing the number of dependent cache loads. Other head sizes use the
+// scalar walk, with the binding element type specialized to match.
+{% if cacheVec4Ok %}
+const HEAD_VEC: u32 = {{ headVec }}u;
+{% endif %}
+const BLOCK_SIZE: u32 = {{ blockSize }}u;
+const MAX_BLOCKS: u32 = {{ maxBlocks }}u;
+const Q_STRIDE: u32 = {{ packedStride if packedQkv else qHidden }}u;
+const Q_PER_KV: u32 = {{ qPerKv }}u;
+// WG is both the key-tile width and the thread count that walks the head axis.
+const WG: u32 = {{ attnWorkgroup }}u;
+{% if splitK %}
+const NUM_SPLITS: u32 = {{ numSplits }}u;
+const Q_HEADS: u32 = {{ numHeads }}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
+// 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));
+}
+{% if cacheVec4Ok %}
+var<workgroup> q_shared: array<vec4<f32>, HEAD_VEC * Q_PER_KV>;
+var<workgroup> running_out: array<vec4<f32>, HEAD_VEC * Q_PER_KV>;
+{% else %}
+var<workgroup> q_shared: array<f32, HEAD_DIM * Q_PER_KV>;
+var<workgroup> running_out: array<f32, HEAD_DIM * Q_PER_KV>;
+{% endif %}
+var<workgroup> probs: array<f32, WG * Q_PER_KV>;
+// One resolved cache row base per key of the current tile.
+var<workgroup> key_rows: array<u32, WG>;
+{% set mdStreams = qPerKv %}
+// Workgroup-cooperative merge of per-thread online-softmax (m, d) partials:
+// mNew = max(m1, m2)
+// dNew = d1 * exp(m1 - mNew) + d2 * exp(m2 - mNew)
+// Both the subgroup and portable barrier-tree engines return the same merged
+// pair to every invocation. Repeated merges require a workgroup barrier between
+// calls before their shared partial storage is reused.
+{% set combineSubgroups = combineSubgroups is defined and combineSubgroups %}
+{% if combineSubgroups %}
+// Per-subgroup partials are published into a deterministic slot: the subgroup's
+// ordinal index within the workgroup (lidx / sgSize). The online (m, d) merge
+// is not float-associative, so thread 0 must fold partials in a fixed order.
+// Subgroups partition a workgroup into contiguous ordinal ranges on supported
+// backends, so the ordinal slot is unique per subgroup and every slot in
+// [0, subgroupCount) is written (each subgroup elects one leader).
+// Sized for the worst case of one partial per invocation.
+var<workgroup> partialM: array<f32, WG>;
+var<workgroup> partialD: array<f32, WG>;
+var<workgroup> combinedMD: vec2<f32>;
+// When the whole workgroup is one subgroup the subgroup reduce already covers
+// it (no barriers, no shared state); otherwise subgroup leaders publish
+// partials through shared memory and thread 0 folds them in ordinal order.
+fn combine_partials(m: f32, d: f32, lidx: u32, sgSize: u32) -> vec2<f32> {
+  let sgM = subgroupMax(m);
+  // A lane with no elements contributes d == 0 (exact identity). A +inf
+  // element made exp(inf - inf) = NaN stick in that lane's d; a NaN element
+  // landed in d via exp(NaN); both survive the merge and are detected by the
+  // code after the reduction.
+  let sgD = subgroupAdd(d * exp_shift(m, sgM));
+  if (sgSize == WG) {
+    return vec2<f32>(sgM, sgD);
+  }
+  let subgroupCount = (WG + sgSize - 1u) / sgSize;
+  // Pre-seed every fold slot with the (max, denom) identity. The fold below reads a
+  // fixed subgroupCount slots in ordinal order (for determinism), but a slot whose
+  // subgroup elects no leader this call — e.g. a fully out-of-window key tile in the
+  // flash-attention loop that re-uses this shared memory each iteration — would
+  // otherwise read stale shared memory. Identity makes such a slot a no-op.
+  // (max identity = -FLT_MAX, denom identity = 0.)
+  if (lidx < subgroupCount) {
+    partialM[lidx] = -FLT_MAX;
+    partialD[lidx] = 0.0;
+  }
+  workgroupBarrier();
+  if (subgroupElect()) {
+    let slot = lidx / sgSize;
+    partialM[slot] = sgM;
+    partialD[slot] = sgD;
+  }
+  workgroupBarrier();
+  if (lidx == 0u) {
+    var accM = -FLT_MAX;
+    var accD = 0.0;
+    for (var i = 0u; i < subgroupCount; i = i + 1u) {
+      let mNew = max(accM, partialM[i]);
+      accD = accD * exp_shift(accM, mNew) + partialD[i] * exp_shift(partialM[i], mNew);
+      accM = mNew;
+    }
+    combinedMD = vec2<f32>(accM, accD);
+  }
+  workgroupBarrier();
+  return combinedMD;
+}
+{% else %}
+{% set mdStreamed = mdStreams is defined %}
+{% set mdStreams = mdStreams if mdStreams is defined else 1 %}
+{% set mdExtent = "WG" if mdStreams == 1 else "WG * " ~ mdStreams ~ "u" %}
+var<workgroup> partialM: array<f32, {{ mdExtent }}>;
+var<workgroup> partialD: array<f32, {{ mdExtent }}>;
+{% if mdStreamed %}
+// In-place fold of {{ mdStreams }} streams. The caller stores its per-thread
+// partials into partialM/partialD first and reads the merged pair of stream s
+// from slot s * WG afterwards.
+fn combine_partials_streams(lidx: u32) {
+  workgroupBarrier();
+  var stride = WG / 2u;
+  loop {
+    if (stride == 0u) {
+      break;
+    }
+    if (lidx < stride) {
+{% for s in range(mdStreams) %}
+      {
+        let slot = {{ s }}u * WG + lidx;
+        let m1 = partialM[slot];
+        let d1 = partialD[slot];
+        let m2 = partialM[slot + stride];
+        let d2 = partialD[slot + stride];
+        let mNew = max(m1, m2);
+        partialD[slot] = d1 * exp_shift(m1, mNew) + d2 * exp_shift(m2, mNew);
+        partialM[slot] = mNew;
+      }
+{% endfor %}
+    }
+    workgroupBarrier();
+    stride = stride / 2u;
+  }
+}
+{% else %}
+fn combine_partials(m: f32, d: f32, lidx: u32) -> vec2<f32> {
+  partialM[lidx] = m;
+  partialD[lidx] = d;
+  workgroupBarrier();
+  var stride = WG / 2u;
+  loop {
+    if (stride == 0u) {
+      break;
+    }
+    if (lidx < stride) {
+      let m1 = partialM[lidx];
+      let d1 = partialD[lidx];
+      let m2 = partialM[lidx + stride];
+      let d2 = partialD[lidx + stride];
+      let mNew = max(m1, m2);
+      partialD[lidx] = d1 * exp_shift(m1, mNew) + d2 * exp_shift(m2, mNew);
+      partialM[lidx] = mNew;
+    }
+    workgroupBarrier();
+    stride = stride / 2u;
+  }
+  let merged = vec2<f32>(partialM[0], partialD[0]);
+  // Trailing barrier so back-to-back calls cannot race a next call's partial
+  // stores against this call's reads of slot 0.
+  workgroupBarrier();
+  return merged;
+}
+{% endif %}
+{% endif %}
+{% if ATTN_SCALE_DIM is not defined %}{% set ATTN_SCALE_DIM = "HEAD_DIM" %}{% endif %}
+fn scale_value() -> f32 {
+  if (params.scale != 0.0) { return params.scale; }
+  return inverseSqrt(f32({{ ATTN_SCALE_DIM }}));
+}
+{% macro varlen_sequence_of_token(buffer="cumulative_sequence_length") %}
+fn sequence_of_token(token: u32, batch_size: u32) -> u32 {
+  for (var b: u32 = 0u; b < batch_size; b = b + 1u) {
+    if (token < u32({{ buffer }}[b + 1u])) {
+      return b;
+    }
+  }
+  return batch_size - 1u;
+}
+{%- endmacro %}
+{% macro paged_sequence_index(withLookup=true) %}
+{% if withLookup %}
+{{ varlen_sequence_of_token() }}
+{% endif %}
+// Physical element base of one (slot, kv head) row of the paged cache, whose
+// layout is (num_blocks, block_size, kv_num_heads, head_size).
+fn paged_row_base(block_id: u32, slot_in_block: u32, kv_head: u32) -> u32 {
+  return ((block_id * BLOCK_SIZE + slot_in_block) * KV_HEADS + kv_head) * HEAD_DIM;
+}
+{%- endmacro %}
+{{ paged_sequence_index() }}
+@compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
+  let token = wg.x;
+  let kv_head = wg.y;
+  let head0 = kv_head * Q_PER_KV;
+  let tid = lid.x;
+  let sequence = sequence_of_token(token, params.batchSize);
+  let local_token = token - u32(cumulative_sequence_length[sequence]);
+  // Keys 0 .. past + local_token are visible to this token; the +1 makes the
+  // bound exclusive.
+  let key_bound = u32(past_seqlens[sequence]) + local_token + 1u;
+  let scale = scale_value();
+{% if splitK %}
+  // Contiguous, equal-length key ranges. The last range is short and a range past
+  // this token's history is empty; both fall out of the tile loop below.
+  let split = wg.z;
+  let splitLen = (key_bound + NUM_SPLITS - 1u) / NUM_SPLITS;
+  let splitStart = split * splitLen;
+  let splitEnd = min(key_bound, splitStart + splitLen);
+{% endif %}
+{% if cacheVec4Ok %}
+  for (var dv: u32 = tid; dv < HEAD_VEC * Q_PER_KV; dv = dv + WG) {
+    let g = dv / HEAD_VEC;
+    let lane = dv - g * HEAD_VEC;
+    let qb = token * Q_STRIDE + (head0 + g) * HEAD_DIM + lane * 4u;
+    q_shared[dv] = vec4<f32>(f32(query[qb]), f32(query[qb + 1u]), f32(query[qb + 2u]), f32(query[qb + 3u]));
+    running_out[dv] = vec4<f32>(0.0);
+  }
+{% else %}
+  for (var d: u32 = tid; d < HEAD_DIM * Q_PER_KV; d = d + WG) {
+    let g = d / HEAD_DIM;
+    let dim = d - g * HEAD_DIM;
+    q_shared[d] = f32(query[token * Q_STRIDE + (head0 + g) * HEAD_DIM + dim]);
+    running_out[d] = 0.0;
+  }
+{% endif %}
+  workgroupBarrier();
+{% for g in range(qPerKv) %}
+  var runningMax_{{ g }} = -FLT_MAX;
+  var runningDenom_{{ g }} = 0.0;
+{% endfor %}
+  var tileBase: u32 = {{ keyLo }};
+  loop {
+    if (tileBase >= {{ keyHi }}) {
+      break;
+    }
+    let kj = tileBase + tid;
+    let keyAllowed = kj < {{ keyHi }};
+{% for g in range(qPerKv) %}
+    var score_{{ g }} = -FLT_MAX;
+{% endfor %}
+    if (keyAllowed) {
+      let block_id = u32(block_table[sequence * MAX_BLOCKS + kj / BLOCK_SIZE]);
+      let row = paged_row_base(block_id, kj % BLOCK_SIZE, kv_head);
+      key_rows[tid] = row;
+{% for g in range(qPerKv) %}
+      var acc_{{ g }} = 0.0;
+{% endfor %}
+      // One cache read, every query head in the group.
+{% if cacheVec4Ok %}
+      let row_v = row / 4u;
+      for (var dv: u32 = 0u; dv < HEAD_VEC; dv = dv + 1u) {
+        let kv = vec4<f32>(key_cache[row_v + dv]);
+{% for g in range(qPerKv) %}
+        acc_{{ g }} = acc_{{ g }} + dot(q_shared[{{ g }}u * HEAD_VEC + dv], kv);
+{% endfor %}
+      }
+{% else %}
+      for (var d: u32 = 0u; d < HEAD_DIM; d = d + 1u) {
+        let kv = f32(key_cache[row + d]);
+{% for g in range(qPerKv) %}
+        acc_{{ g }} = acc_{{ g }} + q_shared[{{ g }}u * HEAD_DIM + d] * kv;
+{% endfor %}
+      }
+{% endif %}
+{% for g in range(qPerKv) %}
+      score_{{ g }} = acc_{{ g }} * scale;
+{% endfor %}
+    }
+    // All Q_PER_KV reductions ride one barrier tree.
+{% for g in range(qPerKv) %}
+    partialM[{{ g }}u * WG + tid] = score_{{ g }};
+    partialD[{{ g }}u * WG + tid] = select(0.0, 1.0, keyAllowed);
+{% endfor %}
+    combine_partials_streams(tid);
+{% for g in range(qPerKv) %}
+    let tile_{{ g }} = vec2<f32>(partialM[{{ g }}u * WG], partialD[{{ g }}u * WG]);
+{% endfor %}
+    workgroupBarrier();
+{% for g in range(qPerKv) %}
+    let newMax_{{ g }} = max(runningMax_{{ g }}, tile_{{ g }}.x);
+    let correction_{{ g }} = exp_shift(runningMax_{{ g }}, newMax_{{ g }});
+    runningDenom_{{ g }} = runningDenom_{{ g }} * correction_{{ g }} + tile_{{ g }}.y * exp_shift(tile_{{ g }}.x, newMax_{{ g }});
+    runningMax_{{ g }} = newMax_{{ g }};
+    probs[{{ g }}u * WG + tid] = select(0.0, exp_shift(score_{{ g }}, newMax_{{ g }}), keyAllowed);
+{% endfor %}
+    workgroupBarrier();
+    // running_out[d] is owned by the same thread across every tile (tid = d mod
+    // WG), so this rescale-and-accumulate needs no further synchronization.
+    let tileCount = min(WG, {{ keyHi }} - tileBase);
+{% set W = "HEAD_VEC" if cacheVec4Ok else "HEAD_DIM" %}
+    for (var dv: u32 = tid; dv < {{ W }}; dv = dv + WG) {
+{% for g in range(qPerKv) %}
+      var vSum_{{ g }} = {{ "vec4<f32>(0.0)" if cacheVec4Ok else "0.0" }};
+{% endfor %}
+      // One cache read, every query head in the group, exactly as the scores.
+      for (var i: u32 = 0u; i < tileCount; i = i + 1u) {
+{% if cacheVec4Ok %}
+        let vv = vec4<f32>(value_cache[key_rows[i] / 4u + dv]);
+{% else %}
+        let vv = f32(value_cache[key_rows[i] + dv]);
+{% endif %}
+{% for g in range(qPerKv) %}
+        vSum_{{ g }} = vSum_{{ g }} + probs[{{ g }}u * WG + i] * vv;
+{% endfor %}
+      }
+{% for g in range(qPerKv) %}
+      running_out[{{ g }}u * {{ W }} + dv] = running_out[{{ g }}u * {{ W }} + dv] * correction_{{ g }} + vSum_{{ g }};
+{% endfor %}
+    }
+    // Orders this tile's probs and key_rows reads before the next tile writes them.
+    workgroupBarrier();
+    tileBase = tileBase + WG;
+  }
+{% if splitK %}
+  // This range's un-normalized state, in the layout the shared merge reads. The
+  // divide by the denominator is the merge's job: it belongs after every range
+  // has been folded, not once per range. Split-K always uses the vec4 cache.
+{% for g in range(qPerKv) %}
+  {
+    let slot_{{ g }} = (token * Q_HEADS + head0 + {{ g }}u) * NUM_SPLITS + split;
+    let pBase_{{ g }} = slot_{{ g }} * HEAD_VEC;
+    for (var dv: u32 = tid; dv < HEAD_VEC; dv = dv + WG) {
+      partial_out[pBase_{{ g }} + dv] = running_out[{{ g }}u * HEAD_VEC + dv];
+    }
+    if (tid == 0u) {
+      partial_stats[slot_{{ g }}] = vec2<f32>(runningMax_{{ g }}, runningDenom_{{ g }});
+    }
+  }
+{% endfor %}
+}
+{% else %}
+  // A token always attends at least its own key, so the denominator is positive
+  // unless every score was non-finite; guard anyway rather than emit NaN.
+{% for g in range(qPerKv) %}
+  let invDenom_{{ g }} = select(0.0, 1.0 / runningDenom_{{ g }}, runningDenom_{{ g }} > 0.0);
+{% endfor %}
+  for (var dv: u32 = tid; dv < {{ W }}; dv = dv + WG) {
+{% for g in range(qPerKv) %}
+    {
+      let out_v = running_out[{{ g }}u * {{ W }} + dv] * invDenom_{{ g }};
+{% if cacheVec4Ok %}
+      let ob = token * {{ qHidden }}u + (head0 + {{ g }}u) * HEAD_DIM + dv * 4u;
+      output[ob] = {{ scalar }}(out_v.x);
+      output[ob + 1u] = {{ scalar }}(out_v.y);
+      output[ob + 2u] = {{ scalar }}(out_v.z);
+      output[ob + 3u] = {{ scalar }}(out_v.w);
+{% else %}
+      output[token * {{ qHidden }}u + (head0 + {{ g }}u) * HEAD_DIM + dv] = {{ scalar }}(out_v);
+{% endif %}
+    }
+{% endfor %}
+  }
+}
+{%- endif %}

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

	@@ -0,0 +1,121 @@

+{% 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;
+  }
+{%- endif %}
+{% endmacro %}
+enable f16;
+{{ env.wgsl.resourceDeclarations }}
+// com.microsoft.PagedAttention, KV scatter pass.
+// Writes this step's keys and values into the block-based cache before the
+// attention pass reads it. One invocation moves one (token, kv head, dim)
+// element of both caches.
+const KV_HEADS: u32 = {{ kvNumHeads }}u;
+const HEAD_DIM: u32 = {{ headSize }}u;
+const BLOCK_SIZE: u32 = {{ blockSize }}u;
+{% if not hasSlotMapping %}
+const MAX_BLOCKS: u32 = {{ maxBlocks }}u;
+{% endif %}
+{% if packedQkv %}
+// Packed layout: query carries [Q | K | V] on one row, so K starts after the
+// query heads and V after the key heads.
+const PACKED_STRIDE: u32 = {{ packedStride }}u;
+const PACKED_K_OFFSET: u32 = {{ qHidden }}u;
+const PACKED_V_OFFSET: u32 = {{ qHidden }}u + {{ kvHidden }}u;
+{% else %}
+const KV_STRIDE: u32 = {{ kvHidden }}u;
+{% endif %}
+{% macro varlen_sequence_of_token(buffer="cumulative_sequence_length") %}
+fn sequence_of_token(token: u32, batch_size: u32) -> u32 {
+  for (var b: u32 = 0u; b < batch_size; b = b + 1u) {
+    if (token < u32({{ buffer }}[b + 1u])) {
+      return b;
+    }
+  }
+  return batch_size - 1u;
+}
+{%- endmacro %}
+{% macro paged_sequence_index(withLookup=true) %}
+{% if withLookup %}
+{{ varlen_sequence_of_token() }}
+{% endif %}
+// Physical element base of one (slot, kv head) row of the paged cache, whose
+// layout is (num_blocks, block_size, kv_num_heads, head_size).
+fn paged_row_base(block_id: u32, slot_in_block: u32, kv_head: u32) -> u32 {
+  return ((block_id * BLOCK_SIZE + slot_in_block) * KV_HEADS + kv_head) * HEAD_DIM;
+}
+{%- endmacro %}
+{{ 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;
+  }
+  let block_id = u32(slot) / BLOCK_SIZE;
+  let slot_in_block = u32(slot) % BLOCK_SIZE;
+{% else %}
+  // Derived slot: the token's absolute position in its sequence's cache is the
+  // past length plus its offset among this step's tokens.
+  let sequence = sequence_of_token(token, params.batchSize);
+  let local_token = token - u32(cumulative_sequence_length[sequence]);
+  let absolute_slot = u32(past_seqlens[sequence]) + local_token;
+  let block_id = u32(block_table[sequence * MAX_BLOCKS + absolute_slot / BLOCK_SIZE]);
+  let slot_in_block = absolute_slot % BLOCK_SIZE;
+{% endif %}
+{% if packedQkv %}
+  let source = token * PACKED_STRIDE + kv_head * HEAD_DIM + dim;
+  let key_value = query[source + PACKED_K_OFFSET];
+  let value_value = query[source + PACKED_V_OFFSET];
+{% else %}
+  let source = token * KV_STRIDE + kv_head * HEAD_DIM + dim;
+  let key_value = key[source];
+  let value_value = value[source];
+{% endif %}
+  let destination = paged_row_base(block_id, slot_in_block, kv_head) + dim;
+  key_cache[destination] = key_value;
+  value_cache[destination] = value_value;
+}

build/webgpu/test.json ADDED Viewed

	@@ -0,0 +1,803 @@

+{
+  "op": "com.microsoft.PagedAttention",
+  "cases": [
+    {
+      "name": "two_sequences_gqa",
+      "provenance": {
+        "notes": "Two packed sequences with different histories. The first sequence's two new tokens land at absolute slots 3 and 4, so the scatter crosses a block boundary and the attention reads two physical blocks."
+      },
+      "attrs": { "num_heads": 4, "kv_num_heads": 2 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [5, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [5, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.13, "scale": 0.8 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [5, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.41, "cosStep": 0.07, "scale": 0.7 }
+        },
+        "keyCacheT": {
+          "dtype": "float16",
+          "shape": [6, 4, 2, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 }
+        },
+        "valueCacheT": {
+          "dtype": "float16",
+          "shape": [6, 4, 2, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.19, "scale": 0.5 }
+        },
+        "cumulativeSequenceLengthT": {
+          "dtype": "int32",
+          "shape": [3],
+          "data": { "kind": "values", "values": [0, 2, 5] }
+        },
+        "pastSeqlensT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [3, 0] } },
+        "blockTableT": { "dtype": "int32", "shape": [2, 3], "data": { "kind": "values", "values": [0, 1, 2, 3, 4, 5] } }
+      },
+      "outputs": { "outputT": { "dtype": "float16", "shape": [5, 32], "tolerance": 0.00001, "relTolerance": 0.0001 } }
+    },
+    {
+      "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": {
+          "dtype": "float16",
+          "shape": [5, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 }
+        },
+        "keyCacheT": {
+          "dtype": "float16",
+          "shape": [6, 4, 2, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 }
+        },
+        "valueCacheT": {
+          "dtype": "float16",
+          "shape": [6, 4, 2, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.19, "scale": 0.5 }
+        },
+        "cumulativeSequenceLengthT": {
+          "dtype": "int32",
+          "shape": [3],
+          "data": { "kind": "values", "values": [0, 2, 5] }
+        },
+        "pastSeqlensT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [3, 0] } },
+        "blockTableT": { "dtype": "int32", "shape": [2, 3], "data": { "kind": "values", "values": [0, 1, 2, 3, 4, 5] } }
+      },
+      "outputs": { "outputT": { "dtype": "float16", "shape": [5, 32], "tolerance": 0.00001, "relTolerance": 0.0001 } }
+    },
+    {
+      "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": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [5, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [5, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.13, "scale": 0.8 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [5, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.41, "cosStep": 0.07, "scale": 0.7 }
+        },
+        "keyCacheT": {
+          "dtype": "float16",
+          "shape": [6, 4, 2, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 }
+        },
+        "valueCacheT": {
+          "dtype": "float16",
+          "shape": [6, 4, 2, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.19, "scale": 0.5 }
+        },
+        "cumulativeSequenceLengthT": {
+          "dtype": "int32",
+          "shape": [3],
+          "data": { "kind": "values", "values": [0, 2, 5] }
+        },
+        "pastSeqlensT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [3, 0] } },
+        "blockTableT": { "dtype": "int32", "shape": [2, 3], "data": { "kind": "values", "values": [0, 1, 2, 3, 4, 5] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float16", "shape": [5, 32], "tolerance": 0.00001, "relTolerance": 0.0001 },
+        "keyCacheT": { "dtype": "float16", "shape": [6, 4, 2, 8], "tolerance": 0.00001, "relTolerance": 0.0001 }
+      }
+    },
+    {
+      "name": "explicit_scale",
+      "attrs": { "num_heads": 4, "kv_num_heads": 2, "scale": 0.125, "is_causal": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [5, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [5, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.13, "scale": 0.8 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [5, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.41, "cosStep": 0.07, "scale": 0.7 }
+        },
+        "keyCacheT": {
+          "dtype": "float16",
+          "shape": [6, 4, 2, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 }
+        },
+        "valueCacheT": {
+          "dtype": "float16",
+          "shape": [6, 4, 2, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.19, "scale": 0.5 }
+        },
+        "cumulativeSequenceLengthT": {
+          "dtype": "int32",
+          "shape": [3],
+          "data": { "kind": "values", "values": [0, 2, 5] }
+        },
+        "pastSeqlensT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [3, 0] } },
+        "blockTableT": { "dtype": "int32", "shape": [2, 3], "data": { "kind": "values", "values": [0, 1, 2, 3, 4, 5] } }
+      },
+      "outputs": { "outputT": { "dtype": "float16", "shape": [5, 32], "tolerance": 0.00001, "relTolerance": 0.0001 } }
+    },
+    {
+      "name": "slot_mapping_explicit",
+      "provenance": {
+        "notes": "slot_mapping names each token's flat cache slot directly; these values reproduce the derived mapping, and this case requests only the updated value-cache alias."
+      },
+      "attrs": { "num_heads": 4, "kv_num_heads": 2 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [5, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [5, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.13, "scale": 0.8 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [5, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.41, "cosStep": 0.07, "scale": 0.7 }
+        },
+        "keyCacheT": {
+          "dtype": "float16",
+          "shape": [6, 4, 2, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 }
+        },
+        "valueCacheT": {
+          "dtype": "float16",
+          "shape": [6, 4, 2, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.19, "scale": 0.5 }
+        },
+        "cumulativeSequenceLengthT": {
+          "dtype": "int32",
+          "shape": [3],
+          "data": { "kind": "values", "values": [0, 2, 5] }
+        },
+        "pastSeqlensT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [3, 0] } },
+        "blockTableT": { "dtype": "int32", "shape": [2, 3], "data": { "kind": "values", "values": [0, 1, 2, 3, 4, 5] } },
+        "slotMappingT": { "dtype": "int32", "shape": [5], "data": { "kind": "values", "values": [3, 4, 12, 13, 14] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float16", "shape": [5, 32], "tolerance": 0.00001, "relTolerance": 0.0001 },
+        "valueCacheT": { "dtype": "float16", "shape": [6, 4, 2, 8], "tolerance": 0.00001, "relTolerance": 0.0001 }
+      }
+    },
+    {
+      "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": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [5, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [5, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.13, "scale": 0.8 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [5, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.41, "cosStep": 0.07, "scale": 0.7 }
+        },
+        "keyCacheT": {
+          "dtype": "float16",
+          "shape": [6, 4, 2, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 }
+        },
+        "valueCacheT": {
+          "dtype": "float16",
+          "shape": [6, 4, 2, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.19, "scale": 0.5 }
+        },
+        "cumulativeSequenceLengthT": {
+          "dtype": "int32",
+          "shape": [3],
+          "data": { "kind": "values", "values": [0, 2, 5] }
+        },
+        "pastSeqlensT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [3, 0] } },
+        "blockTableT": { "dtype": "int32", "shape": [2, 3], "data": { "kind": "values", "values": [0, 1, 2, 3, 4, 5] } },
+        "slotMappingT": { "dtype": "int32", "shape": [5], "data": { "kind": "values", "values": [3, -1, 12, 13, 14] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float16", "shape": [5, 32], "tolerance": 0.00001, "relTolerance": 0.0001 },
+        "keyCacheT": { "dtype": "float16", "shape": [6, 4, 2, 8], "tolerance": 0.00001, "relTolerance": 0.0001 },
+        "valueCacheT": { "dtype": "float16", "shape": [6, 4, 2, 8], "tolerance": 0.00001, "relTolerance": 0.0001 }
+      }
+    },
+    {
+      "name": "packed_qkv_slot_mapping",
+      "provenance": {
+        "notes": "Packed [Q|K|V] with an explicit slot mapping, so the scatter reads K and V out of the query row and binds no schedule tensors at all."
+      },
+      "attrs": { "num_heads": 4, "kv_num_heads": 2 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [5, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 }
+        },
+        "keyCacheT": {
+          "dtype": "float16",
+          "shape": [6, 4, 2, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 }
+        },
+        "valueCacheT": {
+          "dtype": "float16",
+          "shape": [6, 4, 2, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.19, "scale": 0.5 }
+        },
+        "cumulativeSequenceLengthT": {
+          "dtype": "int32",
+          "shape": [3],
+          "data": { "kind": "values", "values": [0, 2, 5] }
+        },
+        "pastSeqlensT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [3, 0] } },
+        "blockTableT": { "dtype": "int32", "shape": [2, 3], "data": { "kind": "values", "values": [0, 1, 2, 3, 4, 5] } },
+        "slotMappingT": { "dtype": "int32", "shape": [5], "data": { "kind": "values", "values": [3, 4, 12, 13, 14] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float16", "shape": [5, 32], "tolerance": 0.00001, "relTolerance": 0.0001 },
+        "keyCacheT": { "dtype": "float16", "shape": [6, 4, 2, 8], "tolerance": 0.00001, "relTolerance": 0.0001 },
+        "valueCacheT": { "dtype": "float16", "shape": [6, 4, 2, 8], "tolerance": 0.00001, "relTolerance": 0.0001 }
+      }
+    },
+    {
+      "name": "single_sequence_decode",
+      "provenance": {
+        "notes": "One decode token whose history spans two full blocks and lands at the start of the third."
+      },
+      "attrs": { "num_heads": 4, "kv_num_heads": 2 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [1, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [1, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.13, "scale": 0.8 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [1, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.41, "cosStep": 0.07, "scale": 0.7 }
+        },
+        "keyCacheT": {
+          "dtype": "float16",
+          "shape": [3, 4, 2, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 }
+        },
+        "valueCacheT": {
+          "dtype": "float16",
+          "shape": [3, 4, 2, 8],
+          "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": [7] } },
+        "blockTableT": { "dtype": "int32", "shape": [1, 3], "data": { "kind": "values", "values": [0, 1, 2] } }
+      },
+      "outputs": { "outputT": { "dtype": "float16", "shape": [1, 32], "tolerance": 0.00001, "relTolerance": 0.0001 } }
+    },
+    {
+      "name": "equal_head_counts",
+      "provenance": { "notes": "num_heads equals kv_num_heads, so the grouped-query mapping is the identity." },
+      "attrs": { "num_heads": 4, "kv_num_heads": 4 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [5, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [5, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.13, "scale": 0.8 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [5, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.41, "cosStep": 0.07, "scale": 0.7 }
+        },
+        "keyCacheT": {
+          "dtype": "float16",
+          "shape": [6, 4, 4, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 }
+        },
+        "valueCacheT": {
+          "dtype": "float16",
+          "shape": [6, 4, 4, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.19, "scale": 0.5 }
+        },
+        "cumulativeSequenceLengthT": {
+          "dtype": "int32",
+          "shape": [3],
+          "data": { "kind": "values", "values": [0, 2, 5] }
+        },
+        "pastSeqlensT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [3, 0] } },
+        "blockTableT": { "dtype": "int32", "shape": [2, 3], "data": { "kind": "values", "values": [0, 1, 2, 3, 4, 5] } }
+      },
+      "outputs": { "outputT": { "dtype": "float16", "shape": [5, 32], "tolerance": 0.00001, "relTolerance": 0.0001 } }
+    },
+    {
+      "name": "long_history_multi_tile",
+      "provenance": {
+        "notes": "A 71-key history exceeds one workgroup tile, and the block table is deliberately not the identity, so the online-softmax merge runs over multiple tiles rather than a single one."
+      },
+      "attrs": { "num_heads": 2, "kv_num_heads": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [1, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [1, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.13, "scale": 0.8 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [1, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.41, "cosStep": 0.07, "scale": 0.7 }
+        },
+        "keyCacheT": {
+          "dtype": "float16",
+          "shape": [9, 8, 1, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 }
+        },
+        "valueCacheT": {
+          "dtype": "float16",
+          "shape": [9, 8, 1, 8],
+          "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": [70] } },
+        "blockTableT": {
+          "dtype": "int32",
+          "shape": [1, 9],
+          "data": { "kind": "values", "values": [5, 2, 8, 0, 7, 1, 6, 3, 4] }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float16", "shape": [1, 16], "tolerance": 0.00001, "relTolerance": 0.0001 } }
+    },
+    {
+      "name": "head_size_16",
+      "provenance": {
+        "notes": "A head dimension wider than the tile's per-thread share, so the accumulator loop strides."
+      },
+      "attrs": { "num_heads": 2, "kv_num_heads": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [3, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [3, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.13, "scale": 0.8 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [3, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.41, "cosStep": 0.07, "scale": 0.7 }
+        },
+        "keyCacheT": {
+          "dtype": "float16",
+          "shape": [4, 4, 1, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 }
+        },
+        "valueCacheT": {
+          "dtype": "float16",
+          "shape": [4, 4, 1, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.19, "scale": 0.5 }
+        },
+        "cumulativeSequenceLengthT": {
+          "dtype": "int32",
+          "shape": [3],
+          "data": { "kind": "values", "values": [0, 2, 3] }
+        },
+        "pastSeqlensT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [1, 5] } },
+        "blockTableT": { "dtype": "int32", "shape": [2, 2], "data": { "kind": "values", "values": [0, 1, 2, 3] } }
+      },
+      "outputs": { "outputT": { "dtype": "float16", "shape": [3, 32], "tolerance": 0.00001, "relTolerance": 0.0001 } }
+    },
+    {
+      "name": "f16_two_sequences",
+      "attrs": { "num_heads": 4, "kv_num_heads": 2 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [5, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [5, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.13, "scale": 0.8 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [5, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.41, "cosStep": 0.07, "scale": 0.7 }
+        },
+        "keyCacheT": {
+          "dtype": "float16",
+          "shape": [6, 4, 2, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 }
+        },
+        "valueCacheT": {
+          "dtype": "float16",
+          "shape": [6, 4, 2, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.19, "scale": 0.5 }
+        },
+        "cumulativeSequenceLengthT": {
+          "dtype": "int32",
+          "shape": [3],
+          "data": { "kind": "values", "values": [0, 2, 5] }
+        },
+        "pastSeqlensT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [3, 0] } },
+        "blockTableT": { "dtype": "int32", "shape": [2, 3], "data": { "kind": "values", "values": [0, 1, 2, 3, 4, 5] } }
+      },
+      "outputs": { "outputT": { "dtype": "float16", "shape": [5, 32], "tolerance": 0.004, "relTolerance": 0.01 } }
+    },
+    {
+      "name": "f16_packed_qkv",
+      "attrs": { "num_heads": 4, "kv_num_heads": 2 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [5, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 }
+        },
+        "keyCacheT": {
+          "dtype": "float16",
+          "shape": [6, 4, 2, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 }
+        },
+        "valueCacheT": {
+          "dtype": "float16",
+          "shape": [6, 4, 2, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.19, "scale": 0.5 }
+        },
+        "cumulativeSequenceLengthT": {
+          "dtype": "int32",
+          "shape": [3],
+          "data": { "kind": "values", "values": [0, 2, 5] }
+        },
+        "pastSeqlensT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [3, 0] } },
+        "blockTableT": { "dtype": "int32", "shape": [2, 3], "data": { "kind": "values", "values": [0, 1, 2, 3, 4, 5] } }
+      },
+      "outputs": { "outputT": { "dtype": "float16", "shape": [5, 32], "tolerance": 0.004, "relTolerance": 0.01 } }
+    },
+    {
+      "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": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [2, 4],
+          "data": { "kind": "values", "values": [0.7854, 1.0234, 1.039, 0.8123, 0.3859, -0.1448, -0.6562, -1.028] }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [2, 2],
+          "data": { "kind": "values", "values": [-0.3046, -0.0781, 0.1597, 0.3877] }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [2, 2],
+          "data": { "kind": "values", "values": [0.5647, 0.756, 0.8391, 0.7958] }
+        },
+        "keyCacheT": {
+          "dtype": "float16",
+          "shape": [3, 2, 1, 2],
+          "data": {
+            "kind": "values",
+            "values": [0.2751, 0.3295, 0.3637, 0.3794, 0.3798, 0.369, 0.3514, 0.3311, 0.3118, 0.296, 0.2849, 0.2785]
+          }
+        },
+        "valueCacheT": {
+          "dtype": "float16",
+          "shape": [3, 2, 1, 2],
+          "data": {
+            "kind": "values",
+            "values": [-0.0755, 0.1377, 0.2257, 0.1768, 0.0339, -0.1232, -0.2093, -0.1675, 0.0079, 0.2677, 0.5238, 0.6804]
+          }
+        },
+        "cumulativeSequenceLengthT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [0, 2] } },
+        "pastSeqlensT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [2] } },
+        "blockTableT": { "dtype": "int32", "shape": [1, 3], "data": { "kind": "values", "values": [2, 0, 1] } }
+      },
+      "outputs": {
+        "outputT": {
+          "dtype": "float16",
+          "shape": [2, 4],
+          "data": {
+            "kind": "values",
+            "values": [0.32763671875, 0.5322265625, 0.32470703125, 0.52978515625, 0.4765625, 0.6181640625, 0.4970703125, 0.64599609375]
+          },
+          "tolerance": 0.000001,
+          "relTolerance": 0.00001
+        },
+        "keyCacheT": {
+          "dtype": "float16",
+          "shape": [3, 2, 1, 2],
+          "data": {
+            "kind": "values",
+            "values": [-0.3046, -0.0781, 0.1597, 0.3877, 0.3798, 0.369, 0.3514, 0.3311, 0.3118, 0.296, 0.2849, 0.2785]
+          },
+          "tolerance": 0.000001,
+          "relTolerance": 0.000001
+        },
+        "valueCacheT": {
+          "dtype": "float16",
+          "shape": [3, 2, 1, 2],
+          "data": {
+            "kind": "values",
+            "values": [0.5647, 0.756, 0.8391, 0.7958, 0.0339, -0.1232, -0.2093, -0.1675, 0.0079, 0.2677, 0.5238, 0.6804]
+          },
+          "tolerance": 0.000001,
+          "relTolerance": 0.000001
+        }
+      }
+    },
+    {
+      "name": "splitk_multi_split_history",
+      "provenance": {
+        "notes": "A 301-key history over a 512-key block table, which is what puts the split-K schedule above its two-split floor: four contiguous key ranges of 76 keys each cross a workgroup tile boundary, so every range runs the tile loop more than once and the merge folds four non-empty online states. Grouped-query (2 query heads per KV head) so the split epilogue writes both streams."
+      },
+      "attrs": { "num_heads": 2, "kv_num_heads": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [1, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [1, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.13, "scale": 0.8 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [1, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.41, "cosStep": 0.07, "scale": 0.7 }
+        },
+        "keyCacheT": {
+          "dtype": "float16",
+          "shape": [64, 8, 1, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 }
+        },
+        "valueCacheT": {
+          "dtype": "float16",
+          "shape": [64, 8, 1, 8],
+          "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": [300] } },
+        "blockTableT": {
+          "dtype": "int32",
+          "shape": [1, 64],
+          "data": {
+            "kind": "values",
+            "values": [3, 10, 17, 24, 31, 38, 45, 52, 59, 2, 9, 16, 23, 30, 37, 44, 51, 58, 1, 8, 15, 22, 29, 36, 43, 50, 57, 0, 7, 14, 21, 28, 35, 42, 49, 56, 63, 6, 13, 20, 27, 34, 41, 48, 55, 62, 5, 12, 19, 26, 33, 40, 47, 54, 61, 4, 11, 18, 25, 32, 39, 46, 53, 60]
+          }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float16", "shape": [1, 16], "tolerance": 0.00001, "relTolerance": 0.0001 } }
+    },
+    {
+      "name": "splitk_empty_trailing_split",
+      "provenance": {
+        "notes": "One key of history over a 256-key block table: the split schedule is two ranges of one key each, so the second range covers nothing. It is the split that must write the (-FLT_MAX, 0) identity rather than a normalized row, and the merge must fold that identity away."
+      },
+      "attrs": { "num_heads": 2, "kv_num_heads": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [1, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [1, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.13, "scale": 0.8 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [1, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.41, "cosStep": 0.07, "scale": 0.7 }
+        },
+        "keyCacheT": {
+          "dtype": "float16",
+          "shape": [32, 8, 1, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 }
+        },
+        "valueCacheT": {
+          "dtype": "float16",
+          "shape": [32, 8, 1, 8],
+          "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": [0] } },
+        "blockTableT": {
+          "dtype": "int32",
+          "shape": [1, 32],
+          "data": {
+            "kind": "values",
+            "values": [5, 12, 19, 26, 1, 8, 15, 22, 29, 4, 11, 18, 25, 0, 7, 14, 21, 28, 3, 10, 17, 24, 31, 6, 13, 20, 27, 2, 9, 16, 23, 30]
+          }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float16", "shape": [1, 16], "tolerance": 0.00001, "relTolerance": 0.0001 } }
+    },
+    {
+      "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": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [1, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [1, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.13, "scale": 0.8 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [1, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.41, "cosStep": 0.07, "scale": 0.7 }
+        },
+        "keyCacheT": {
+          "dtype": "float16",
+          "shape": [32, 8, 1, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 }
+        },
+        "valueCacheT": {
+          "dtype": "float16",
+          "shape": [32, 8, 1, 8],
+          "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": [200] } },
+        "blockTableT": {
+          "dtype": "int32",
+          "shape": [1, 32],
+          "data": {
+            "kind": "values",
+            "values": [7, 14, 21, 28, 3, 10, 17, 24, 31, 6, 13, 20, 27, 2, 9, 16, 23, 30, 5, 12, 19, 26, 1, 8, 15, 22, 29, 4, 11, 18, 25, 0]
+          }
+        },
+        "slotMappingT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [176] } }
+      },
+      "outputs": { "outputT": { "dtype": "float16", "shape": [1, 16], "tolerance": 0.00001, "relTolerance": 0.0001 } }
+    },
+    {
+      "name": "splitk_packed_qkv",
+      "provenance": {
+        "notes": "The split-K schedule reading a packed [Q; K; V] query row, so the split pass strides the packed query layout while the merge writes the unpacked output."
+      },
+      "attrs": { "num_heads": 2, "kv_num_heads": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [1, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 }
+        },
+        "keyCacheT": {
+          "dtype": "float16",
+          "shape": [32, 8, 1, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 }
+        },
+        "valueCacheT": {
+          "dtype": "float16",
+          "shape": [32, 8, 1, 8],
+          "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": [200] } },
+        "blockTableT": {
+          "dtype": "int32",
+          "shape": [1, 32],
+          "data": {
+            "kind": "values",
+            "values": [9, 16, 23, 30, 5, 12, 19, 26, 1, 8, 15, 22, 29, 4, 11, 18, 25, 0, 7, 14, 21, 28, 3, 10, 17, 24, 31, 6, 13, 20, 27, 2]
+          }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float16", "shape": [1, 16], "tolerance": 0.00001, "relTolerance": 0.0001 } }
+    },
+    {
+      "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": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [1, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 }
+        },
+        "keyCacheT": {
+          "dtype": "float16",
+          "shape": [32, 8, 1, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 }
+        },
+        "valueCacheT": {
+          "dtype": "float16",
+          "shape": [32, 8, 1, 8],
+          "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": [200] } },
+        "blockTableT": {
+          "dtype": "int32",
+          "shape": [1, 32],
+          "data": {
+            "kind": "values",
+            "values": [11, 18, 25, 0, 7, 14, 21, 28, 3, 10, 17, 24, 31, 6, 13, 20, 27, 2, 9, 16, 23, 30, 5, 12, 19, 26, 1, 8, 15, 22, 29, 4]
+          }
+        },
+        "slotMappingT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [208] } }
+      },
+      "outputs": { "outputT": { "dtype": "float16", "shape": [1, 16], "tolerance": 0.00001, "relTolerance": 0.0001 } }
+    }
+  ]
+}