sync 6fdf6301e2bb

Browse files

Files changed (8) hide show

README.md +109 -0
build/webgpu/attn-flash-decode-splitk-merge.wgsl.jinja +93 -0
build/webgpu/attn-flash-decode-splitk.wgsl.jinja +412 -0
build/webgpu/bench.json +971 -0
build/webgpu/dmmha-present.wgsl.jinja +79 -0
build/webgpu/manifest.json +1535 -0
build/webgpu/metadata.json +48 -0
build/webgpu/test.json +1580 -0

README.md CHANGED Viewed

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

 ---
+library_name: kernels
 license: apache-2.0
+tags:
+- kernel
+- webgpu
+- wgsl
 ---
+# com.microsoft.DecoderMaskedMultiHeadAttention
+`com.microsoft`  ·  ONNX Runtime contrib operator  ·  contrib since_version 1
+## Description
+Single-token decode attention, self and cross, excluding the QKV projection. Cross attention reads whole encoder key and value tensors; self attention appends this token to a capacity-sized cache whose live length arrives in `past_sequence_length`. The mask is additive: a zero entry adds `mask_filter_value` rather than removing the key. The optional `qk` output is the scaled score row before the softmax, and is produced for cross attention only.
+See the [ONNX Runtime `DecoderMaskedMultiHeadAttention` contrib-operator spec](https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.DecoderMaskedMultiHeadAttention) for the reference semantics.
+## Inputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `queryT` | `query` | `T` | `3` | — | Query row for the single decoded token, with shape `(batch_size, 1, num_heads * head_size)`. | required |
+| `keyT` | `key` | `T` | — | — | Cross attention: the full encoder keys with shape `(batch_size, num_heads, kv_sequence_length, head_size)`. Self attention: this token's key row with shape `(batch_size, 1, num_heads * head_size)`. | required |
+| `valueT` | `value` | `T` | — | — | Cross attention: the full encoder values, shaped like `key`. Self attention: this token's value row, shaped like `key`. | required |
+| `maskIndexT` | `mask_index` | `M` | `2` | — | Key-padding mask with shape `(batch_size, total_sequence_length)`. A zero entry ADDS `mask_filter_value` to that key's score; it does not remove the key. | optional |
+| `attentionBiasT` | `attention_bias` | `T` | `4` | — | Additive score bias with shape `(batch_size or 1, num_heads or 1, 1, total_sequence_length)`. Self attention only. | optional |
+| `pastKeyT` | `past_key` | `T` | `4` | — | Key cache at capacity, with shape `(batch_size, num_heads, max_sequence_length, head_size)`. Its presence selects self attention. | optional |
+| `pastValueT` | `past_value` | `T` | `4` | — | Value cache at capacity, shaped like `past_key`. | optional |
+| `pastSequenceLengthT` | `past_sequence_length` | `M` | `1` | — | Live cache length before this token, as a single value. It selects the cache row this token is written to and bounds the key loop, so it is read from tensor data rather than baked. | optional |
+## Outputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `outputT` | `output` | `T` | `3` | `[queryT[0], 1, queryT[2]]` | Attention result with shape `(batch_size, 1, num_heads * head_size)`. | required |
+| `presentKeyT` | `present_key` | `T` | `4` | same as `pastKeyT` | Key cache after appending this token, at capacity. Every row other than the appended one keeps its past value. | optional |
+| `presentValueT` | `present_value` | `T` | `4` | same as `pastValueT` | Value cache after appending this token, shaped like `present_key`. | optional |
+| `qkT` | `qk` | `QK` | `4` | `[queryT[0], num_heads, 1, keyT[2]]` | Scaled `Q * K^T` with shape `(batch_size, num_heads, 1, kv_sequence_length)`, taken after the mask and BEFORE the softmax despite the schema calling it normalized. Cross attention only. | optional |
+## Attributes
+Attributes and default values (overridable per request):
+| Attribute | Default | Description |
+| --- | --- | --- |
+| `mask_filter_value` | `-10000` | Value added to the score of a key whose mask entry is zero. |
+| `num_heads` | — | Number of attention heads. The head size is the query hidden size divided by this count. |
+| `output_qk` | — | Set to 1 to emit the pre-softmax `qk` output. It must agree with whether the caller supplies that output. |
+| `past_present_share_buffer` | — | Set to 1 when the past and present caches are one capacity-sized buffer. Self attention requires it; cross attention has no cache and leaves it at 0. |
+| `scale` | — | Multiplier applied to `Q * K^T`. Omitted or zero selects the default `1 / sqrt(head_size)`. |
+## Type constraints
+| Variable | Allowed dtypes |
+| --- | --- |
+| `T` | `float32`, `float16` |
+| `QK` | `float32`, `float16` |
+| `M` | `int32` |
+## Implementation variants
+One implementation is selected per call from the device capabilities, the request shapes and the dtypes; these notes say what each one covers.
+- `cross_direct` — A single key partition writes normalized attention directly, removing partial buffers and the merge dispatch while preserving masks and diagnostic-score outputs.
+- `cross_qk_direct` — A single key partition writes normalized attention directly, removing partial buffers and the merge dispatch while preserving masks and diagnostic-score outputs.
+- `cross_mask_direct` — A single key partition writes normalized attention directly, removing partial buffers and the merge dispatch while preserving masks and diagnostic-score outputs.
+- `cross_mask_qk_direct` — A single key partition writes normalized attention directly, removing partial buffers and the merge dispatch while preserving masks and diagnostic-score outputs.
+- `cross_nosg_direct` — A single key partition writes normalized attention directly, removing partial buffers and the merge dispatch while preserving masks and diagnostic-score outputs.
+- `cross_qk_nosg_direct` — A single key partition writes normalized attention directly, removing partial buffers and the merge dispatch while preserving masks and diagnostic-score outputs.
+- `cross_mask_nosg_direct` — A single key partition writes normalized attention directly, removing partial buffers and the merge dispatch while preserving masks and diagnostic-score outputs.
+- `cross_mask_qk_nosg_direct` — A single key partition writes normalized attention directly, removing partial buffers and the merge dispatch while preserving masks and diagnostic-score outputs.
+## Device requirements
+Some implementation variants require `subgroups`. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.
+## Files
+- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
+- [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
+- [`test.json`](build/webgpu/test.json) — correctness cases
+- [`bench.json`](build/webgpu/bench.json) — benchmark cases
+- [`attn-flash-decode-splitk-merge.wgsl.jinja`](build/webgpu/attn-flash-decode-splitk-merge.wgsl.jinja)
+- [`attn-flash-decode-splitk.wgsl.jinja`](build/webgpu/attn-flash-decode-splitk.wgsl.jinja)
+- [`dmmha-present.wgsl.jinja`](build/webgpu/dmmha-present.wgsl.jinja)
+## Use with `@huggingface/kernels`
+```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.3
+```
+Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
+The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
+Replace each `*Data` placeholder with a typed array containing the corresponding input data.
+```js
+import { getKernel } from "@huggingface/kernels";
+const kernel = await getKernel("webgpu-kernels/com.microsoft.DecoderMaskedMultiHeadAttention", { version: 1 });
+const { outputT } = await kernel({
+  queryT: { data: queryTData, shape: [1, 1, 32] },
+  keyT: { data: keyTData, shape: [1, 2, 1, 16] },
+  valueT: { data: valueTData, shape: [1, 2, 1, 16] },
+}, {
+  attrs: { num_heads: 2 },
+});
+```

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

	@@ -0,0 +1,93 @@

+{% if splitQueries is not defined %}{% set splitQueries = false %}{% endif %}
+{% set qSeq = 0 %}
+{{ 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: rank-4 [batch, heads, seq, 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 %}
+// 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_HEADS + h) * Q_SEQ + queryToken) * HEAD_DIM_V4;
+{% else %}
+  let qBaseV4 = (b * Q_HEADS + 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);
+  output[qBaseV4 + d4] = vec4<{{ scalar }}>(outValue);
+}

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

	@@ -0,0 +1,412 @@

+{% if useSubgroups is not defined %}{% set useSubgroups = true %}{% endif %}
+{% if splitQueries is not defined %}{% set splitQueries = false %}{% endif %}
+{% if cacheSeqlens is not defined %}{% set cacheSeqlens = false %}{% endif %}
+{% if hasMask is not defined %}{% set hasMask = false %}{% endif %}
+{% if hasKeyPadMask is not defined %}{% set hasKeyPadMask = false %}{% endif %}
+{% if emitQk is not defined %}{% set emitQk = false %}{% endif %}
+{% if cacheSeqlensIndex is not defined %}{% set cacheSeqlensIndex = "b" %}{% endif %}
+{% if maskFilterValue is not defined %}{% set maskFilterValue = 0.0 %}{% endif %}
+{% if qkScalar is not defined %}{% set qkScalar = "f32" %}{% endif %}
+{% set qSeq = 0 %}
+{% set splitKWorkgroupSize = workgroupSizeSpec if workgroupSizeSpec is defined else tunables.WORKGROUP_SIZE %}
+{% if useSubgroups %}
+enable subgroups;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+// Split-K flash attention. Single-partition direct output skips the merge pass.
+// Otherwise this is the first of two passes. This geometry
+// handles decode and short-query, long-context prefill inputs.
+//
+// The non-split flash decode launches only `batch * numHeads` workgroups, each
+// sweeping the whole KV sequence serially in WG-key tiles. This pass splits the
+// KV sequence into `NUM_SPLITS` contiguous ranges and gives each range its own
+// workgroup, so `batch * numHeads * NUM_SPLITS` workgroups run the tiled online
+// softmax in parallel. Each workgroup emits the *un-normalized* online state for
+// its range — the running (max, denom) and the softmax-weighted V sum before the
+// final divide — and the merge pass combines the per-split states with the online
+// rule.
+// Layout: rank-4 [batch, heads, seq, headDim] for Q/K/V.
+const HEAD_DIM: u32 = {{ headDim }}u;
+const HEAD_DIM_V4: u32 = {{ headDimV4 }}u;
+const Q_HEADS: u32 = {{ qNumHeads }}u;
+const KV_HEADS: u32 = {{ kvNumHeads }}u;
+const WG: u32 = {{ splitKWorkgroupSize }}u;
+const NUM_SPLITS: u32 = {{ numSplits }}u;
+{% if splitQueries %}
+const Q_SEQ: u32 = {{ qSeq }}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));
+}
+var<workgroup> q_shared: array<vec4<f32>, HEAD_DIM_V4>;
+var<workgroup> running_out: array<vec4<f32>, HEAD_DIM_V4>;
+var<workgroup> probs: array<f32, WG>;
+{% set coopQk = useSubgroups and headDimV4 >= 8 and not (usesF16 and headDimV4 <= 32) and (allowCooperativeQk if allowCooperativeQk is defined else true) %}
+{% set jGroups = (splitKWorkgroupSize / headDimV4)|int %}
+{% set jSplitV = (splitKWorkgroupSize % headDimV4 == 0) and (jGroups >= 2) %}
+{% if coopQk %}
+var<workgroup> sval_sh: array<f32, WG>;
+{% endif %}
+{% if jSplitV %}
+var<workgroup> vacc_sh: array<vec4<f32>, WG>;
+{% endif %}
+{% set combineSubgroups = useSubgroups %}
+// 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 %}
+// Cross-subgroup merge that assumes nothing about which invocations share a
+// subgroup or how many subgroups there are: each subgroup's elected lane
+// publishes the subgroup pair in the slot at its OWN invocation index and sets
+// that index's bit in a workgroup bitmask; thread 0 then folds exactly the
+// published slots, in ascending index order (the online (m, d) merge is not
+// float-associative, so the order is fixed), and clears the mask for the next
+// call as it reads it. Workgroup memory starts zeroed, so the mask needs no
+// setup. Same three collectives as a single-subgroup reduce, two barriers.
+var<workgroup> partialM: array<f32, WG>;
+var<workgroup> partialD: array<f32, WG>;
+var<workgroup> leaderMask: array<atomic<u32>, (WG + 31u) / 32u>;
+var<workgroup> combinedMD: vec2<f32>;
+// When the whole workgroup is one subgroup the subgroup reduce already covers
+// it (no barriers, no shared state). `subgroup_size` is the size of the current
+// subgroup and uniform, so the test is exact and may guard the barriers below.
+fn combine_partials(m: f32, d: f32, lidx: u32, sgSize: u32) -> vec2<f32> {
+  let sgM = subgroupMax(m);
+  // A lane with no elements contributes d == 0 (exact identity). A +inf
+  // element made exp(inf - inf) = NaN stick in that lane's d; a NaN element
+  // landed in d via exp(NaN); both survive the merge and are detected by the
+  // code after the reduction.
+  let sgD = subgroupAdd(d * exp_shift(m, sgM));
+  if (sgSize == WG) {
+    return vec2<f32>(sgM, sgD);
+  }
+  if (subgroupElect()) {
+    partialM[lidx] = sgM;
+    partialD[lidx] = sgD;
+    atomicOr(&leaderMask[lidx / 32u], 1u << (lidx % 32u));
+  }
+  workgroupBarrier();
+  if (lidx == 0u) {
+    var accM = -FLT_MAX;
+    var accD = 0.0;
+    for (var w = 0u; w < (WG + 31u) / 32u; w = w + 1u) {
+      var bits = atomicExchange(&leaderMask[w], 0u);
+      while (bits != 0u) {
+        let slot = w * 32u + firstTrailingBit(bits);
+        bits = bits & (bits - 1u);
+        let mNew = max(accM, partialM[slot]);
+        accD = accD * exp_shift(accM, mNew) + partialD[slot] * exp_shift(partialM[slot], mNew);
+        accM = mNew;
+      }
+    }
+    combinedMD = vec2<f32>(accM, accD);
+  }
+  workgroupBarrier();
+  return combinedMD;
+}
+{% else %}
+{% set mdExtent = "WG" %}
+var<workgroup> partialM: array<f32, {{ mdExtent }}>;
+var<workgroup> partialD: array<f32, {{ mdExtent }}>;
+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 %}
+{% set ATTN_SCALE_DIM = "HEAD_DIM" %}fn scale_value() -> f32 {
+  if (params.scale != 0.0) { return params.scale; }
+  return inverseSqrt(f32({{ ATTN_SCALE_DIM }}));
+}
+fn load_key4(indexV4: u32) -> vec4<f32> {
+  return vec4<f32>(key[indexV4]);
+}
+fn load_value4(indexV4: u32) -> vec4<f32> {
+  return vec4<f32>(value[indexV4]);
+}
+@compute @workgroup_size(WG, 1, 1)
+fn main(
+  @builtin(workgroup_id) wg: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>{% if useSubgroups %},
+  @builtin(subgroup_size) sgSize: u32{% endif %}
+) {
+{% if useSubgroups %}
+  // Subgroup tiles partition the fixed workgroup exactly. The advertised range
+  // is validated before dispatch; retain this uniform guard for implementations that
+  // choose an intermediate width at pipeline execution time.
+  if (sgSize == 0u || sgSize > WG || WG % sgSize != 0u) { return; }
+{% endif %}
+{% if splitQueries %}
+  let queryToken = wg.x / NUM_SPLITS;
+  let split = wg.x % NUM_SPLITS;
+{% else %}
+  let split = wg.x;
+{% endif %}
+  let h = wg.y;
+  let b = wg.z;
+  if (h >= Q_HEADS || split >= NUM_SPLITS{% if splitQueries %} || queryToken >= Q_SEQ{% endif %}) {
+    return;
+  }
+  let tid = lid.x;
+  let hKv = h / (Q_HEADS / KV_HEADS);
+  let cacheSeq = params.kvSeq;
+{% if cacheSeqlens %}
+  // Buffer-sharing caches retain their capacity in the physical BNSH stride;
+  // seqlens_k supplies the active end independently for each batch.
+  let kvSeq = min(cacheSeq, u32(seqlens_k[{{ cacheSeqlensIndex }}]) + 1u);
+{% else %}
+  let kvSeq = cacheSeq;
+{% endif %}
+  // Query row (decode uses token zero; short-query prefill folds the token into wg.x).
+{% if splitQueries %}
+  let qBaseV4 = ((b * Q_HEADS + h) * Q_SEQ + queryToken) * HEAD_DIM_V4;
+{% else %}
+  let qBaseV4 = (b * Q_HEADS + h) * HEAD_DIM_V4;
+{% endif %}
+  let kvBaseV4 = (b * KV_HEADS + hKv) * cacheSeq * HEAD_DIM_V4;
+  let kvTokenStrideV4 = HEAD_DIM_V4;
+  // Contiguous KV range owned by this split. Ceil division lets the last split
+  // absorb any remainder; empty ranges write identity partials and are ignored
+  // by the merge pass.
+  let keysPerSplit = (kvSeq + NUM_SPLITS - 1u) / NUM_SPLITS;
+  let splitStart = split * keysPerSplit;
+  var splitEnd = splitStart + keysPerSplit;
+  if (splitEnd > kvSeq) {
+    splitEnd = kvSeq;
+  }
+  for (var d4 = tid; d4 < HEAD_DIM_V4; d4 = d4 + WG) {
+    var qv = vec4<f32>(query[qBaseV4 + d4]);
+    q_shared[d4] = qv;
+    running_out[d4] = vec4<f32>(0.0);
+  }
+  workgroupBarrier();
+  let scale = scale_value();
+  var runningMax = -FLT_MAX;
+  var runningDenom = 0.0;
+  var kjBase = splitStart;
+  loop {
+    if (kjBase >= splitEnd) {
+      break;
+    }
+    let kj = kjBase + tid;
+    var keyAllowed = kj < splitEnd;
+    let tileCount = min(WG, splitEnd - kjBase);
+    var score = -FLT_MAX;
+    var m = -FLT_MAX;
+    var dPart = 0.0;
+{% if coopQk %}
+    // Cooperative Q.K: one subgroup per key, lanes splitting HEAD_DIM_V4, then a hardware
+    // subgroupAdd — turns the per-thread HEAD_DIM_V4-long dependent dot chain into a few
+    // strided vec4 dots + one reduce. Uniform trip count keeps subgroupAdd in uniform flow.
+    let sgPerWg = WG / sgSize;
+    let qkRounds = (tileCount + sgPerWg - 1u) / sgPerWg;
+    let lane = tid % sgSize;
+    let sgInWg = tid / sgSize;
+    for (var rr: u32 = 0u; rr < qkRounds; rr = rr + 1u) {
+      let j = rr * sgPerWg + sgInWg;
+      var accS: f32 = 0.0;
+      if (j < tileCount) {
+        let kRowV4 = kvBaseV4 + (kjBase + j) * kvTokenStrideV4;
+        for (var d4: u32 = lane; d4 < HEAD_DIM_V4; d4 = d4 + sgSize) {
+          accS = accS + dot(q_shared[d4], load_key4(kRowV4 + d4));
+        }
+      }
+      let sj = subgroupAdd(accS);
+      if (lane == 0u && j < tileCount) {
+        sval_sh[j] = sj;
+      }
+    }
+    workgroupBarrier();
+    if (keyAllowed) {
+      score = sval_sh[tid] * scale;
+      m = score;
+      dPart = 1.0;
+    }
+{% else %}
+    if (keyAllowed) {
+      let kRowV4 = kvBaseV4 + kj * kvTokenStrideV4;
+      var acc: f32 = 0.0;
+      for (var d4: u32 = 0u; d4 < HEAD_DIM_V4; d4 = d4 + 1u) {
+        acc = acc + dot(q_shared[d4], load_key4(kRowV4 + d4));
+      }
+      score = acc * scale;
+      m = score;
+      dPart = 1.0;
+    }
+{% endif %}
+{% if hasMask %}
+    if (keyAllowed) {
+{% if splitQueries %}
+      let maskQuery = queryToken;
+{% else %}
+      let maskQuery = 0u;
+{% endif %}
+      let maskIndex = b * params.maskBatchStride + h * params.maskHeadStride + maskQuery * params.maskSeqStride + kj;
+      score = score + f32(attn_mask[maskIndex]);
+      m = score;
+    }
+{% endif %}{% if hasKeyPadMask %}
+    if (keyAllowed) {
+      // A zero entry is padding: it ADDS the filter value rather than removing the
+      // key, so the row stays a (heavily skewed) average instead of collapsing.
+      // The mask row width is the live total length, which is exactly kvSeq.
+      score = score + select(f32({{ maskFilterValue }}), 0.0, key_pad_mask[b * kvSeq + kj] > 0);
+      m = score;
+    }
+{% endif %}{% if emitQk %}
+    if (keyAllowed) {
+      // The pre-softmax score row, after scaling and every additive term. Splits own
+      // disjoint key ranges, so the whole row is written exactly once per (b, h).
+      qk[(b * Q_HEADS + h) * kvSeq + kj] = {{ qkScalar }}(score);
+    }
+{% endif %}
+    let tile = combine_partials(m, dPart, tid{% if useSubgroups %}, sgSize{% endif %});
+// Merge one key tile's online-softmax (maximum, denominator) partial into the
+// running state, then store the per-key probabilities consumed by V accumulation.
+    let newMax = max(runningMax, tile.x);
+    let correction = exp_shift(runningMax, newMax);
+    runningDenom = runningDenom * correction + tile.y * exp_shift(tile.x, newMax);
+    runningMax = newMax;
+    var prob = 0.0;
+    if (keyAllowed) {
+      prob = exp_shift(score, newMax);
+    }
+    probs[tid] = prob;
+    workgroupBarrier();
+{% if jSplitV %}
+    // j-split V accumulation: thread (jg, d4v) sums keys j == jg mod
+    // J_GROUPS for dim block d4v into a register, then the groups combine
+    // through shared memory so all lanes participate.
+    const J_GROUPS: u32 = {{ jGroups }}u;
+    let jg = tid / HEAD_DIM_V4;
+    let d4v = tid % HEAD_DIM_V4;
+    var vacc = vec4<f32>(0.0);
+    var jj = jg;
+    loop {
+      if (jj >= tileCount) { break; }
+      vacc = vacc + probs[jj] * load_value4(
+        kvBaseV4 + (kjBase + jj) * kvTokenStrideV4 + d4v      );
+      jj = jj + J_GROUPS;
+    }
+    vacc_sh[tid] = vacc;
+    workgroupBarrier();
+    for (var d4: u32 = tid; d4 < HEAD_DIM_V4; d4 = d4 + WG) {
+      var a4 = running_out[d4] * correction;
+      for (var g: u32 = 0u; g < J_GROUPS; g = g + 1u) {
+        a4 = a4 + vacc_sh[g * HEAD_DIM_V4 + d4];
+      }
+      running_out[d4] = a4;
+    }
+    workgroupBarrier();
+{% else %}
+    for (var d4: u32 = tid; d4 < HEAD_DIM_V4; d4 = d4 + WG) {
+      var vSum = vec4<f32>(0.0);
+      for (var i: u32 = 0u; i < tileCount; i = i + 1u) {
+        vSum = vSum + probs[i] * load_value4(
+          kvBaseV4 + (kjBase + i) * kvTokenStrideV4 + d4        );
+      }
+      running_out[d4] = running_out[d4] * correction + vSum;
+    }
+    workgroupBarrier();
+{% endif %}
+    kjBase = kjBase + WG;
+  }
+{% if directOutput is defined and directOutput %}
+  // The manifest admits this epilogue only for one partition and no post-merge
+  // bias or gate. Preserve the merge's finite/nonfinite weighting and zero-row
+  // behavior before narrowing once to the output dtype.
+  let finalWeight = exp_shift(runningMax, max(-FLT_MAX, runningMax));
+  let finalDenom = runningDenom * finalWeight;
+  for (var d4: u32 = tid; d4 < HEAD_DIM_V4; d4 = d4 + WG) {
+    let finalOut = running_out[d4] * finalWeight;
+    output[qBaseV4 + d4] = vec4<{{ "f16" if usesF16 else "f32" }}>(
+      select(vec4<f32>(0.0), finalOut / finalDenom, finalDenom > 0.0));
+  }
+{% else %}
+  // Emit un-normalized partials for (b, h, split): the merge pass divides.
+{% if splitQueries %}
+  let partialBase = (((b * Q_SEQ + queryToken) * Q_HEADS + h) * NUM_SPLITS + split) * HEAD_DIM_V4;
+{% else %}
+  let partialBase = ((b * Q_HEADS + h) * NUM_SPLITS + split) * HEAD_DIM_V4;
+{% endif %}
+  for (var d4: u32 = tid; d4 < HEAD_DIM_V4; d4 = d4 + WG) {
+    partial_out[partialBase + d4] = running_out[d4];
+  }
+  if (tid == 0u) {
+{% if splitQueries %}
+    let mdBase = ((b * Q_SEQ + queryToken) * Q_HEADS + h) * NUM_SPLITS + split;
+{% else %}
+    let mdBase = (b * Q_HEADS + h) * NUM_SPLITS + split;
+{% endif %}
+    // (max, denom) travel together to the merge, so they share one buffer as an
+    // interleaved vec2 rather than costing two bindings. Interleaved, not two
+    // halves, so the index needs no region size — and the merge reads both
+    // fields of a split in a single load.
+    partial_stats[mdBase] = vec2<f32>(runningMax, runningDenom);
+  }
+{% endif %}
+}

build/webgpu/bench.json ADDED Viewed

	@@ -0,0 +1,971 @@

+{
+  "tunableSpace": { "THREAD_QK_OUTPUT_WORKGROUPS": [0, 2] },
+  "cases": [
+    {
+      "name": "cross-whisper-tiny-b1-h6-kv1500-d64-qk-f32",
+      "preset": "model",
+      "vars": { "elemBytes": 4, "qkBytes": 4 },
+      "attrs": { "num_heads": 6, "output_qk": 1 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 384], "dtype": "float32", "dist": "normal", "seed": 301, "scale": 0.2 },
+        "keyT": { "shape": [1, 6, 1500, 64], "dtype": "float32", "dist": "normal", "seed": 302, "scale": 0.2 },
+        "valueT": { "shape": [1, 6, 1500, 64], "dtype": "float32", "dist": "normal", "seed": 303, "scale": 0.2 }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 1, 384], "dtype": "float32" },
+        "qkT": { "shape": [1, 6, 1, 1500], "dtype": "float32" }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "(1 * 6 * 1500 * 64 * 2 + 1 * 384 * 2) * args.elemBytes + 1 * 6 * 1500 * args.qkBytes"
+          }
+        ]
+      }
+    },
+    {
+      "name": "cross-whisper-large-b1-h20-kv1500-d64-qk-f16",
+      "preset": "model",
+      "vars": { "elemBytes": 2, "qkBytes": 4 },
+      "attrs": { "num_heads": 20, "output_qk": 1 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 1280], "dtype": "float16", "dist": "normal", "seed": 304, "scale": 0.2 },
+        "keyT": { "shape": [1, 20, 1500, 64], "dtype": "float16", "dist": "normal", "seed": 305, "scale": 0.2 },
+        "valueT": { "shape": [1, 20, 1500, 64], "dtype": "float16", "dist": "normal", "seed": 306, "scale": 0.2 }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 1, 1280], "dtype": "float16" },
+        "qkT": { "shape": [1, 20, 1, 1500], "dtype": "float32" }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "(1 * 20 * 1500 * 64 * 2 + 1 * 1280 * 2) * args.elemBytes + 1 * 20 * 1500 * args.qkBytes"
+          }
+        ]
+      }
+    },
+    {
+      "name": "cross-whisper-large-b4-h20-kv1500-d64-qk-f16",
+      "preset": "model",
+      "vars": { "elemBytes": 2, "qkBytes": 4 },
+      "attrs": { "num_heads": 20, "output_qk": 1 },
+      "inputs": {
+        "queryT": { "shape": [4, 1, 1280], "dtype": "float16", "dist": "normal", "seed": 307, "scale": 0.2 },
+        "keyT": { "shape": [4, 20, 1500, 64], "dtype": "float16", "dist": "normal", "seed": 308, "scale": 0.2 },
+        "valueT": { "shape": [4, 20, 1500, 64], "dtype": "float16", "dist": "normal", "seed": 309, "scale": 0.2 }
+      },
+      "outputs": {
+        "outputT": { "shape": [4, 1, 1280], "dtype": "float16" },
+        "qkT": { "shape": [4, 20, 1, 1500], "dtype": "float32" }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "(4 * 20 * 1500 * 64 * 2 + 4 * 1280 * 2) * args.elemBytes + 4 * 20 * 1500 * args.qkBytes"
+          }
+        ]
+      }
+    },
+    {
+      "name": "cross-whisper-large-b8-h20-kv1500-d64-f16",
+      "preset": "model",
+      "vars": { "elemBytes": 2, "qkBytes": 0 },
+      "attrs": { "num_heads": 20 },
+      "inputs": {
+        "queryT": { "shape": [8, 1, 1280], "dtype": "float16", "dist": "normal", "seed": 310, "scale": 0.2 },
+        "keyT": { "shape": [8, 20, 1500, 64], "dtype": "float16", "dist": "normal", "seed": 311, "scale": 0.2 },
+        "valueT": { "shape": [8, 20, 1500, 64], "dtype": "float16", "dist": "normal", "seed": 312, "scale": 0.2 }
+      },
+      "outputs": { "outputT": { "shape": [8, 1, 1280], "dtype": "float16" } },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "(8 * 20 * 1500 * 64 * 2 + 8 * 1280 * 2) * args.elemBytes" }]
+      }
+    },
+    {
+      "name": "cross-whisper-large-b8-h20-kv1500-d64-f32",
+      "preset": "model",
+      "vars": { "elemBytes": 4, "qkBytes": 0 },
+      "attrs": { "num_heads": 20 },
+      "inputs": {
+        "queryT": { "shape": [8, 1, 1280], "dtype": "float32", "dist": "normal", "seed": 313, "scale": 0.2 },
+        "keyT": { "shape": [8, 20, 1500, 64], "dtype": "float32", "dist": "normal", "seed": 314, "scale": 0.2 },
+        "valueT": { "shape": [8, 20, 1500, 64], "dtype": "float32", "dist": "normal", "seed": 315, "scale": 0.2 }
+      },
+      "outputs": { "outputT": { "shape": [8, 1, 1280], "dtype": "float32" } },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "(8 * 20 * 1500 * 64 * 2 + 8 * 1280 * 2) * args.elemBytes" }]
+      }
+    },
+    {
+      "name": "cross-whisper-large-b16-h20-kv1500-d64-f16",
+      "preset": "stress",
+      "vars": { "elemBytes": 2, "qkBytes": 0 },
+      "attrs": { "num_heads": 20 },
+      "inputs": {
+        "queryT": { "shape": [16, 1, 1280], "dtype": "float16", "dist": "normal", "seed": 316, "scale": 0.2 },
+        "keyT": { "shape": [16, 20, 1500, 64], "dtype": "float16", "dist": "normal", "seed": 317, "scale": 0.2 },
+        "valueT": { "shape": [16, 20, 1500, 64], "dtype": "float16", "dist": "normal", "seed": 318, "scale": 0.2 }
+      },
+      "outputs": { "outputT": { "shape": [16, 1, 1280], "dtype": "float16" } },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "(16 * 20 * 1500 * 64 * 2 + 16 * 1280 * 2) * args.elemBytes" }]
+      }
+    },
+    {
+      "name": "cross-whisper-large-b8-h20-kv1500-d64-mask-f16",
+      "preset": "model",
+      "vars": { "elemBytes": 2, "qkBytes": 0 },
+      "attrs": { "num_heads": 20 },
+      "inputs": {
+        "queryT": { "shape": [8, 1, 1280], "dtype": "float16", "dist": "normal", "seed": 319, "scale": 0.2 },
+        "keyT": { "shape": [8, 20, 1500, 64], "dtype": "float16", "dist": "normal", "seed": 320, "scale": 0.2 },
+        "valueT": { "shape": [8, 20, 1500, 64], "dtype": "float16", "dist": "normal", "seed": 321, "scale": 0.2 },
+        "maskIndexT": { "shape": [8, 1500], "dtype": "int32", "data": { "kind": "cycle", "values": [1, 1, 1, 0] } }
+      },
+      "outputs": { "outputT": { "shape": [8, 1, 1280], "dtype": "float16" } },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "(8 * 20 * 1500 * 64 * 2 + 8 * 1280 * 2) * args.elemBytes" }]
+      }
+    },
+    {
+      "name": "cross-whisper-large-b8-h20-kv1500-d64-qk-f16",
+      "preset": "model",
+      "vars": { "elemBytes": 2, "qkBytes": 4 },
+      "attrs": { "num_heads": 20, "output_qk": 1 },
+      "inputs": {
+        "queryT": { "shape": [8, 1, 1280], "dtype": "float16", "dist": "normal", "seed": 322, "scale": 0.2 },
+        "keyT": { "shape": [8, 20, 1500, 64], "dtype": "float16", "dist": "normal", "seed": 323, "scale": 0.2 },
+        "valueT": { "shape": [8, 20, 1500, 64], "dtype": "float16", "dist": "normal", "seed": 324, "scale": 0.2 }
+      },
+      "outputs": {
+        "outputT": { "shape": [8, 1, 1280], "dtype": "float16" },
+        "qkT": { "shape": [8, 20, 1, 1500], "dtype": "float32" }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "(8 * 20 * 1500 * 64 * 2 + 8 * 1280 * 2) * args.elemBytes + 8 * 20 * 1500 * args.qkBytes"
+          }
+        ]
+      }
+    },
+    {
+      "name": "cross-b1-h6-kv256-d64-f32",
+      "preset": "smoke",
+      "vars": { "elemBytes": 4, "qkBytes": 0 },
+      "attrs": { "num_heads": 6 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 384], "dtype": "float32", "dist": "normal", "seed": 325, "scale": 0.2 },
+        "keyT": { "shape": [1, 6, 256, 64], "dtype": "float32", "dist": "normal", "seed": 326, "scale": 0.2 },
+        "valueT": { "shape": [1, 6, 256, 64], "dtype": "float32", "dist": "normal", "seed": 327, "scale": 0.2 }
+      },
+      "outputs": { "outputT": { "shape": [1, 1, 384], "dtype": "float32" } },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "(1 * 6 * 256 * 64 * 2 + 1 * 384 * 2) * args.elemBytes" }]
+      }
+    },
+    {
+      "name": "cross-b2-h8-kv512-d128-f32",
+      "preset": "smoke",
+      "vars": { "elemBytes": 4, "qkBytes": 0 },
+      "attrs": { "num_heads": 8 },
+      "inputs": {
+        "queryT": { "shape": [2, 1, 1024], "dtype": "float32", "dist": "normal", "seed": 328, "scale": 0.2 },
+        "keyT": { "shape": [2, 8, 512, 128], "dtype": "float32", "dist": "normal", "seed": 329, "scale": 0.2 },
+        "valueT": { "shape": [2, 8, 512, 128], "dtype": "float32", "dist": "normal", "seed": 330, "scale": 0.2 }
+      },
+      "outputs": { "outputT": { "shape": [2, 1, 1024], "dtype": "float32" } },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "(2 * 8 * 512 * 128 * 2 + 2 * 1024 * 2) * args.elemBytes" }]
+      }
+    },
+    {
+      "name": "self-whisper-base-b1-h8-cap448-d64-past1-f32",
+      "preset": "model",
+      "vars": { "elemBytes": 4 },
+      "attrs": { "num_heads": 8, "past_present_share_buffer": 1 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 512], "dtype": "float32", "dist": "normal", "seed": 331, "scale": 0.2 },
+        "keyT": { "shape": [1, 1, 512], "dtype": "float32", "dist": "normal", "seed": 332, "scale": 0.2 },
+        "valueT": { "shape": [1, 1, 512], "dtype": "float32", "dist": "normal", "seed": 333, "scale": 0.2 },
+        "pastKeyT": { "shape": [1, 8, 448, 64], "dtype": "float32", "dist": "normal", "seed": 334, "scale": 0.2 },
+        "pastValueT": { "shape": [1, 8, 448, 64], "dtype": "float32", "dist": "normal", "seed": 335, "scale": 0.2 },
+        "pastSequenceLengthT": { "shape": [1], "dtype": "int32", "data": { "kind": "values", "values": [1] } }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 1, 512], "dtype": "float32" },
+        "presentKeyT": { "shape": [1, 8, 448, 64], "dtype": "float32" },
+        "presentValueT": { "shape": [1, 8, 448, 64], "dtype": "float32" }
+      },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "(4 * 229376 + 2 * 1024 + 2 * 1 * 512) * args.elemBytes" }]
+      }
+    },
+    {
+      "name": "self-whisper-base-b1-h8-cap448-d64-past223-f32",
+      "preset": "model",
+      "vars": { "elemBytes": 4 },
+      "attrs": { "num_heads": 8, "past_present_share_buffer": 1 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 512], "dtype": "float32", "dist": "normal", "seed": 336, "scale": 0.2 },
+        "keyT": { "shape": [1, 1, 512], "dtype": "float32", "dist": "normal", "seed": 337, "scale": 0.2 },
+        "valueT": { "shape": [1, 1, 512], "dtype": "float32", "dist": "normal", "seed": 338, "scale": 0.2 },
+        "pastKeyT": { "shape": [1, 8, 448, 64], "dtype": "float32", "dist": "normal", "seed": 339, "scale": 0.2 },
+        "pastValueT": { "shape": [1, 8, 448, 64], "dtype": "float32", "dist": "normal", "seed": 340, "scale": 0.2 },
+        "pastSequenceLengthT": { "shape": [1], "dtype": "int32", "data": { "kind": "values", "values": [223] } }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 1, 512], "dtype": "float32" },
+        "presentKeyT": { "shape": [1, 8, 448, 64], "dtype": "float32" },
+        "presentValueT": { "shape": [1, 8, 448, 64], "dtype": "float32" }
+      },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "(4 * 229376 + 2 * 114688 + 2 * 1 * 512) * args.elemBytes" }]
+      }
+    },
+    {
+      "name": "self-whisper-base-b1-h8-cap448-d64-past447-f32",
+      "preset": "model",
+      "vars": { "elemBytes": 4 },
+      "attrs": { "num_heads": 8, "past_present_share_buffer": 1 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 512], "dtype": "float32", "dist": "normal", "seed": 341, "scale": 0.2 },
+        "keyT": { "shape": [1, 1, 512], "dtype": "float32", "dist": "normal", "seed": 342, "scale": 0.2 },
+        "valueT": { "shape": [1, 1, 512], "dtype": "float32", "dist": "normal", "seed": 343, "scale": 0.2 },
+        "pastKeyT": { "shape": [1, 8, 448, 64], "dtype": "float32", "dist": "normal", "seed": 344, "scale": 0.2 },
+        "pastValueT": { "shape": [1, 8, 448, 64], "dtype": "float32", "dist": "normal", "seed": 345, "scale": 0.2 },
+        "pastSequenceLengthT": { "shape": [1], "dtype": "int32", "data": { "kind": "values", "values": [446] } }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 1, 512], "dtype": "float32" },
+        "presentKeyT": { "shape": [1, 8, 448, 64], "dtype": "float32" },
+        "presentValueT": { "shape": [1, 8, 448, 64], "dtype": "float32" }
+      },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "(4 * 229376 + 2 * 228864 + 2 * 1 * 512) * args.elemBytes" }]
+      }
+    },
+    {
+      "name": "self-whisper-base-b4-h8-cap448-d64-past223-f16",
+      "preset": "model",
+      "vars": { "elemBytes": 2 },
+      "attrs": { "num_heads": 8, "past_present_share_buffer": 1 },
+      "inputs": {
+        "queryT": { "shape": [4, 1, 512], "dtype": "float16", "dist": "normal", "seed": 346, "scale": 0.2 },
+        "keyT": { "shape": [4, 1, 512], "dtype": "float16", "dist": "normal", "seed": 347, "scale": 0.2 },
+        "valueT": { "shape": [4, 1, 512], "dtype": "float16", "dist": "normal", "seed": 348, "scale": 0.2 },
+        "pastKeyT": { "shape": [4, 8, 448, 64], "dtype": "float16", "dist": "normal", "seed": 349, "scale": 0.2 },
+        "pastValueT": { "shape": [4, 8, 448, 64], "dtype": "float16", "dist": "normal", "seed": 350, "scale": 0.2 },
+        "pastSequenceLengthT": { "shape": [1], "dtype": "int32", "data": { "kind": "values", "values": [223] } }
+      },
+      "outputs": {
+        "outputT": { "shape": [4, 1, 512], "dtype": "float16" },
+        "presentKeyT": { "shape": [4, 8, 448, 64], "dtype": "float16" },
+        "presentValueT": { "shape": [4, 8, 448, 64], "dtype": "float16" }
+      },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "(4 * 917504 + 2 * 458752 + 2 * 4 * 512) * args.elemBytes" }]
+      }
+    },
+    {
+      "name": "self-b32-h8-cap448-d64-past223-f32",
+      "preset": "model",
+      "vars": { "elemBytes": 4 },
+      "attrs": { "num_heads": 8, "past_present_share_buffer": 1 },
+      "inputs": {
+        "queryT": { "shape": [32, 1, 512], "dtype": "float32", "dist": "normal", "seed": 351, "scale": 0.2 },
+        "keyT": { "shape": [32, 1, 512], "dtype": "float32", "dist": "normal", "seed": 352, "scale": 0.2 },
+        "valueT": { "shape": [32, 1, 512], "dtype": "float32", "dist": "normal", "seed": 353, "scale": 0.2 },
+        "pastKeyT": { "shape": [32, 8, 448, 64], "dtype": "float32", "dist": "normal", "seed": 354, "scale": 0.2 },
+        "pastValueT": { "shape": [32, 8, 448, 64], "dtype": "float32", "dist": "normal", "seed": 355, "scale": 0.2 },
+        "pastSequenceLengthT": { "shape": [1], "dtype": "int32", "data": { "kind": "values", "values": [223] } }
+      },
+      "outputs": {
+        "outputT": { "shape": [32, 1, 512], "dtype": "float32" },
+        "presentKeyT": { "shape": [32, 8, 448, 64], "dtype": "float32" },
+        "presentValueT": { "shape": [32, 8, 448, 64], "dtype": "float32" }
+      },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "(4 * 7340032 + 2 * 3670016 + 2 * 32 * 512) * args.elemBytes" }]
+      }
+    },
+    {
+      "name": "self-b32-h8-cap448-d64-past223-f16",
+      "preset": "model",
+      "vars": { "elemBytes": 2 },
+      "attrs": { "num_heads": 8, "past_present_share_buffer": 1 },
+      "inputs": {
+        "queryT": { "shape": [32, 1, 512], "dtype": "float16", "dist": "normal", "seed": 356, "scale": 0.2 },
+        "keyT": { "shape": [32, 1, 512], "dtype": "float16", "dist": "normal", "seed": 357, "scale": 0.2 },
+        "valueT": { "shape": [32, 1, 512], "dtype": "float16", "dist": "normal", "seed": 358, "scale": 0.2 },
+        "pastKeyT": { "shape": [32, 8, 448, 64], "dtype": "float16", "dist": "normal", "seed": 359, "scale": 0.2 },
+        "pastValueT": { "shape": [32, 8, 448, 64], "dtype": "float16", "dist": "normal", "seed": 360, "scale": 0.2 },
+        "pastSequenceLengthT": { "shape": [1], "dtype": "int32", "data": { "kind": "values", "values": [223] } }
+      },
+      "outputs": {
+        "outputT": { "shape": [32, 1, 512], "dtype": "float16" },
+        "presentKeyT": { "shape": [32, 8, 448, 64], "dtype": "float16" },
+        "presentValueT": { "shape": [32, 8, 448, 64], "dtype": "float16" }
+      },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "(4 * 7340032 + 2 * 3670016 + 2 * 32 * 512) * args.elemBytes" }]
+      }
+    },
+    {
+      "name": "self-b32-h8-cap448-d64-past447-f32",
+      "preset": "stress",
+      "vars": { "elemBytes": 4 },
+      "attrs": { "num_heads": 8, "past_present_share_buffer": 1 },
+      "inputs": {
+        "queryT": { "shape": [32, 1, 512], "dtype": "float32", "dist": "normal", "seed": 361, "scale": 0.2 },
+        "keyT": { "shape": [32, 1, 512], "dtype": "float32", "dist": "normal", "seed": 362, "scale": 0.2 },
+        "valueT": { "shape": [32, 1, 512], "dtype": "float32", "dist": "normal", "seed": 363, "scale": 0.2 },
+        "pastKeyT": { "shape": [32, 8, 448, 64], "dtype": "float32", "dist": "normal", "seed": 364, "scale": 0.2 },
+        "pastValueT": { "shape": [32, 8, 448, 64], "dtype": "float32", "dist": "normal", "seed": 365, "scale": 0.2 },
+        "pastSequenceLengthT": { "shape": [1], "dtype": "int32", "data": { "kind": "values", "values": [446] } }
+      },
+      "outputs": {
+        "outputT": { "shape": [32, 1, 512], "dtype": "float32" },
+        "presentKeyT": { "shape": [32, 8, 448, 64], "dtype": "float32" },
+        "presentValueT": { "shape": [32, 8, 448, 64], "dtype": "float32" }
+      },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "(4 * 7340032 + 2 * 7323648 + 2 * 32 * 512) * args.elemBytes" }]
+      }
+    },
+    {
+      "name": "self-b32-h8-cap448-d64-past223-mask-f32",
+      "preset": "model",
+      "vars": { "elemBytes": 4 },
+      "attrs": { "num_heads": 8, "past_present_share_buffer": 1 },
+      "inputs": {
+        "queryT": { "shape": [32, 1, 512], "dtype": "float32", "dist": "normal", "seed": 366, "scale": 0.2 },
+        "keyT": { "shape": [32, 1, 512], "dtype": "float32", "dist": "normal", "seed": 367, "scale": 0.2 },
+        "valueT": { "shape": [32, 1, 512], "dtype": "float32", "dist": "normal", "seed": 368, "scale": 0.2 },
+        "maskIndexT": { "shape": [32, 224], "dtype": "int32", "data": { "kind": "cycle", "values": [1, 1, 1, 0] } },
+        "pastKeyT": { "shape": [32, 8, 448, 64], "dtype": "float32", "dist": "normal", "seed": 369, "scale": 0.2 },
+        "pastValueT": { "shape": [32, 8, 448, 64], "dtype": "float32", "dist": "normal", "seed": 370, "scale": 0.2 },
+        "pastSequenceLengthT": { "shape": [1], "dtype": "int32", "data": { "kind": "values", "values": [223] } }
+      },
+      "outputs": {
+        "outputT": { "shape": [32, 1, 512], "dtype": "float32" },
+        "presentKeyT": { "shape": [32, 8, 448, 64], "dtype": "float32" },
+        "presentValueT": { "shape": [32, 8, 448, 64], "dtype": "float32" }
+      },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "(4 * 7340032 + 2 * 3670016 + 2 * 32 * 512) * args.elemBytes" }]
+      }
+    },
+    {
+      "name": "self-b16-h12-cap2048-d64-past1023-f16",
+      "preset": "stress",
+      "vars": { "elemBytes": 2 },
+      "attrs": { "num_heads": 12, "past_present_share_buffer": 1 },
+      "inputs": {
+        "queryT": { "shape": [16, 1, 768], "dtype": "float16", "dist": "normal", "seed": 371, "scale": 0.2 },
+        "keyT": { "shape": [16, 1, 768], "dtype": "float16", "dist": "normal", "seed": 372, "scale": 0.2 },
+        "valueT": { "shape": [16, 1, 768], "dtype": "float16", "dist": "normal", "seed": 373, "scale": 0.2 },
+        "pastKeyT": { "shape": [16, 12, 2048, 64], "dtype": "float16", "dist": "normal", "seed": 374, "scale": 0.2 },
+        "pastValueT": { "shape": [16, 12, 2048, 64], "dtype": "float16", "dist": "normal", "seed": 375, "scale": 0.2 },
+        "pastSequenceLengthT": { "shape": [1], "dtype": "int32", "data": { "kind": "values", "values": [1023] } }
+      },
+      "outputs": {
+        "outputT": { "shape": [16, 1, 768], "dtype": "float16" },
+        "presentKeyT": { "shape": [16, 12, 2048, 64], "dtype": "float16" },
+        "presentValueT": { "shape": [16, 12, 2048, 64], "dtype": "float16" }
+      },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "(4 * 25165824 + 2 * 12582912 + 2 * 16 * 768) * args.elemBytes" }]
+      }
+    },
+    {
+      "name": "self-b1-h4-cap64-d64-past31-f32",
+      "preset": "smoke",
+      "vars": { "elemBytes": 4 },
+      "attrs": { "num_heads": 4, "past_present_share_buffer": 1 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 256], "dtype": "float32", "dist": "normal", "seed": 376, "scale": 0.2 },
+        "keyT": { "shape": [1, 1, 256], "dtype": "float32", "dist": "normal", "seed": 377, "scale": 0.2 },
+        "valueT": { "shape": [1, 1, 256], "dtype": "float32", "dist": "normal", "seed": 378, "scale": 0.2 },
+        "pastKeyT": { "shape": [1, 4, 64, 64], "dtype": "float32", "dist": "normal", "seed": 379, "scale": 0.2 },
+        "pastValueT": { "shape": [1, 4, 64, 64], "dtype": "float32", "dist": "normal", "seed": 380, "scale": 0.2 },
+        "pastSequenceLengthT": { "shape": [1], "dtype": "int32", "data": { "kind": "values", "values": [31] } }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 1, 256], "dtype": "float32" },
+        "presentKeyT": { "shape": [1, 4, 64, 64], "dtype": "float32" },
+        "presentValueT": { "shape": [1, 4, 64, 64], "dtype": "float32" }
+      },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "(4 * 16384 + 2 * 8192 + 2 * 1 * 256) * args.elemBytes" }]
+      }
+    },
+    {
+      "name": "cross-b1-h6-kv16-d64-f32",
+      "preset": "smoke",
+      "vars": { "elemBytes": 4, "qkBytes": 0, "batch": 1, "heads": 6, "keys": 16, "headDim": 64, "maskBytes": 0 },
+      "attrs": { "num_heads": 6 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 384], "dtype": "float32", "dist": "normal", "seed": 325, "scale": 0.2 },
+        "keyT": { "shape": [1, 6, 16, 64], "dtype": "float32", "dist": "normal", "seed": 326, "scale": 0.2 },
+        "valueT": { "shape": [1, 6, 16, 64], "dtype": "float32", "dist": "normal", "seed": 327, "scale": 0.2 }
+      },
+      "outputs": { "outputT": { "shape": [1, 1, 384], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * args.batch * args.heads * args.headDim * (args.keys + 1) * args.elemBytes + args.maskBytes + args.qkBytes",
+            "provenance": "logical"
+          }
+        ]
+      }
+    },
+    {
+      "name": "cross-b1-h6-kv16-d64-f16",
+      "preset": "smoke",
+      "vars": { "elemBytes": 2, "qkBytes": 0, "batch": 1, "heads": 6, "keys": 16, "headDim": 64, "maskBytes": 0 },
+      "attrs": { "num_heads": 6 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 384], "dtype": "float16", "dist": "normal", "seed": 325, "scale": 0.2 },
+        "keyT": { "shape": [1, 6, 16, 64], "dtype": "float16", "dist": "normal", "seed": 326, "scale": 0.2 },
+        "valueT": { "shape": [1, 6, 16, 64], "dtype": "float16", "dist": "normal", "seed": 327, "scale": 0.2 }
+      },
+      "outputs": { "outputT": { "shape": [1, 1, 384], "dtype": "float16" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * args.batch * args.heads * args.headDim * (args.keys + 1) * args.elemBytes + args.maskBytes + args.qkBytes",
+            "provenance": "logical"
+          }
+        ]
+      }
+    },
+    {
+      "name": "cross-b1-h6-kv64-d64-f32",
+      "preset": "smoke",
+      "vars": { "elemBytes": 4, "qkBytes": 0, "batch": 1, "heads": 6, "keys": 64, "headDim": 64, "maskBytes": 0 },
+      "attrs": { "num_heads": 6 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 384], "dtype": "float32", "dist": "normal", "seed": 325, "scale": 0.2 },
+        "keyT": { "shape": [1, 6, 64, 64], "dtype": "float32", "dist": "normal", "seed": 326, "scale": 0.2 },
+        "valueT": { "shape": [1, 6, 64, 64], "dtype": "float32", "dist": "normal", "seed": 327, "scale": 0.2 }
+      },
+      "outputs": { "outputT": { "shape": [1, 1, 384], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * args.batch * args.heads * args.headDim * (args.keys + 1) * args.elemBytes + args.maskBytes + args.qkBytes",
+            "provenance": "logical"
+          }
+        ]
+      }
+    },
+    {
+      "name": "cross-b1-h6-kv64-d64-f16",
+      "preset": "smoke",
+      "vars": { "elemBytes": 2, "qkBytes": 0, "batch": 1, "heads": 6, "keys": 64, "headDim": 64, "maskBytes": 0 },
+      "attrs": { "num_heads": 6 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 384], "dtype": "float16", "dist": "normal", "seed": 325, "scale": 0.2 },
+        "keyT": { "shape": [1, 6, 64, 64], "dtype": "float16", "dist": "normal", "seed": 326, "scale": 0.2 },
+        "valueT": { "shape": [1, 6, 64, 64], "dtype": "float16", "dist": "normal", "seed": 327, "scale": 0.2 }
+      },
+      "outputs": { "outputT": { "shape": [1, 1, 384], "dtype": "float16" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * args.batch * args.heads * args.headDim * (args.keys + 1) * args.elemBytes + args.maskBytes + args.qkBytes",
+            "provenance": "logical"
+          }
+        ]
+      }
+    },
+    {
+      "name": "cross-b1-h6-kv255-d64-f32",
+      "preset": "smoke",
+      "vars": { "elemBytes": 4, "qkBytes": 0, "batch": 1, "heads": 6, "keys": 255, "headDim": 64, "maskBytes": 0 },
+      "attrs": { "num_heads": 6 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 384], "dtype": "float32", "dist": "normal", "seed": 325, "scale": 0.2 },
+        "keyT": { "shape": [1, 6, 255, 64], "dtype": "float32", "dist": "normal", "seed": 326, "scale": 0.2 },
+        "valueT": { "shape": [1, 6, 255, 64], "dtype": "float32", "dist": "normal", "seed": 327, "scale": 0.2 }
+      },
+      "outputs": { "outputT": { "shape": [1, 1, 384], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * args.batch * args.heads * args.headDim * (args.keys + 1) * args.elemBytes + args.maskBytes + args.qkBytes",
+            "provenance": "logical"
+          }
+        ]
+      }
+    },
+    {
+      "name": "cross-b1-h6-kv255-d64-f16",
+      "preset": "smoke",
+      "vars": { "elemBytes": 2, "qkBytes": 0, "batch": 1, "heads": 6, "keys": 255, "headDim": 64, "maskBytes": 0 },
+      "attrs": { "num_heads": 6 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 384], "dtype": "float16", "dist": "normal", "seed": 325, "scale": 0.2 },
+        "keyT": { "shape": [1, 6, 255, 64], "dtype": "float16", "dist": "normal", "seed": 326, "scale": 0.2 },
+        "valueT": { "shape": [1, 6, 255, 64], "dtype": "float16", "dist": "normal", "seed": 327, "scale": 0.2 }
+      },
+      "outputs": { "outputT": { "shape": [1, 1, 384], "dtype": "float16" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * args.batch * args.heads * args.headDim * (args.keys + 1) * args.elemBytes + args.maskBytes + args.qkBytes",
+            "provenance": "logical"
+          }
+        ]
+      }
+    },
+    {
+      "name": "cross-b1-h6-kv257-d64-f32",
+      "preset": "smoke",
+      "vars": { "elemBytes": 4, "qkBytes": 0, "batch": 1, "heads": 6, "keys": 257, "headDim": 64, "maskBytes": 0 },
+      "attrs": { "num_heads": 6 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 384], "dtype": "float32", "dist": "normal", "seed": 325, "scale": 0.2 },
+        "keyT": { "shape": [1, 6, 257, 64], "dtype": "float32", "dist": "normal", "seed": 326, "scale": 0.2 },
+        "valueT": { "shape": [1, 6, 257, 64], "dtype": "float32", "dist": "normal", "seed": 327, "scale": 0.2 }
+      },
+      "outputs": { "outputT": { "shape": [1, 1, 384], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * args.batch * args.heads * args.headDim * (args.keys + 1) * args.elemBytes + args.maskBytes + args.qkBytes",
+            "provenance": "logical"
+          }
+        ]
+      }
+    },
+    {
+      "name": "cross-b1-h6-kv257-d64-f16",
+      "preset": "smoke",
+      "vars": { "elemBytes": 2, "qkBytes": 0, "batch": 1, "heads": 6, "keys": 257, "headDim": 64, "maskBytes": 0 },
+      "attrs": { "num_heads": 6 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 384], "dtype": "float16", "dist": "normal", "seed": 325, "scale": 0.2 },
+        "keyT": { "shape": [1, 6, 257, 64], "dtype": "float16", "dist": "normal", "seed": 326, "scale": 0.2 },
+        "valueT": { "shape": [1, 6, 257, 64], "dtype": "float16", "dist": "normal", "seed": 327, "scale": 0.2 }
+      },
+      "outputs": { "outputT": { "shape": [1, 1, 384], "dtype": "float16" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * args.batch * args.heads * args.headDim * (args.keys + 1) * args.elemBytes + args.maskBytes + args.qkBytes",
+            "provenance": "logical"
+          }
+        ]
+      }
+    },
+    {
+      "name": "cross-mask-qk-b1-h6-kv64-d64-f32",
+      "preset": "smoke",
+      "vars": { "elemBytes": 4, "qkBytes": 1536, "batch": 1, "heads": 6, "keys": 64, "headDim": 64, "maskBytes": 256 },
+      "attrs": { "num_heads": 6, "output_qk": 1, "mask_filter_value": -4 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 384], "dtype": "float32", "dist": "normal", "seed": 325, "scale": 0.2 },
+        "keyT": { "shape": [1, 6, 64, 64], "dtype": "float32", "dist": "normal", "seed": 326, "scale": 0.2 },
+        "valueT": { "shape": [1, 6, 64, 64], "dtype": "float32", "dist": "normal", "seed": 327, "scale": 0.2 },
+        "maskIndexT": { "dtype": "int32", "shape": [1, 64], "data": { "kind": "cycle", "values": [1, 0, 1] } }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 1, 384], "dtype": "float32" },
+        "qkT": { "dtype": "float32", "shape": [1, 6, 1, 64] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * args.batch * args.heads * args.headDim * (args.keys + 1) * args.elemBytes + args.maskBytes + args.qkBytes",
+            "provenance": "logical"
+          }
+        ]
+      }
+    },
+    {
+      "name": "cross-mask-qk-b1-h6-kv64-d64-f16",
+      "preset": "smoke",
+      "vars": { "elemBytes": 2, "qkBytes": 768, "batch": 1, "heads": 6, "keys": 64, "headDim": 64, "maskBytes": 256 },
+      "attrs": { "num_heads": 6, "output_qk": 1, "mask_filter_value": -4 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 384], "dtype": "float16", "dist": "normal", "seed": 325, "scale": 0.2 },
+        "keyT": { "shape": [1, 6, 64, 64], "dtype": "float16", "dist": "normal", "seed": 326, "scale": 0.2 },
+        "valueT": { "shape": [1, 6, 64, 64], "dtype": "float16", "dist": "normal", "seed": 327, "scale": 0.2 },
+        "maskIndexT": { "dtype": "int32", "shape": [1, 64], "data": { "kind": "cycle", "values": [1, 0, 1] } }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 1, 384], "dtype": "float16" },
+        "qkT": { "dtype": "float16", "shape": [1, 6, 1, 64] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * args.batch * args.heads * args.headDim * (args.keys + 1) * args.elemBytes + args.maskBytes + args.qkBytes",
+            "provenance": "logical"
+          }
+        ]
+      }
+    },
+    {
+      "name": "cross-mask-qk-b1-h6-kv257-d64-f32",
+      "preset": "smoke",
+      "vars": { "elemBytes": 4, "qkBytes": 6168, "batch": 1, "heads": 6, "keys": 257, "headDim": 64, "maskBytes": 1028 },
+      "attrs": { "num_heads": 6, "output_qk": 1, "mask_filter_value": -4 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 384], "dtype": "float32", "dist": "normal", "seed": 325, "scale": 0.2 },
+        "keyT": { "shape": [1, 6, 257, 64], "dtype": "float32", "dist": "normal", "seed": 326, "scale": 0.2 },
+        "valueT": { "shape": [1, 6, 257, 64], "dtype": "float32", "dist": "normal", "seed": 327, "scale": 0.2 },
+        "maskIndexT": { "dtype": "int32", "shape": [1, 257], "data": { "kind": "cycle", "values": [1, 0, 1] } }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 1, 384], "dtype": "float32" },
+        "qkT": { "dtype": "float32", "shape": [1, 6, 1, 257] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * args.batch * args.heads * args.headDim * (args.keys + 1) * args.elemBytes + args.maskBytes + args.qkBytes",
+            "provenance": "logical"
+          }
+        ]
+      }
+    },
+    {
+      "name": "cross-mask-qk-b1-h6-kv257-d64-f16",
+      "preset": "smoke",
+      "vars": { "elemBytes": 2, "qkBytes": 3084, "batch": 1, "heads": 6, "keys": 257, "headDim": 64, "maskBytes": 1028 },
+      "attrs": { "num_heads": 6, "output_qk": 1, "mask_filter_value": -4 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 384], "dtype": "float16", "dist": "normal", "seed": 325, "scale": 0.2 },
+        "keyT": { "shape": [1, 6, 257, 64], "dtype": "float16", "dist": "normal", "seed": 326, "scale": 0.2 },
+        "valueT": { "shape": [1, 6, 257, 64], "dtype": "float16", "dist": "normal", "seed": 327, "scale": 0.2 },
+        "maskIndexT": { "dtype": "int32", "shape": [1, 257], "data": { "kind": "cycle", "values": [1, 0, 1] } }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 1, 384], "dtype": "float16" },
+        "qkT": { "dtype": "float16", "shape": [1, 6, 1, 257] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * args.batch * args.heads * args.headDim * (args.keys + 1) * args.elemBytes + args.maskBytes + args.qkBytes",
+            "provenance": "logical"
+          }
+        ]
+      }
+    },
+    {
+      "name": "cross-b1-h6-kv256-d32-f32",
+      "preset": "smoke",
+      "vars": { "elemBytes": 4, "qkBytes": 0, "batch": 1, "heads": 6, "keys": 256, "headDim": 32 },
+      "attrs": { "num_heads": 6 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 192], "dtype": "float32", "dist": "normal", "seed": 325, "scale": 0.2 },
+        "keyT": { "shape": [1, 6, 256, 32], "dtype": "float32", "dist": "normal", "seed": 326, "scale": 0.2 },
+        "valueT": { "shape": [1, 6, 256, 32], "dtype": "float32", "dist": "normal", "seed": 327, "scale": 0.2 }
+      },
+      "outputs": { "outputT": { "shape": [1, 1, 192], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * args.batch * args.heads * args.headDim * (args.keys + 1) * args.elemBytes",
+            "provenance": "logical"
+          }
+        ]
+      }
+    },
+    {
+      "name": "cross-b1-h6-kv256-d128-f32",
+      "preset": "smoke",
+      "vars": { "elemBytes": 4, "qkBytes": 0, "batch": 1, "heads": 6, "keys": 256, "headDim": 128 },
+      "attrs": { "num_heads": 6 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 768], "dtype": "float32", "dist": "normal", "seed": 325, "scale": 0.2 },
+        "keyT": { "shape": [1, 6, 256, 128], "dtype": "float32", "dist": "normal", "seed": 326, "scale": 0.2 },
+        "valueT": { "shape": [1, 6, 256, 128], "dtype": "float32", "dist": "normal", "seed": 327, "scale": 0.2 }
+      },
+      "outputs": { "outputT": { "shape": [1, 1, 768], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * args.batch * args.heads * args.headDim * (args.keys + 1) * args.elemBytes",
+            "provenance": "logical"
+          }
+        ]
+      }
+    },
+    {
+      "name": "cross-b1-h6-kv256-d256-f32",
+      "preset": "smoke",
+      "vars": { "elemBytes": 4, "qkBytes": 0, "batch": 1, "heads": 6, "keys": 256, "headDim": 256 },
+      "attrs": { "num_heads": 6 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 1536], "dtype": "float32", "dist": "normal", "seed": 325, "scale": 0.2 },
+        "keyT": { "shape": [1, 6, 256, 256], "dtype": "float32", "dist": "normal", "seed": 326, "scale": 0.2 },
+        "valueT": { "shape": [1, 6, 256, 256], "dtype": "float32", "dist": "normal", "seed": 327, "scale": 0.2 }
+      },
+      "outputs": { "outputT": { "shape": [1, 1, 1536], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * args.batch * args.heads * args.headDim * (args.keys + 1) * args.elemBytes",
+            "provenance": "logical"
+          }
+        ]
+      }
+    },
+    {
+      "name": "cross-b1-h6-kv1500-d64-f32",
+      "preset": "smoke",
+      "vars": { "elemBytes": 4, "qkBytes": 0, "batch": 1, "heads": 6, "keys": 1500, "headDim": 64 },
+      "attrs": { "num_heads": 6 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 384], "dtype": "float32", "dist": "normal", "seed": 325, "scale": 0.2 },
+        "keyT": { "shape": [1, 6, 1500, 64], "dtype": "float32", "dist": "normal", "seed": 326, "scale": 0.2 },
+        "valueT": { "shape": [1, 6, 1500, 64], "dtype": "float32", "dist": "normal", "seed": 327, "scale": 0.2 }
+      },
+      "outputs": { "outputT": { "shape": [1, 1, 384], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * args.batch * args.heads * args.headDim * (args.keys + 1) * args.elemBytes",
+            "provenance": "logical"
+          }
+        ]
+      }
+    },
+    {
+      "name": "cross-b2-h6-kv1500-d64-f32",
+      "preset": "smoke",
+      "vars": { "elemBytes": 4, "qkBytes": 0, "batch": 2, "heads": 6, "keys": 1500, "headDim": 64 },
+      "attrs": { "num_heads": 6 },
+      "inputs": {
+        "queryT": { "shape": [2, 1, 384], "dtype": "float32", "dist": "normal", "seed": 325, "scale": 0.2 },
+        "keyT": { "shape": [2, 6, 1500, 64], "dtype": "float32", "dist": "normal", "seed": 326, "scale": 0.2 },
+        "valueT": { "shape": [2, 6, 1500, 64], "dtype": "float32", "dist": "normal", "seed": 327, "scale": 0.2 }
+      },
+      "outputs": { "outputT": { "shape": [2, 1, 384], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * args.batch * args.heads * args.headDim * (args.keys + 1) * args.elemBytes",
+            "provenance": "logical"
+          }
+        ]
+      }
+    },
+    {
+      "name": "cross-b4-h6-kv1500-d64-f32",
+      "preset": "smoke",
+      "vars": { "elemBytes": 4, "qkBytes": 0, "batch": 4, "heads": 6, "keys": 1500, "headDim": 64 },
+      "attrs": { "num_heads": 6 },
+      "inputs": {
+        "queryT": { "shape": [4, 1, 384], "dtype": "float32", "dist": "normal", "seed": 325, "scale": 0.2 },
+        "keyT": { "shape": [4, 6, 1500, 64], "dtype": "float32", "dist": "normal", "seed": 326, "scale": 0.2 },
+        "valueT": { "shape": [4, 6, 1500, 64], "dtype": "float32", "dist": "normal", "seed": 327, "scale": 0.2 }
+      },
+      "outputs": { "outputT": { "shape": [4, 1, 384], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * args.batch * args.heads * args.headDim * (args.keys + 1) * args.elemBytes",
+            "provenance": "logical"
+          }
+        ]
+      }
+    },
+    {
+      "name": "qk-policy-b1-h4-kv256-d512-f32",
+      "preset": "smoke",
+      "vars": { "elemBytes": 4, "qkBytes": 0, "batch": 1, "heads": 4, "keys": 256, "headDim": 512 },
+      "attrs": { "num_heads": 4 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 2048], "dtype": "float32", "dist": "normal", "seed": 325, "scale": 0.2 },
+        "keyT": { "shape": [1, 4, 256, 512], "dtype": "float32", "dist": "normal", "seed": 326, "scale": 0.2 },
+        "valueT": { "shape": [1, 4, 256, 512], "dtype": "float32", "dist": "normal", "seed": 327, "scale": 0.2 }
+      },
+      "outputs": { "outputT": { "shape": [1, 1, 2048], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * args.batch * args.heads * args.headDim * (args.keys + 1) * args.elemBytes",
+            "provenance": "logical"
+          }
+        ]
+      }
+    },
+    {
+      "name": "qk-policy-b1-h2-kv256-d1024-f32",
+      "preset": "smoke",
+      "vars": { "elemBytes": 4, "qkBytes": 0, "batch": 1, "heads": 2, "keys": 256, "headDim": 1024 },
+      "attrs": { "num_heads": 2 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 2048], "dtype": "float32", "dist": "normal", "seed": 325, "scale": 0.2 },
+        "keyT": { "shape": [1, 2, 256, 1024], "dtype": "float32", "dist": "normal", "seed": 326, "scale": 0.2 },
+        "valueT": { "shape": [1, 2, 256, 1024], "dtype": "float32", "dist": "normal", "seed": 327, "scale": 0.2 }
+      },
+      "outputs": { "outputT": { "shape": [1, 1, 2048], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * args.batch * args.heads * args.headDim * (args.keys + 1) * args.elemBytes",
+            "provenance": "logical"
+          }
+        ]
+      }
+    },
+    {
+      "name": "qk-policy-b1-h15-kv256-d128-f32",
+      "preset": "smoke",
+      "vars": { "elemBytes": 4, "qkBytes": 0, "batch": 1, "heads": 15, "keys": 256, "headDim": 128 },
+      "attrs": { "num_heads": 15 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 1920], "dtype": "float32", "dist": "normal", "seed": 325, "scale": 0.2 },
+        "keyT": { "shape": [1, 15, 256, 128], "dtype": "float32", "dist": "normal", "seed": 326, "scale": 0.2 },
+        "valueT": { "shape": [1, 15, 256, 128], "dtype": "float32", "dist": "normal", "seed": 327, "scale": 0.2 }
+      },
+      "outputs": { "outputT": { "shape": [1, 1, 1920], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * args.batch * args.heads * args.headDim * (args.keys + 1) * args.elemBytes",
+            "provenance": "logical"
+          }
+        ]
+      }
+    },
+    {
+      "name": "qk-policy-b1-h16-kv256-d128-f32",
+      "preset": "smoke",
+      "vars": { "elemBytes": 4, "qkBytes": 0, "batch": 1, "heads": 16, "keys": 256, "headDim": 128 },
+      "attrs": { "num_heads": 16 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 2048], "dtype": "float32", "dist": "normal", "seed": 325, "scale": 0.2 },
+        "keyT": { "shape": [1, 16, 256, 128], "dtype": "float32", "dist": "normal", "seed": 326, "scale": 0.2 },
+        "valueT": { "shape": [1, 16, 256, 128], "dtype": "float32", "dist": "normal", "seed": 327, "scale": 0.2 }
+      },
+      "outputs": { "outputT": { "shape": [1, 1, 2048], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * args.batch * args.heads * args.headDim * (args.keys + 1) * args.elemBytes",
+            "provenance": "logical"
+          }
+        ]
+      }
+    },
+    {
+      "name": "qk-policy-b1-h17-kv256-d128-f32",
+      "preset": "smoke",
+      "vars": { "elemBytes": 4, "qkBytes": 0, "batch": 1, "heads": 17, "keys": 256, "headDim": 128 },
+      "attrs": { "num_heads": 17 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 2176], "dtype": "float32", "dist": "normal", "seed": 325, "scale": 0.2 },
+        "keyT": { "shape": [1, 17, 256, 128], "dtype": "float32", "dist": "normal", "seed": 326, "scale": 0.2 },
+        "valueT": { "shape": [1, 17, 256, 128], "dtype": "float32", "dist": "normal", "seed": 327, "scale": 0.2 }
+      },
+      "outputs": { "outputT": { "shape": [1, 1, 2176], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * args.batch * args.heads * args.headDim * (args.keys + 1) * args.elemBytes",
+            "provenance": "logical"
+          }
+        ]
+      }
+    },
+    {
+      "name": "qk-policy-b1-h15-kv1500-d128-f32",
+      "preset": "smoke",
+      "vars": { "elemBytes": 4, "qkBytes": 0, "batch": 1, "heads": 15, "keys": 1500, "headDim": 128 },
+      "attrs": { "num_heads": 15 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 1920], "dtype": "float32", "dist": "normal", "seed": 325, "scale": 0.2 },
+        "keyT": { "shape": [1, 15, 1500, 128], "dtype": "float32", "dist": "normal", "seed": 326, "scale": 0.2 },
+        "valueT": { "shape": [1, 15, 1500, 128], "dtype": "float32", "dist": "normal", "seed": 327, "scale": 0.2 }
+      },
+      "outputs": { "outputT": { "shape": [1, 1, 1920], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * args.batch * args.heads * args.headDim * (args.keys + 1) * args.elemBytes",
+            "provenance": "logical"
+          }
+        ]
+      }
+    },
+    {
+      "name": "qk-policy-b1-h16-kv1500-d128-f32",
+      "preset": "smoke",
+      "vars": { "elemBytes": 4, "qkBytes": 0, "batch": 1, "heads": 16, "keys": 1500, "headDim": 128 },
+      "attrs": { "num_heads": 16 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 2048], "dtype": "float32", "dist": "normal", "seed": 325, "scale": 0.2 },
+        "keyT": { "shape": [1, 16, 1500, 128], "dtype": "float32", "dist": "normal", "seed": 326, "scale": 0.2 },
+        "valueT": { "shape": [1, 16, 1500, 128], "dtype": "float32", "dist": "normal", "seed": 327, "scale": 0.2 }
+      },
+      "outputs": { "outputT": { "shape": [1, 1, 2048], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * args.batch * args.heads * args.headDim * (args.keys + 1) * args.elemBytes",
+            "provenance": "logical"
+          }
+        ]
+      }
+    },
+    {
+      "name": "qk-policy-b1-h17-kv1500-d128-f32",
+      "preset": "smoke",
+      "vars": { "elemBytes": 4, "qkBytes": 0, "batch": 1, "heads": 17, "keys": 1500, "headDim": 128 },
+      "attrs": { "num_heads": 17 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 2176], "dtype": "float32", "dist": "normal", "seed": 325, "scale": 0.2 },
+        "keyT": { "shape": [1, 17, 1500, 128], "dtype": "float32", "dist": "normal", "seed": 326, "scale": 0.2 },
+        "valueT": { "shape": [1, 17, 1500, 128], "dtype": "float32", "dist": "normal", "seed": 327, "scale": 0.2 }
+      },
+      "outputs": { "outputT": { "shape": [1, 1, 2176], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * args.batch * args.heads * args.headDim * (args.keys + 1) * args.elemBytes",
+            "provenance": "logical"
+          }
+        ]
+      }
+    }
+  ]
+}

build/webgpu/dmmha-present.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,79 @@

+{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
+{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
+  if ({{ name }} >= {{ bound }}) {
+    return;
+  }{% endmacro %}
+{{ env.wgsl.resourceDeclarations }}
+// Cache maintenance for decode self attention, in two roles from one template.
+//
+//   retain — copy the whole past buffer into the present buffer at capacity.
+//            The rows beyond the live length are copied too: the operator's
+//            contract is that present equals past everywhere except the one
+//            appended row, and upstream's own kernel memcpys the full buffer.
+//            The pass is a pure copy, so it declares `viewAlias`: it is elided
+//            when `past_key` and `present_key` resolve to one buffer.
+//   append — write this token's key and value rows at index past_sequence_length,
+//            which arrives as tensor data and so is read here rather than baked.
+//
+// Key and value share one dispatch: they have identical geometry, and one
+// invocation touching both halves keeps the index arithmetic in registers.
+const WG: u32 = {{ workgroupSize }}u;
+{% if appendRow %}
+const HEAD_SIZE_V4: u32 = {{ headSizeV4 }}u;
+const NUM_HEADS: u32 = {{ numHeads }}u;
+const HIDDEN_V4: u32 = NUM_HEADS * HEAD_SIZE_V4;
+@compute @workgroup_size(WG)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+  // One invocation owns one vec4 of one (batch, head) row.
+{{ flat_index_2d("WG", "index", "params.rowCountV4") }}
+  let d4 = index % HEAD_SIZE_V4;
+  let head = (index / HEAD_SIZE_V4) % NUM_HEADS;
+  let batch = index / HIDDEN_V4;
+  // past_sequence_length is a single value shared by every batch row.
+  let slot = u32(past_sequence_length[0]);
+  if (slot >= params.capacity) {
+    return;
+  }
+  let cacheIndex = ((batch * NUM_HEADS + head) * params.capacity + slot) * HEAD_SIZE_V4 + d4;
+  present_key[cacheIndex] = key[index];
+  present_value[cacheIndex] = value[index];
+}
+{% else %}
+// Each lane stages ITEMS words of each cache before storing any of them, so that
+// many loads are in flight before it has to wait. One word per lane leaves the
+// copy latency-bound wherever the caches stay resident in cache; where the copy is
+// DRAM-bound the depth is neutral. Consecutive lanes still read consecutive words;
+// the stride between a lane's own items is a whole workgroup. ITEMS is bounded so
+// the dispatch keeps a floor number of workgroups, since raising it divides them.
+const ITEMS: u32 = {{ itemsPerLane }}u;
+@compute @workgroup_size(WG)
+fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
+  let block = (wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u) * WG * ITEMS + lid.x;
+  var stagedKey: array<vec4<{{ cacheScalar }}>, ITEMS>;
+  var stagedValue: array<vec4<{{ cacheScalar }}>, ITEMS>;
+{% for i in range(itemsPerLane) %}
+  {
+    let index = block + {{ i }}u * WG;
+    if (index < params.countV4) {
+      stagedKey[{{ i }}] = past_key[index];
+      stagedValue[{{ i }}] = past_value[index];
+    }
+  }
+{% endfor %}
+{% for i in range(itemsPerLane) %}
+  {
+    let index = block + {{ i }}u * WG;
+    if (index < params.countV4) {
+      present_key[index] = stagedKey[{{ i }}];
+      present_value[index] = stagedValue[{{ i }}];
+    }
+  }
+{% endfor %}
+}
+{% endif %}

build/webgpu/manifest.json ADDED Viewed

	@@ -0,0 +1,1535 @@

+{
+  "domain": "com.microsoft",
+  "name": "DecoderMaskedMultiHeadAttention",
+  "sinceVersion": 1,
+  "inputs": {
+    "queryT": { "onnx": "query", "dtype": "T", "rank": 3 },
+    "keyT": { "onnx": "key", "dtype": "T" },
+    "valueT": { "onnx": "value", "dtype": "T" },
+    "maskIndexT": { "onnx": "mask_index", "dtype": "M", "rank": 2, "optional": true },
+    "attentionBiasT": { "onnx": "attention_bias", "dtype": "T", "rank": 4, "optional": true },
+    "pastKeyT": { "onnx": "past_key", "dtype": "T", "rank": 4, "optional": true },
+    "pastValueT": { "onnx": "past_value", "dtype": "T", "rank": 4, "optional": true },
+    "pastSequenceLengthT": { "onnx": "past_sequence_length", "dtype": "M", "rank": 1, "optional": true }
+  },
+  "outputs": {
+    "outputT": {
+      "onnx": "output",
+      "dtype": "T",
+      "rank": 3,
+      "shape": "[dim(shapes.queryT, 0), 1, dim(shapes.queryT, 2)]"
+    },
+    "presentKeyT": { "onnx": "present_key", "dtype": "T", "rank": 4, "optional": true, "shape": "shapes.pastKeyT" },
+    "presentValueT": {
+      "onnx": "present_value",
+      "dtype": "T",
+      "rank": 4,
+      "optional": true,
+      "shape": "shapes.pastValueT"
+    },
+    "qkT": {
+      "onnx": "qk",
+      "dtype": "QK",
+      "rank": 4,
+      "optional": true,
+      "shape": "[dim(shapes.queryT, 0), attrs.num_heads, 1, dim(shapes.keyT, 2)]"
+    }
+  },
+  "attributes": {
+    "num_heads": {},
+    "past_present_share_buffer": {},
+    "scale": {},
+    "mask_filter_value": { "default": -10000 },
+    "output_qk": {}
+  },
+  "attributeConstraints": { "num_heads": { "required": true } },
+  "typeConstraints": { "T": ["float32", "float16"], "QK": ["float32", "float16"], "M": ["int32"] },
+  "tunables": {
+    "WORKGROUP_SIZE": { "default": 256 },
+    "COPY_WORKGROUP_SIZE": { "default": 256 },
+    "COPY_ITEMS_PER_LANE": { "default": 8 },
+    "COPY_MIN_WORKGROUPS": { "default": 64 },
+    "MAX_SPLITS": { "default": 8 },
+    "KEYS_PER_SPLIT": { "default": 256 },
+    "THREAD_QK_OUTPUT_WORKGROUPS": { "default": 2 }
+  },
+  "derive": {
+    "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
+    "narrowSubgroupRange": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize < device.adapterInfo.subgroupMaxSize and device.adapterInfo.subgroupMaxSize <= 16",
+    "batchSize": "dim(shapes.queryT, 0)",
+    "hiddenSize": "dim(shapes.queryT, 2)",
+    "numHeads": "attrs.num_heads",
+    "headSize": "hiddenSize / max(1, numHeads)",
+    "headDimV4": "headSize / 4",
+    "isCross": "not present.pastKeyT and not present.pastValueT",
+    "crossKvLen": "dim(shapes.keyT, 2) if (isCross and ranks.keyT == 4) else 0",
+    "cacheCapacity": "dim(shapes.pastKeyT, 2) if present.pastKeyT else 0",
+    "physicalKeys": "crossKvLen if isCross else cacheCapacity",
+    "qkRequested": "attrs.output_qk != 0 if has(attrs, \"output_qk\") else false",
+    "floatDtypeOk": "(tensorDtypes.queryT == \"float32\" or tensorDtypes.queryT == \"float16\") and f16Ok(tensorDtypes.queryT) and tensorDtypes.keyT == tensorDtypes.queryT and tensorDtypes.valueT == tensorDtypes.queryT and tensorDtypes.outputT == tensorDtypes.queryT",
+    "geometryOk": "ranks.queryT == 3 and dim(shapes.queryT, 1) == 1 and numHeads > 0 and hiddenSize > 0 and hiddenSize % numHeads == 0 and headSize > 0 and headSize % 4 == 0",
+    "outputShapeOk": "ranks.outputT == 3 and dim(shapes.outputT, 0) == batchSize and dim(shapes.outputT, 1) == 1 and dim(shapes.outputT, 2) == hiddenSize",
+    "workgroupOk": "tunables.WORKGROUP_SIZE > 0 and tunables.WORKGROUP_SIZE <= deviceWorkgroupCap and pow2ceil(tunables.WORKGROUP_SIZE) == tunables.WORKGROUP_SIZE and headDimV4 <= deviceWorkgroupCap and tunables.COPY_WORKGROUP_SIZE > 0 and tunables.COPY_WORKGROUP_SIZE <= deviceWorkgroupCap",
+    "workgroupStorageOk": "32 * headDimV4 + 24 * tunables.WORKGROUP_SIZE <= device.limits.maxComputeWorkgroupStorageSize",
+    "numSplits": "max(1, min(tunables.MAX_SPLITS, ceilDiv(physicalKeys, tunables.KEYS_PER_SPLIT)))",
+    "splitScratchBytes": "batchSize * numHeads * numSplits * headSize * 4",
+    "splitStatsBytes": "2 * batchSize * numHeads * numSplits * 4",
+    "scratchFits": "splitScratchBytes <= device.limits.maxStorageBufferBindingSize and splitScratchBytes <= device.limits.maxBufferSize and splitStatsBytes <= device.limits.maxStorageBufferBindingSize and splitStatsBytes <= device.limits.maxBufferSize",
+    "dispatchFits": "batchSize <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and numHeads <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and numSplits <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
+    "coreOk": "floatDtypeOk and geometryOk and outputShapeOk and workgroupOk and workgroupStorageOk and scratchFits and dispatchFits and physicalKeys > 0",
+    "crossOk": "coreOk and isCross and ranks.keyT == 4 and ranks.valueT == 4 and dim(shapes.keyT, 0) == batchSize and dim(shapes.keyT, 1) == numHeads and dim(shapes.keyT, 3) == headSize and sameShape(shapes.valueT, shapes.keyT) and not present.attentionBiasT and not present.pastSequenceLengthT and not present.presentKeyT and not present.presentValueT",
+    "cacheDtypeOk": "tensorDtypes.pastKeyT == tensorDtypes.queryT and tensorDtypes.pastValueT == tensorDtypes.queryT and tensorDtypes.presentKeyT == tensorDtypes.queryT and tensorDtypes.presentValueT == tensorDtypes.queryT and tensorDtypes.pastSequenceLengthT == \"int32\"",
+    "selfOk": "coreOk and not isCross and present.pastValueT and present.pastSequenceLengthT and present.presentKeyT and present.presentValueT and attrs.past_present_share_buffer == 1 and ranks.keyT == 3 and ranks.valueT == 3 and dim(shapes.keyT, 0) == batchSize and dim(shapes.keyT, 1) == 1 and dim(shapes.keyT, 2) == hiddenSize and sameShape(shapes.valueT, shapes.keyT) and ranks.pastKeyT == 4 and dim(shapes.pastKeyT, 0) == batchSize and dim(shapes.pastKeyT, 1) == numHeads and dim(shapes.pastKeyT, 3) == headSize and cacheCapacity > 1 and sameShape(shapes.pastValueT, shapes.pastKeyT) and sameShape(shapes.presentKeyT, shapes.pastKeyT) and sameShape(shapes.presentValueT, shapes.pastKeyT) and numel(shapes.pastSequenceLengthT) == 1 and cacheDtypeOk and not present.qkT and not qkRequested",
+    "maskShapeOk": "present.maskIndexT and ranks.maskIndexT == 2 and dim(shapes.maskIndexT, 0) == batchSize and tensorDtypes.maskIndexT == \"int32\" and (dim(shapes.maskIndexT, 1) == crossKvLen if isCross else (dim(shapes.maskIndexT, 1) >= 1 and dim(shapes.maskIndexT, 1) <= cacheCapacity))",
+    "attnBiasShapeOk": "present.attentionBiasT and ranks.attentionBiasT == 4 and tensorDtypes.attentionBiasT == tensorDtypes.queryT and dim(shapes.attentionBiasT, 2) == 1 and (dim(shapes.attentionBiasT, 0) == 1 or dim(shapes.attentionBiasT, 0) == batchSize) and (dim(shapes.attentionBiasT, 1) == 1 or dim(shapes.attentionBiasT, 1) == numHeads) and dim(shapes.attentionBiasT, 3) >= 1 and dim(shapes.attentionBiasT, 3) <= cacheCapacity",
+    "qkShapeOk": "present.qkT and qkRequested and ranks.qkT == 4 and (tensorDtypes.qkT == \"float32\" or tensorDtypes.qkT == \"float16\") and f16Ok(tensorDtypes.qkT) and dim(shapes.qkT, 0) == batchSize and dim(shapes.qkT, 1) == numHeads and dim(shapes.qkT, 2) == 1 and dim(shapes.qkT, 3) == crossKvLen",
+    "crossPlainOk": "crossOk and not present.maskIndexT and not present.qkT and not qkRequested",
+    "crossMaskOk": "crossOk and maskShapeOk and not present.qkT and not qkRequested",
+    "crossQkOk": "crossOk and not present.maskIndexT and qkShapeOk",
+    "crossMaskQkOk": "crossOk and maskShapeOk and qkShapeOk",
+    "selfPlainOk": "selfOk and not present.maskIndexT and not present.attentionBiasT",
+    "selfMaskOk": "selfOk and maskShapeOk and not present.attentionBiasT",
+    "selfBiasOk": "selfOk and not present.maskIndexT and attnBiasShapeOk",
+    "selfMaskBiasOk": "selfOk and maskShapeOk and attnBiasShapeOk",
+    "subgroupSplitKOk": "device.features.has(\"subgroups\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize >= 8 and device.adapterInfo.subgroupMinSize % 8 == 0 and device.adapterInfo.subgroupMaxSize % 8 == 0 and device.adapterInfo.subgroupMaxSize <= tunables.WORKGROUP_SIZE and tunables.WORKGROUP_SIZE % device.adapterInfo.subgroupMinSize == 0 and tunables.WORKGROUP_SIZE % device.adapterInfo.subgroupMaxSize == 0",
+    "subgroupSplitKPreferred": "subgroupSplitKOk and not narrowSubgroupRange",
+    "attnScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
+    "attnVec4": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
+    "attnMaskElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
+    "qkScalarType": "\"f16\" if tensorDtypes.qkT == \"float16\" else \"f32\"",
+    "attnBiasWidth": "dim(shapes.attentionBiasT, 3) if present.attentionBiasT else 0",
+    "attnBiasHeads": "dim(shapes.attentionBiasT, 1) if present.attentionBiasT else 0",
+    "attnBiasBatch": "dim(shapes.attentionBiasT, 0) if present.attentionBiasT else 0",
+    "cacheWordsV4": "numel(shapes.pastKeyT) / 4 if present.pastKeyT else 0",
+    "retainBlockWords": "tunables.COPY_WORKGROUP_SIZE * tunables.COPY_MIN_WORKGROUPS",
+    "retainOccupancyCap": "pick([[cacheWordsV4 >= 8 * retainBlockWords, 8], [cacheWordsV4 >= 4 * retainBlockWords, 4], [cacheWordsV4 >= 2 * retainBlockWords, 2]], 1)",
+    "retainItemsPerLane": "max(1, min(tunables.COPY_ITEMS_PER_LANE, retainOccupancyCap))"
+  },
+  "bindings": {
+    "query": { "arg": "queryT", "elementType": "$attnVec4" },
+    "key_cross": { "arg": "keyT", "name": "key", "elementType": "$attnVec4" },
+    "value_cross": { "arg": "valueT", "name": "value", "elementType": "$attnVec4" },
+    "key_cache": { "arg": "presentKeyT", "name": "key", "buffer": "read-only-storage", "elementType": "$attnVec4" },
+    "value_cache": {
+      "arg": "presentValueT",
+      "name": "value",
+      "buffer": "read-only-storage",
+      "elementType": "$attnVec4"
+    },
+    "key_row": { "arg": "keyT", "name": "key", "elementType": "$attnVec4" },
+    "value_row": { "arg": "valueT", "name": "value", "elementType": "$attnVec4" },
+    "past_key": { "arg": "pastKeyT", "elementType": "$attnVec4" },
+    "past_value": { "arg": "pastValueT", "elementType": "$attnVec4" },
+    "present_key": { "arg": "presentKeyT", "elementType": "$attnVec4" },
+    "present_value": { "arg": "presentValueT", "elementType": "$attnVec4" },
+    "seqlens_k": { "arg": "pastSequenceLengthT", "name": "seqlens_k", "elementType": "i32" },
+    "past_sequence_length": { "arg": "pastSequenceLengthT", "elementType": "i32" },
+    "key_pad_mask": { "arg": "maskIndexT", "name": "key_pad_mask", "elementType": "i32" },
+    "attn_mask": { "arg": "attentionBiasT", "name": "attn_mask", "elementType": "$attnMaskElement" },
+    "qk": { "arg": "qkT", "elementType": "$qkScalarType" },
+    "partial_out": { "scratch": "partialOut", "elementType": "vec4<f32>" },
+    "partial_stats": { "scratch": "partialStats", "elementType": "vec2<f32>" },
+    "partial_out_read": {
+      "scratch": "partialOut",
+      "name": "partial_out",
+      "buffer": "read-only-storage",
+      "elementType": "vec4<f32>"
+    },
+    "partial_stats_read": {
+      "scratch": "partialStats",
+      "name": "partial_stats",
+      "buffer": "read-only-storage",
+      "elementType": "vec2<f32>"
+    },
+    "output": { "arg": "outputT", "elementType": "$attnVec4" },
+    "params_cross": {
+      "name": "params",
+      "struct": [
+        { "name": "kvSeq", "type": "u32", "value": "crossKvLen" },
+        { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" }
+      ]
+    },
+    "params_self": {
+      "name": "params",
+      "struct": [
+        { "name": "kvSeq", "type": "u32", "value": "cacheCapacity" },
+        { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" }
+      ]
+    },
+    "params_self_bias": {
+      "name": "params",
+      "struct": [
+        { "name": "kvSeq", "type": "u32", "value": "cacheCapacity" },
+        { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" },
+        {
+          "name": "maskBatchStride",
+          "type": "u32",
+          "value": "0 if attnBiasBatch == 1 else attnBiasHeads * attnBiasWidth"
+        },
+        { "name": "maskHeadStride", "type": "u32", "value": "0 if attnBiasHeads == 1 else attnBiasWidth" },
+        { "name": "maskSeqStride", "type": "u32", "value": "0" }
+      ]
+    },
+    "params_retain": {
+      "name": "params",
+      "struct": [{ "name": "countV4", "type": "u32", "value": "numel(shapes.pastKeyT) / 4" }]
+    },
+    "params_append": {
+      "name": "params",
+      "struct": [
+        { "name": "rowCountV4", "type": "u32", "value": "batchSize * hiddenSize / 4" },
+        { "name": "capacity", "type": "u32", "value": "cacheCapacity" }
+      ]
+    }
+  },
+  "variants": [
+    {
+      "id": "cross_direct",
+      "priority": 31,
+      "when": ["crossPlainOk", "subgroupSplitKOk", "numSplits == 1"],
+      "demoteWhen": ["not subgroupSplitKPreferred"],
+      "requires": { "features": ["subgroups"] },
+      "passes": [
+        {
+          "id": "split_attention",
+          "name": "DecoderMaskedMultiHeadAttention.CrossSplitK",
+          "shader": "attn-flash-decode-splitk.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "qNumHeads": "numHeads",
+            "kvNumHeads": "numHeads",
+            "headDim": "headSize",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits",
+            "usesF16": "tensorDtypes.queryT == \"float16\"",
+            "hasKeyPadMask": "present.maskIndexT",
+            "emitQk": "present.qkT",
+            "useSubgroups": true,
+            "allowCooperativeQk": "not useSubgroups or usesF16 or headDimV4 < 8 or batchSize * hiddenSize / 4 > tunables.THREAD_QK_OUTPUT_WORKGROUPS * tunables.WORKGROUP_SIZE"
+          },
+          "bindings": ["query", "key_cross", "value_cross", "params_cross", "output"],
+          "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        }
+      ],
+      "derive": { "directOutput": true }
+    },
+    {
+      "id": "cross",
+      "priority": 30,
+      "when": ["crossPlainOk", "subgroupSplitKOk"],
+      "demoteWhen": ["not subgroupSplitKPreferred"],
+      "requires": { "features": ["subgroups"] },
+      "intermediates": [
+        { "id": "partialOut", "dtype": "float32", "shape": "[batchSize * numHeads * numSplits * headSize]" },
+        { "id": "partialStats", "dtype": "float32", "shape": "[2 * batchSize * numHeads * numSplits]" }
+      ],
+      "passes": [
+        {
+          "id": "split_attention",
+          "name": "DecoderMaskedMultiHeadAttention.CrossSplitK",
+          "shader": "attn-flash-decode-splitk.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "qNumHeads": "numHeads",
+            "kvNumHeads": "numHeads",
+            "headDim": "headSize",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits",
+            "usesF16": "tensorDtypes.queryT == \"float16\"",
+            "hasKeyPadMask": "present.maskIndexT",
+            "emitQk": "present.qkT",
+            "useSubgroups": true,
+            "allowCooperativeQk": "not useSubgroups or usesF16 or headDimV4 < 8 or batchSize * hiddenSize / 4 > tunables.THREAD_QK_OUTPUT_WORKGROUPS * tunables.WORKGROUP_SIZE"
+          },
+          "bindings": ["query", "key_cross", "value_cross", "partial_out", "partial_stats", "params_cross"],
+          "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        },
+        {
+          "id": "merge",
+          "name": "DecoderMaskedMultiHeadAttention.CrossMerge",
+          "shader": "attn-flash-decode-splitk-merge.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "scalar": "attnScalar",
+            "qNumHeads": "numHeads",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits"
+          },
+          "bindings": ["partial_out_read", "partial_stats_read", "output"],
+          "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        }
+      ]
+    },
+    {
+      "id": "cross_qk_direct",
+      "priority": 31,
+      "when": ["crossQkOk", "subgroupSplitKOk", "numSplits == 1"],
+      "demoteWhen": ["not subgroupSplitKPreferred"],
+      "requires": { "features": ["subgroups"] },
+      "passes": [
+        {
+          "id": "split_attention",
+          "name": "DecoderMaskedMultiHeadAttention.CrossSplitKQk",
+          "shader": "attn-flash-decode-splitk.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "qNumHeads": "numHeads",
+            "kvNumHeads": "numHeads",
+            "headDim": "headSize",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits",
+            "usesF16": "tensorDtypes.queryT == \"float16\"",
+            "hasKeyPadMask": "present.maskIndexT",
+            "emitQk": "present.qkT",
+            "qkScalar": "qkScalarType",
+            "useSubgroups": true,
+            "allowCooperativeQk": "not useSubgroups or usesF16 or headDimV4 < 8 or batchSize * hiddenSize / 4 > tunables.THREAD_QK_OUTPUT_WORKGROUPS * tunables.WORKGROUP_SIZE"
+          },
+          "bindings": ["query", "key_cross", "value_cross", "qk", "params_cross", "output"],
+          "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        }
+      ],
+      "derive": { "directOutput": true }
+    },
+    {
+      "id": "cross_qk",
+      "priority": 30,
+      "when": ["crossQkOk", "subgroupSplitKOk"],
+      "demoteWhen": ["not subgroupSplitKPreferred"],
+      "requires": { "features": ["subgroups"] },
+      "intermediates": [
+        { "id": "partialOut", "dtype": "float32", "shape": "[batchSize * numHeads * numSplits * headSize]" },
+        { "id": "partialStats", "dtype": "float32", "shape": "[2 * batchSize * numHeads * numSplits]" }
+      ],
+      "passes": [
+        {
+          "id": "split_attention",
+          "name": "DecoderMaskedMultiHeadAttention.CrossSplitKQk",
+          "shader": "attn-flash-decode-splitk.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "qNumHeads": "numHeads",
+            "kvNumHeads": "numHeads",
+            "headDim": "headSize",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits",
+            "usesF16": "tensorDtypes.queryT == \"float16\"",
+            "hasKeyPadMask": "present.maskIndexT",
+            "emitQk": "present.qkT",
+            "qkScalar": "qkScalarType",
+            "useSubgroups": true,
+            "allowCooperativeQk": "not useSubgroups or usesF16 or headDimV4 < 8 or batchSize * hiddenSize / 4 > tunables.THREAD_QK_OUTPUT_WORKGROUPS * tunables.WORKGROUP_SIZE"
+          },
+          "bindings": ["query", "key_cross", "value_cross", "qk", "partial_out", "partial_stats", "params_cross"],
+          "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        },
+        {
+          "id": "merge",
+          "name": "DecoderMaskedMultiHeadAttention.CrossMerge",
+          "shader": "attn-flash-decode-splitk-merge.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "scalar": "attnScalar",
+            "qNumHeads": "numHeads",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits"
+          },
+          "bindings": ["partial_out_read", "partial_stats_read", "output"],
+          "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        }
+      ]
+    },
+    {
+      "id": "cross_mask_direct",
+      "priority": 31,
+      "when": ["crossMaskOk", "subgroupSplitKOk", "numSplits == 1"],
+      "demoteWhen": ["not subgroupSplitKPreferred"],
+      "requires": { "features": ["subgroups"] },
+      "passes": [
+        {
+          "id": "split_attention",
+          "name": "DecoderMaskedMultiHeadAttention.CrossSplitK",
+          "shader": "attn-flash-decode-splitk.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "qNumHeads": "numHeads",
+            "kvNumHeads": "numHeads",
+            "headDim": "headSize",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits",
+            "usesF16": "tensorDtypes.queryT == \"float16\"",
+            "hasKeyPadMask": "present.maskIndexT",
+            "maskFilterValue": "attrs.mask_filter_value",
+            "emitQk": "present.qkT",
+            "useSubgroups": true,
+            "allowCooperativeQk": "not useSubgroups or usesF16 or headDimV4 < 8 or batchSize * hiddenSize / 4 > tunables.THREAD_QK_OUTPUT_WORKGROUPS * tunables.WORKGROUP_SIZE"
+          },
+          "bindings": ["query", "key_cross", "value_cross", "key_pad_mask", "params_cross", "output"],
+          "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        }
+      ],
+      "derive": { "directOutput": true }
+    },
+    {
+      "id": "cross_mask",
+      "priority": 30,
+      "when": ["crossMaskOk", "subgroupSplitKOk"],
+      "demoteWhen": ["not subgroupSplitKPreferred"],
+      "requires": { "features": ["subgroups"] },
+      "intermediates": [
+        { "id": "partialOut", "dtype": "float32", "shape": "[batchSize * numHeads * numSplits * headSize]" },
+        { "id": "partialStats", "dtype": "float32", "shape": "[2 * batchSize * numHeads * numSplits]" }
+      ],
+      "passes": [
+        {
+          "id": "split_attention",
+          "name": "DecoderMaskedMultiHeadAttention.CrossSplitK",
+          "shader": "attn-flash-decode-splitk.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "qNumHeads": "numHeads",
+            "kvNumHeads": "numHeads",
+            "headDim": "headSize",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits",
+            "usesF16": "tensorDtypes.queryT == \"float16\"",
+            "hasKeyPadMask": "present.maskIndexT",
+            "maskFilterValue": "attrs.mask_filter_value",
+            "emitQk": "present.qkT",
+            "useSubgroups": true,
+            "allowCooperativeQk": "not useSubgroups or usesF16 or headDimV4 < 8 or batchSize * hiddenSize / 4 > tunables.THREAD_QK_OUTPUT_WORKGROUPS * tunables.WORKGROUP_SIZE"
+          },
+          "bindings": ["query", "key_cross", "value_cross", "key_pad_mask", "partial_out", "partial_stats", "params_cross"],
+          "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        },
+        {
+          "id": "merge",
+          "name": "DecoderMaskedMultiHeadAttention.CrossMerge",
+          "shader": "attn-flash-decode-splitk-merge.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "scalar": "attnScalar",
+            "qNumHeads": "numHeads",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits"
+          },
+          "bindings": ["partial_out_read", "partial_stats_read", "output"],
+          "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        }
+      ]
+    },
+    {
+      "id": "cross_mask_qk_direct",
+      "priority": 31,
+      "when": ["crossMaskQkOk", "subgroupSplitKOk", "numSplits == 1"],
+      "demoteWhen": ["not subgroupSplitKPreferred"],
+      "requires": { "features": ["subgroups"] },
+      "passes": [
+        {
+          "id": "split_attention",
+          "name": "DecoderMaskedMultiHeadAttention.CrossSplitKQk",
+          "shader": "attn-flash-decode-splitk.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "qNumHeads": "numHeads",
+            "kvNumHeads": "numHeads",
+            "headDim": "headSize",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits",
+            "usesF16": "tensorDtypes.queryT == \"float16\"",
+            "hasKeyPadMask": "present.maskIndexT",
+            "maskFilterValue": "attrs.mask_filter_value",
+            "emitQk": "present.qkT",
+            "qkScalar": "qkScalarType",
+            "useSubgroups": true,
+            "allowCooperativeQk": "not useSubgroups or usesF16 or headDimV4 < 8 or batchSize * hiddenSize / 4 > tunables.THREAD_QK_OUTPUT_WORKGROUPS * tunables.WORKGROUP_SIZE"
+          },
+          "bindings": ["query", "key_cross", "value_cross", "key_pad_mask", "qk", "params_cross", "output"],
+          "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        }
+      ],
+      "derive": { "directOutput": true }
+    },
+    {
+      "id": "cross_mask_qk",
+      "priority": 30,
+      "when": ["crossMaskQkOk", "subgroupSplitKOk"],
+      "demoteWhen": ["not subgroupSplitKPreferred"],
+      "requires": { "features": ["subgroups"] },
+      "intermediates": [
+        { "id": "partialOut", "dtype": "float32", "shape": "[batchSize * numHeads * numSplits * headSize]" },
+        { "id": "partialStats", "dtype": "float32", "shape": "[2 * batchSize * numHeads * numSplits]" }
+      ],
+      "passes": [
+        {
+          "id": "split_attention",
+          "name": "DecoderMaskedMultiHeadAttention.CrossSplitKQk",
+          "shader": "attn-flash-decode-splitk.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "qNumHeads": "numHeads",
+            "kvNumHeads": "numHeads",
+            "headDim": "headSize",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits",
+            "usesF16": "tensorDtypes.queryT == \"float16\"",
+            "hasKeyPadMask": "present.maskIndexT",
+            "maskFilterValue": "attrs.mask_filter_value",
+            "emitQk": "present.qkT",
+            "qkScalar": "qkScalarType",
+            "useSubgroups": true,
+            "allowCooperativeQk": "not useSubgroups or usesF16 or headDimV4 < 8 or batchSize * hiddenSize / 4 > tunables.THREAD_QK_OUTPUT_WORKGROUPS * tunables.WORKGROUP_SIZE"
+          },
+          "bindings": ["query", "key_cross", "value_cross", "key_pad_mask", "qk", "partial_out", "partial_stats", "params_cross"],
+          "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        },
+        {
+          "id": "merge",
+          "name": "DecoderMaskedMultiHeadAttention.CrossMerge",
+          "shader": "attn-flash-decode-splitk-merge.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "scalar": "attnScalar",
+            "qNumHeads": "numHeads",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits"
+          },
+          "bindings": ["partial_out_read", "partial_stats_read", "output"],
+          "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        }
+      ]
+    },
+    {
+      "id": "cross_nosg_direct",
+      "priority": 21,
+      "when": ["crossPlainOk", "numSplits == 1"],
+      "requires": {},
+      "passes": [
+        {
+          "id": "split_attention",
+          "name": "DecoderMaskedMultiHeadAttention.CrossSplitKNoSg",
+          "shader": "attn-flash-decode-splitk.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "qNumHeads": "numHeads",
+            "kvNumHeads": "numHeads",
+            "headDim": "headSize",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits",
+            "usesF16": "tensorDtypes.queryT == \"float16\"",
+            "hasKeyPadMask": "present.maskIndexT",
+            "emitQk": "present.qkT",
+            "useSubgroups": false,
+            "allowCooperativeQk": "not useSubgroups or usesF16 or headDimV4 < 8 or batchSize * hiddenSize / 4 > tunables.THREAD_QK_OUTPUT_WORKGROUPS * tunables.WORKGROUP_SIZE"
+          },
+          "bindings": ["query", "key_cross", "value_cross", "params_cross", "output"],
+          "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        }
+      ],
+      "derive": { "directOutput": true }
+    },
+    {
+      "id": "cross_nosg",
+      "priority": 20,
+      "when": ["crossPlainOk"],
+      "requires": {},
+      "intermediates": [
+        { "id": "partialOut", "dtype": "float32", "shape": "[batchSize * numHeads * numSplits * headSize]" },
+        { "id": "partialStats", "dtype": "float32", "shape": "[2 * batchSize * numHeads * numSplits]" }
+      ],
+      "passes": [
+        {
+          "id": "split_attention",
+          "name": "DecoderMaskedMultiHeadAttention.CrossSplitKNoSg",
+          "shader": "attn-flash-decode-splitk.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "qNumHeads": "numHeads",
+            "kvNumHeads": "numHeads",
+            "headDim": "headSize",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits",
+            "usesF16": "tensorDtypes.queryT == \"float16\"",
+            "hasKeyPadMask": "present.maskIndexT",
+            "emitQk": "present.qkT",
+            "useSubgroups": false,
+            "allowCooperativeQk": "not useSubgroups or usesF16 or headDimV4 < 8 or batchSize * hiddenSize / 4 > tunables.THREAD_QK_OUTPUT_WORKGROUPS * tunables.WORKGROUP_SIZE"
+          },
+          "bindings": ["query", "key_cross", "value_cross", "partial_out", "partial_stats", "params_cross"],
+          "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        },
+        {
+          "id": "merge",
+          "name": "DecoderMaskedMultiHeadAttention.CrossMergeNoSg",
+          "shader": "attn-flash-decode-splitk-merge.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "scalar": "attnScalar",
+            "qNumHeads": "numHeads",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits"
+          },
+          "bindings": ["partial_out_read", "partial_stats_read", "output"],
+          "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        }
+      ]
+    },
+    {
+      "id": "cross_qk_nosg_direct",
+      "priority": 21,
+      "when": ["crossQkOk", "numSplits == 1"],
+      "requires": {},
+      "passes": [
+        {
+          "id": "split_attention",
+          "name": "DecoderMaskedMultiHeadAttention.CrossSplitKQkNoSg",
+          "shader": "attn-flash-decode-splitk.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "qNumHeads": "numHeads",
+            "kvNumHeads": "numHeads",
+            "headDim": "headSize",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits",
+            "usesF16": "tensorDtypes.queryT == \"float16\"",
+            "hasKeyPadMask": "present.maskIndexT",
+            "emitQk": "present.qkT",
+            "qkScalar": "qkScalarType",
+            "useSubgroups": false,
+            "allowCooperativeQk": "not useSubgroups or usesF16 or headDimV4 < 8 or batchSize * hiddenSize / 4 > tunables.THREAD_QK_OUTPUT_WORKGROUPS * tunables.WORKGROUP_SIZE"
+          },
+          "bindings": ["query", "key_cross", "value_cross", "qk", "params_cross", "output"],
+          "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        }
+      ],
+      "derive": { "directOutput": true }
+    },
+    {
+      "id": "cross_qk_nosg",
+      "priority": 20,
+      "when": ["crossQkOk"],
+      "requires": {},
+      "intermediates": [
+        { "id": "partialOut", "dtype": "float32", "shape": "[batchSize * numHeads * numSplits * headSize]" },
+        { "id": "partialStats", "dtype": "float32", "shape": "[2 * batchSize * numHeads * numSplits]" }
+      ],
+      "passes": [
+        {
+          "id": "split_attention",
+          "name": "DecoderMaskedMultiHeadAttention.CrossSplitKQkNoSg",
+          "shader": "attn-flash-decode-splitk.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "qNumHeads": "numHeads",
+            "kvNumHeads": "numHeads",
+            "headDim": "headSize",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits",
+            "usesF16": "tensorDtypes.queryT == \"float16\"",
+            "hasKeyPadMask": "present.maskIndexT",
+            "emitQk": "present.qkT",
+            "qkScalar": "qkScalarType",
+            "useSubgroups": false,
+            "allowCooperativeQk": "not useSubgroups or usesF16 or headDimV4 < 8 or batchSize * hiddenSize / 4 > tunables.THREAD_QK_OUTPUT_WORKGROUPS * tunables.WORKGROUP_SIZE"
+          },
+          "bindings": ["query", "key_cross", "value_cross", "qk", "partial_out", "partial_stats", "params_cross"],
+          "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        },
+        {
+          "id": "merge",
+          "name": "DecoderMaskedMultiHeadAttention.CrossMergeNoSg",
+          "shader": "attn-flash-decode-splitk-merge.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "scalar": "attnScalar",
+            "qNumHeads": "numHeads",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits"
+          },
+          "bindings": ["partial_out_read", "partial_stats_read", "output"],
+          "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        }
+      ]
+    },
+    {
+      "id": "cross_mask_nosg_direct",
+      "priority": 21,
+      "when": ["crossMaskOk", "numSplits == 1"],
+      "requires": {},
+      "passes": [
+        {
+          "id": "split_attention",
+          "name": "DecoderMaskedMultiHeadAttention.CrossSplitKNoSg",
+          "shader": "attn-flash-decode-splitk.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "qNumHeads": "numHeads",
+            "kvNumHeads": "numHeads",
+            "headDim": "headSize",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits",
+            "usesF16": "tensorDtypes.queryT == \"float16\"",
+            "hasKeyPadMask": "present.maskIndexT",
+            "maskFilterValue": "attrs.mask_filter_value",
+            "emitQk": "present.qkT",
+            "useSubgroups": false,
+            "allowCooperativeQk": "not useSubgroups or usesF16 or headDimV4 < 8 or batchSize * hiddenSize / 4 > tunables.THREAD_QK_OUTPUT_WORKGROUPS * tunables.WORKGROUP_SIZE"
+          },
+          "bindings": ["query", "key_cross", "value_cross", "key_pad_mask", "params_cross", "output"],
+          "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        }
+      ],
+      "derive": { "directOutput": true }
+    },
+    {
+      "id": "cross_mask_nosg",
+      "priority": 20,
+      "when": ["crossMaskOk"],
+      "requires": {},
+      "intermediates": [
+        { "id": "partialOut", "dtype": "float32", "shape": "[batchSize * numHeads * numSplits * headSize]" },
+        { "id": "partialStats", "dtype": "float32", "shape": "[2 * batchSize * numHeads * numSplits]" }
+      ],
+      "passes": [
+        {
+          "id": "split_attention",
+          "name": "DecoderMaskedMultiHeadAttention.CrossSplitKNoSg",
+          "shader": "attn-flash-decode-splitk.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "qNumHeads": "numHeads",
+            "kvNumHeads": "numHeads",
+            "headDim": "headSize",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits",
+            "usesF16": "tensorDtypes.queryT == \"float16\"",
+            "hasKeyPadMask": "present.maskIndexT",
+            "maskFilterValue": "attrs.mask_filter_value",
+            "emitQk": "present.qkT",
+            "useSubgroups": false,
+            "allowCooperativeQk": "not useSubgroups or usesF16 or headDimV4 < 8 or batchSize * hiddenSize / 4 > tunables.THREAD_QK_OUTPUT_WORKGROUPS * tunables.WORKGROUP_SIZE"
+          },
+          "bindings": ["query", "key_cross", "value_cross", "key_pad_mask", "partial_out", "partial_stats", "params_cross"],
+          "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        },
+        {
+          "id": "merge",
+          "name": "DecoderMaskedMultiHeadAttention.CrossMergeNoSg",
+          "shader": "attn-flash-decode-splitk-merge.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "scalar": "attnScalar",
+            "qNumHeads": "numHeads",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits"
+          },
+          "bindings": ["partial_out_read", "partial_stats_read", "output"],
+          "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        }
+      ]
+    },
+    {
+      "id": "cross_mask_qk_nosg_direct",
+      "priority": 21,
+      "when": ["crossMaskQkOk", "numSplits == 1"],
+      "requires": {},
+      "passes": [
+        {
+          "id": "split_attention",
+          "name": "DecoderMaskedMultiHeadAttention.CrossSplitKQkNoSg",
+          "shader": "attn-flash-decode-splitk.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "qNumHeads": "numHeads",
+            "kvNumHeads": "numHeads",
+            "headDim": "headSize",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits",
+            "usesF16": "tensorDtypes.queryT == \"float16\"",
+            "hasKeyPadMask": "present.maskIndexT",
+            "maskFilterValue": "attrs.mask_filter_value",
+            "emitQk": "present.qkT",
+            "qkScalar": "qkScalarType",
+            "useSubgroups": false,
+            "allowCooperativeQk": "not useSubgroups or usesF16 or headDimV4 < 8 or batchSize * hiddenSize / 4 > tunables.THREAD_QK_OUTPUT_WORKGROUPS * tunables.WORKGROUP_SIZE"
+          },
+          "bindings": ["query", "key_cross", "value_cross", "key_pad_mask", "qk", "params_cross", "output"],
+          "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        }
+      ],
+      "derive": { "directOutput": true }
+    },
+    {
+      "id": "cross_mask_qk_nosg",
+      "priority": 20,
+      "when": ["crossMaskQkOk"],
+      "requires": {},
+      "intermediates": [
+        { "id": "partialOut", "dtype": "float32", "shape": "[batchSize * numHeads * numSplits * headSize]" },
+        { "id": "partialStats", "dtype": "float32", "shape": "[2 * batchSize * numHeads * numSplits]" }
+      ],
+      "passes": [
+        {
+          "id": "split_attention",
+          "name": "DecoderMaskedMultiHeadAttention.CrossSplitKQkNoSg",
+          "shader": "attn-flash-decode-splitk.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "qNumHeads": "numHeads",
+            "kvNumHeads": "numHeads",
+            "headDim": "headSize",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits",
+            "usesF16": "tensorDtypes.queryT == \"float16\"",
+            "hasKeyPadMask": "present.maskIndexT",
+            "maskFilterValue": "attrs.mask_filter_value",
+            "emitQk": "present.qkT",
+            "qkScalar": "qkScalarType",
+            "useSubgroups": false,
+            "allowCooperativeQk": "not useSubgroups or usesF16 or headDimV4 < 8 or batchSize * hiddenSize / 4 > tunables.THREAD_QK_OUTPUT_WORKGROUPS * tunables.WORKGROUP_SIZE"
+          },
+          "bindings": ["query", "key_cross", "value_cross", "key_pad_mask", "qk", "partial_out", "partial_stats", "params_cross"],
+          "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        },
+        {
+          "id": "merge",
+          "name": "DecoderMaskedMultiHeadAttention.CrossMergeNoSg",
+          "shader": "attn-flash-decode-splitk-merge.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "scalar": "attnScalar",
+            "qNumHeads": "numHeads",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits"
+          },
+          "bindings": ["partial_out_read", "partial_stats_read", "output"],
+          "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        }
+      ]
+    },
+    {
+      "id": "self",
+      "priority": 30,
+      "when": ["selfPlainOk", "subgroupSplitKOk"],
+      "demoteWhen": ["not subgroupSplitKPreferred"],
+      "requires": { "features": ["subgroups"] },
+      "intermediates": [
+        { "id": "partialOut", "dtype": "float32", "shape": "[batchSize * numHeads * numSplits * headSize]" },
+        { "id": "partialStats", "dtype": "float32", "shape": "[2 * batchSize * numHeads * numSplits]" }
+      ],
+      "passes": [
+        {
+          "id": "present_retain",
+          "name": "DecoderMaskedMultiHeadAttention.PresentRetain",
+          "shader": "dmmha-present.wgsl.jinja",
+          "derive": {
+            "appendRow": false,
+            "workgroupSize": "tunables.COPY_WORKGROUP_SIZE",
+            "itemsPerLane": "retainItemsPerLane",
+            "cacheScalar": "attnScalar"
+          },
+          "bindings": ["past_key", "past_value", "present_key", "present_value", "params_retain"],
+          "dispatch": {
+            "x": "min(ceilDiv(cacheWordsV4, tunables.COPY_WORKGROUP_SIZE * retainItemsPerLane), 65535)",
+            "y": "ceilDiv(ceilDiv(cacheWordsV4, tunables.COPY_WORKGROUP_SIZE * retainItemsPerLane), 65535)",
+            "z": 1
+          },
+          "viewAlias": [
+            { "input": "past_key", "output": "present_key" },
+            { "input": "past_value", "output": "present_value" }
+          ]
+        },
+        {
+          "id": "present_append",
+          "name": "DecoderMaskedMultiHeadAttention.PresentAppend",
+          "shader": "dmmha-present.wgsl.jinja",
+          "derive": {
+            "appendRow": true,
+            "workgroupSize": "tunables.COPY_WORKGROUP_SIZE",
+            "headSizeV4": "headDimV4",
+            "numHeads": "numHeads"
+          },
+          "bindings": ["key_row", "value_row", "past_sequence_length", "present_key", "present_value", "params_append"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize * hiddenSize / 4), (tunables.COPY_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize * hiddenSize / 4), (tunables.COPY_WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
+        },
+        {
+          "id": "split_attention",
+          "name": "DecoderMaskedMultiHeadAttention.SelfSplitK",
+          "shader": "attn-flash-decode-splitk.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "qNumHeads": "numHeads",
+            "kvNumHeads": "numHeads",
+            "headDim": "headSize",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits",
+            "usesF16": "tensorDtypes.queryT == \"float16\"",
+            "cacheSeqlens": true,
+            "cacheSeqlensIndex": "\"0u\"",
+            "hasKeyPadMask": "present.maskIndexT",
+            "hasMask": "present.attentionBiasT"
+          },
+          "bindings": ["query", "key_cache", "value_cache", "seqlens_k", "partial_out", "partial_stats", "params_self"],
+          "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        },
+        {
+          "id": "merge",
+          "name": "DecoderMaskedMultiHeadAttention.SelfMerge",
+          "shader": "attn-flash-decode-splitk-merge.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "scalar": "attnScalar",
+            "qNumHeads": "numHeads",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits"
+          },
+          "bindings": ["partial_out_read", "partial_stats_read", "output"],
+          "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        }
+      ]
+    },
+    {
+      "id": "self_bias",
+      "priority": 30,
+      "when": ["selfBiasOk", "subgroupSplitKOk"],
+      "demoteWhen": ["not subgroupSplitKPreferred"],
+      "requires": { "features": ["subgroups"] },
+      "intermediates": [
+        { "id": "partialOut", "dtype": "float32", "shape": "[batchSize * numHeads * numSplits * headSize]" },
+        { "id": "partialStats", "dtype": "float32", "shape": "[2 * batchSize * numHeads * numSplits]" }
+      ],
+      "passes": [
+        {
+          "id": "present_retain",
+          "name": "DecoderMaskedMultiHeadAttention.PresentRetain",
+          "shader": "dmmha-present.wgsl.jinja",
+          "derive": {
+            "appendRow": false,
+            "workgroupSize": "tunables.COPY_WORKGROUP_SIZE",
+            "itemsPerLane": "retainItemsPerLane",
+            "cacheScalar": "attnScalar"
+          },
+          "bindings": ["past_key", "past_value", "present_key", "present_value", "params_retain"],
+          "dispatch": {
+            "x": "min(ceilDiv(cacheWordsV4, tunables.COPY_WORKGROUP_SIZE * retainItemsPerLane), 65535)",
+            "y": "ceilDiv(ceilDiv(cacheWordsV4, tunables.COPY_WORKGROUP_SIZE * retainItemsPerLane), 65535)",
+            "z": 1
+          },
+          "viewAlias": [
+            { "input": "past_key", "output": "present_key" },
+            { "input": "past_value", "output": "present_value" }
+          ]
+        },
+        {
+          "id": "present_append",
+          "name": "DecoderMaskedMultiHeadAttention.PresentAppend",
+          "shader": "dmmha-present.wgsl.jinja",
+          "derive": {
+            "appendRow": true,
+            "workgroupSize": "tunables.COPY_WORKGROUP_SIZE",
+            "headSizeV4": "headDimV4",
+            "numHeads": "numHeads"
+          },
+          "bindings": ["key_row", "value_row", "past_sequence_length", "present_key", "present_value", "params_append"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize * hiddenSize / 4), (tunables.COPY_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize * hiddenSize / 4), (tunables.COPY_WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
+        },
+        {
+          "id": "split_attention",
+          "name": "DecoderMaskedMultiHeadAttention.SelfSplitK",
+          "shader": "attn-flash-decode-splitk.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "qNumHeads": "numHeads",
+            "kvNumHeads": "numHeads",
+            "headDim": "headSize",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits",
+            "usesF16": "tensorDtypes.queryT == \"float16\"",
+            "cacheSeqlens": true,
+            "cacheSeqlensIndex": "\"0u\"",
+            "hasKeyPadMask": "present.maskIndexT",
+            "hasMask": "present.attentionBiasT"
+          },
+          "bindings": ["query", "key_cache", "value_cache", "seqlens_k", "attn_mask", "partial_out", "partial_stats", "params_self_bias"],
+          "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        },
+        {
+          "id": "merge",
+          "name": "DecoderMaskedMultiHeadAttention.SelfMerge",
+          "shader": "attn-flash-decode-splitk-merge.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "scalar": "attnScalar",
+            "qNumHeads": "numHeads",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits"
+          },
+          "bindings": ["partial_out_read", "partial_stats_read", "output"],
+          "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        }
+      ]
+    },
+    {
+      "id": "self_mask",
+      "priority": 30,
+      "when": ["selfMaskOk", "subgroupSplitKOk"],
+      "demoteWhen": ["not subgroupSplitKPreferred"],
+      "requires": { "features": ["subgroups"] },
+      "intermediates": [
+        { "id": "partialOut", "dtype": "float32", "shape": "[batchSize * numHeads * numSplits * headSize]" },
+        { "id": "partialStats", "dtype": "float32", "shape": "[2 * batchSize * numHeads * numSplits]" }
+      ],
+      "passes": [
+        {
+          "id": "present_retain",
+          "name": "DecoderMaskedMultiHeadAttention.PresentRetain",
+          "shader": "dmmha-present.wgsl.jinja",
+          "derive": {
+            "appendRow": false,
+            "workgroupSize": "tunables.COPY_WORKGROUP_SIZE",
+            "itemsPerLane": "retainItemsPerLane",
+            "cacheScalar": "attnScalar"
+          },
+          "bindings": ["past_key", "past_value", "present_key", "present_value", "params_retain"],
+          "dispatch": {
+            "x": "min(ceilDiv(cacheWordsV4, tunables.COPY_WORKGROUP_SIZE * retainItemsPerLane), 65535)",
+            "y": "ceilDiv(ceilDiv(cacheWordsV4, tunables.COPY_WORKGROUP_SIZE * retainItemsPerLane), 65535)",
+            "z": 1
+          },
+          "viewAlias": [
+            { "input": "past_key", "output": "present_key" },
+            { "input": "past_value", "output": "present_value" }
+          ]
+        },
+        {
+          "id": "present_append",
+          "name": "DecoderMaskedMultiHeadAttention.PresentAppend",
+          "shader": "dmmha-present.wgsl.jinja",
+          "derive": {
+            "appendRow": true,
+            "workgroupSize": "tunables.COPY_WORKGROUP_SIZE",
+            "headSizeV4": "headDimV4",
+            "numHeads": "numHeads"
+          },
+          "bindings": ["key_row", "value_row", "past_sequence_length", "present_key", "present_value", "params_append"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize * hiddenSize / 4), (tunables.COPY_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize * hiddenSize / 4), (tunables.COPY_WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
+        },
+        {
+          "id": "split_attention",
+          "name": "DecoderMaskedMultiHeadAttention.SelfSplitK",
+          "shader": "attn-flash-decode-splitk.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "qNumHeads": "numHeads",
+            "kvNumHeads": "numHeads",
+            "headDim": "headSize",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits",
+            "usesF16": "tensorDtypes.queryT == \"float16\"",
+            "cacheSeqlens": true,
+            "cacheSeqlensIndex": "\"0u\"",
+            "hasKeyPadMask": "present.maskIndexT",
+            "maskFilterValue": "attrs.mask_filter_value",
+            "hasMask": "present.attentionBiasT"
+          },
+          "bindings": ["query", "key_cache", "value_cache", "seqlens_k", "key_pad_mask", "partial_out", "partial_stats", "params_self"],
+          "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        },
+        {
+          "id": "merge",
+          "name": "DecoderMaskedMultiHeadAttention.SelfMerge",
+          "shader": "attn-flash-decode-splitk-merge.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "scalar": "attnScalar",
+            "qNumHeads": "numHeads",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits"
+          },
+          "bindings": ["partial_out_read", "partial_stats_read", "output"],
+          "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        }
+      ]
+    },
+    {
+      "id": "self_mask_bias",
+      "priority": 30,
+      "when": ["selfMaskBiasOk", "subgroupSplitKOk"],
+      "demoteWhen": ["not subgroupSplitKPreferred"],
+      "requires": { "features": ["subgroups"] },
+      "intermediates": [
+        { "id": "partialOut", "dtype": "float32", "shape": "[batchSize * numHeads * numSplits * headSize]" },
+        { "id": "partialStats", "dtype": "float32", "shape": "[2 * batchSize * numHeads * numSplits]" }
+      ],
+      "passes": [
+        {
+          "id": "present_retain",
+          "name": "DecoderMaskedMultiHeadAttention.PresentRetain",
+          "shader": "dmmha-present.wgsl.jinja",
+          "derive": {
+            "appendRow": false,
+            "workgroupSize": "tunables.COPY_WORKGROUP_SIZE",
+            "itemsPerLane": "retainItemsPerLane",
+            "cacheScalar": "attnScalar"
+          },
+          "bindings": ["past_key", "past_value", "present_key", "present_value", "params_retain"],
+          "dispatch": {
+            "x": "min(ceilDiv(cacheWordsV4, tunables.COPY_WORKGROUP_SIZE * retainItemsPerLane), 65535)",
+            "y": "ceilDiv(ceilDiv(cacheWordsV4, tunables.COPY_WORKGROUP_SIZE * retainItemsPerLane), 65535)",
+            "z": 1
+          },
+          "viewAlias": [
+            { "input": "past_key", "output": "present_key" },
+            { "input": "past_value", "output": "present_value" }
+          ]
+        },
+        {
+          "id": "present_append",
+          "name": "DecoderMaskedMultiHeadAttention.PresentAppend",
+          "shader": "dmmha-present.wgsl.jinja",
+          "derive": {
+            "appendRow": true,
+            "workgroupSize": "tunables.COPY_WORKGROUP_SIZE",
+            "headSizeV4": "headDimV4",
+            "numHeads": "numHeads"
+          },
+          "bindings": ["key_row", "value_row", "past_sequence_length", "present_key", "present_value", "params_append"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize * hiddenSize / 4), (tunables.COPY_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize * hiddenSize / 4), (tunables.COPY_WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
+        },
+        {
+          "id": "split_attention",
+          "name": "DecoderMaskedMultiHeadAttention.SelfSplitK",
+          "shader": "attn-flash-decode-splitk.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "qNumHeads": "numHeads",
+            "kvNumHeads": "numHeads",
+            "headDim": "headSize",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits",
+            "usesF16": "tensorDtypes.queryT == \"float16\"",
+            "cacheSeqlens": true,
+            "cacheSeqlensIndex": "\"0u\"",
+            "hasKeyPadMask": "present.maskIndexT",
+            "maskFilterValue": "attrs.mask_filter_value",
+            "hasMask": "present.attentionBiasT"
+          },
+          "bindings": ["query", "key_cache", "value_cache", "seqlens_k", "key_pad_mask", "attn_mask", "partial_out", "partial_stats", "params_self_bias"],
+          "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        },
+        {
+          "id": "merge",
+          "name": "DecoderMaskedMultiHeadAttention.SelfMerge",
+          "shader": "attn-flash-decode-splitk-merge.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "scalar": "attnScalar",
+            "qNumHeads": "numHeads",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits"
+          },
+          "bindings": ["partial_out_read", "partial_stats_read", "output"],
+          "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        }
+      ]
+    },
+    {
+      "id": "self_nosg",
+      "priority": 20,
+      "when": ["selfPlainOk"],
+      "requires": {},
+      "intermediates": [
+        { "id": "partialOut", "dtype": "float32", "shape": "[batchSize * numHeads * numSplits * headSize]" },
+        { "id": "partialStats", "dtype": "float32", "shape": "[2 * batchSize * numHeads * numSplits]" }
+      ],
+      "passes": [
+        {
+          "id": "present_retain",
+          "name": "DecoderMaskedMultiHeadAttention.PresentRetain",
+          "shader": "dmmha-present.wgsl.jinja",
+          "derive": {
+            "appendRow": false,
+            "workgroupSize": "tunables.COPY_WORKGROUP_SIZE",
+            "itemsPerLane": "retainItemsPerLane",
+            "cacheScalar": "attnScalar"
+          },
+          "bindings": ["past_key", "past_value", "present_key", "present_value", "params_retain"],
+          "dispatch": {
+            "x": "min(ceilDiv(cacheWordsV4, tunables.COPY_WORKGROUP_SIZE * retainItemsPerLane), 65535)",
+            "y": "ceilDiv(ceilDiv(cacheWordsV4, tunables.COPY_WORKGROUP_SIZE * retainItemsPerLane), 65535)",
+            "z": 1
+          },
+          "viewAlias": [
+            { "input": "past_key", "output": "present_key" },
+            { "input": "past_value", "output": "present_value" }
+          ]
+        },
+        {
+          "id": "present_append",
+          "name": "DecoderMaskedMultiHeadAttention.PresentAppend",
+          "shader": "dmmha-present.wgsl.jinja",
+          "derive": {
+            "appendRow": true,
+            "workgroupSize": "tunables.COPY_WORKGROUP_SIZE",
+            "headSizeV4": "headDimV4",
+            "numHeads": "numHeads"
+          },
+          "bindings": ["key_row", "value_row", "past_sequence_length", "present_key", "present_value", "params_append"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize * hiddenSize / 4), (tunables.COPY_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize * hiddenSize / 4), (tunables.COPY_WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
+        },
+        {
+          "id": "split_attention",
+          "name": "DecoderMaskedMultiHeadAttention.SelfSplitKNoSg",
+          "shader": "attn-flash-decode-splitk.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "qNumHeads": "numHeads",
+            "kvNumHeads": "numHeads",
+            "headDim": "headSize",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits",
+            "usesF16": "tensorDtypes.queryT == \"float16\"",
+            "cacheSeqlens": true,
+            "cacheSeqlensIndex": "\"0u\"",
+            "hasKeyPadMask": "present.maskIndexT",
+            "hasMask": "present.attentionBiasT",
+            "useSubgroups": false
+          },
+          "bindings": ["query", "key_cache", "value_cache", "seqlens_k", "partial_out", "partial_stats", "params_self"],
+          "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        },
+        {
+          "id": "merge",
+          "name": "DecoderMaskedMultiHeadAttention.SelfMergeNoSg",
+          "shader": "attn-flash-decode-splitk-merge.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "scalar": "attnScalar",
+            "qNumHeads": "numHeads",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits"
+          },
+          "bindings": ["partial_out_read", "partial_stats_read", "output"],
+          "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        }
+      ]
+    },
+    {
+      "id": "self_bias_nosg",
+      "priority": 20,
+      "when": ["selfBiasOk"],
+      "requires": {},
+      "intermediates": [
+        { "id": "partialOut", "dtype": "float32", "shape": "[batchSize * numHeads * numSplits * headSize]" },
+        { "id": "partialStats", "dtype": "float32", "shape": "[2 * batchSize * numHeads * numSplits]" }
+      ],
+      "passes": [
+        {
+          "id": "present_retain",
+          "name": "DecoderMaskedMultiHeadAttention.PresentRetain",
+          "shader": "dmmha-present.wgsl.jinja",
+          "derive": {
+            "appendRow": false,
+            "workgroupSize": "tunables.COPY_WORKGROUP_SIZE",
+            "itemsPerLane": "retainItemsPerLane",
+            "cacheScalar": "attnScalar"
+          },
+          "bindings": ["past_key", "past_value", "present_key", "present_value", "params_retain"],
+          "dispatch": {
+            "x": "min(ceilDiv(cacheWordsV4, tunables.COPY_WORKGROUP_SIZE * retainItemsPerLane), 65535)",
+            "y": "ceilDiv(ceilDiv(cacheWordsV4, tunables.COPY_WORKGROUP_SIZE * retainItemsPerLane), 65535)",
+            "z": 1
+          },
+          "viewAlias": [
+            { "input": "past_key", "output": "present_key" },
+            { "input": "past_value", "output": "present_value" }
+          ]
+        },
+        {
+          "id": "present_append",
+          "name": "DecoderMaskedMultiHeadAttention.PresentAppend",
+          "shader": "dmmha-present.wgsl.jinja",
+          "derive": {
+            "appendRow": true,
+            "workgroupSize": "tunables.COPY_WORKGROUP_SIZE",
+            "headSizeV4": "headDimV4",
+            "numHeads": "numHeads"
+          },
+          "bindings": ["key_row", "value_row", "past_sequence_length", "present_key", "present_value", "params_append"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize * hiddenSize / 4), (tunables.COPY_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize * hiddenSize / 4), (tunables.COPY_WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
+        },
+        {
+          "id": "split_attention",
+          "name": "DecoderMaskedMultiHeadAttention.SelfSplitKNoSg",
+          "shader": "attn-flash-decode-splitk.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "qNumHeads": "numHeads",
+            "kvNumHeads": "numHeads",
+            "headDim": "headSize",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits",
+            "usesF16": "tensorDtypes.queryT == \"float16\"",
+            "cacheSeqlens": true,
+            "cacheSeqlensIndex": "\"0u\"",
+            "hasKeyPadMask": "present.maskIndexT",
+            "hasMask": "present.attentionBiasT",
+            "useSubgroups": false
+          },
+          "bindings": ["query", "key_cache", "value_cache", "seqlens_k", "attn_mask", "partial_out", "partial_stats", "params_self_bias"],
+          "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        },
+        {
+          "id": "merge",
+          "name": "DecoderMaskedMultiHeadAttention.SelfMergeNoSg",
+          "shader": "attn-flash-decode-splitk-merge.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "scalar": "attnScalar",
+            "qNumHeads": "numHeads",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits"
+          },
+          "bindings": ["partial_out_read", "partial_stats_read", "output"],
+          "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        }
+      ]
+    },
+    {
+      "id": "self_mask_nosg",
+      "priority": 20,
+      "when": ["selfMaskOk"],
+      "requires": {},
+      "intermediates": [
+        { "id": "partialOut", "dtype": "float32", "shape": "[batchSize * numHeads * numSplits * headSize]" },
+        { "id": "partialStats", "dtype": "float32", "shape": "[2 * batchSize * numHeads * numSplits]" }
+      ],
+      "passes": [
+        {
+          "id": "present_retain",
+          "name": "DecoderMaskedMultiHeadAttention.PresentRetain",
+          "shader": "dmmha-present.wgsl.jinja",
+          "derive": {
+            "appendRow": false,
+            "workgroupSize": "tunables.COPY_WORKGROUP_SIZE",
+            "itemsPerLane": "retainItemsPerLane",
+            "cacheScalar": "attnScalar"
+          },
+          "bindings": ["past_key", "past_value", "present_key", "present_value", "params_retain"],
+          "dispatch": {
+            "x": "min(ceilDiv(cacheWordsV4, tunables.COPY_WORKGROUP_SIZE * retainItemsPerLane), 65535)",
+            "y": "ceilDiv(ceilDiv(cacheWordsV4, tunables.COPY_WORKGROUP_SIZE * retainItemsPerLane), 65535)",
+            "z": 1
+          },
+          "viewAlias": [
+            { "input": "past_key", "output": "present_key" },
+            { "input": "past_value", "output": "present_value" }
+          ]
+        },
+        {
+          "id": "present_append",
+          "name": "DecoderMaskedMultiHeadAttention.PresentAppend",
+          "shader": "dmmha-present.wgsl.jinja",
+          "derive": {
+            "appendRow": true,
+            "workgroupSize": "tunables.COPY_WORKGROUP_SIZE",
+            "headSizeV4": "headDimV4",
+            "numHeads": "numHeads"
+          },
+          "bindings": ["key_row", "value_row", "past_sequence_length", "present_key", "present_value", "params_append"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize * hiddenSize / 4), (tunables.COPY_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize * hiddenSize / 4), (tunables.COPY_WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
+        },
+        {
+          "id": "split_attention",
+          "name": "DecoderMaskedMultiHeadAttention.SelfSplitKNoSg",
+          "shader": "attn-flash-decode-splitk.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "qNumHeads": "numHeads",
+            "kvNumHeads": "numHeads",
+            "headDim": "headSize",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits",
+            "usesF16": "tensorDtypes.queryT == \"float16\"",
+            "cacheSeqlens": true,
+            "cacheSeqlensIndex": "\"0u\"",
+            "hasKeyPadMask": "present.maskIndexT",
+            "maskFilterValue": "attrs.mask_filter_value",
+            "hasMask": "present.attentionBiasT",
+            "useSubgroups": false
+          },
+          "bindings": ["query", "key_cache", "value_cache", "seqlens_k", "key_pad_mask", "partial_out", "partial_stats", "params_self"],
+          "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        },
+        {
+          "id": "merge",
+          "name": "DecoderMaskedMultiHeadAttention.SelfMergeNoSg",
+          "shader": "attn-flash-decode-splitk-merge.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "scalar": "attnScalar",
+            "qNumHeads": "numHeads",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits"
+          },
+          "bindings": ["partial_out_read", "partial_stats_read", "output"],
+          "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        }
+      ]
+    },
+    {
+      "id": "self_mask_bias_nosg",
+      "priority": 20,
+      "when": ["selfMaskBiasOk"],
+      "requires": {},
+      "intermediates": [
+        { "id": "partialOut", "dtype": "float32", "shape": "[batchSize * numHeads * numSplits * headSize]" },
+        { "id": "partialStats", "dtype": "float32", "shape": "[2 * batchSize * numHeads * numSplits]" }
+      ],
+      "passes": [
+        {
+          "id": "present_retain",
+          "name": "DecoderMaskedMultiHeadAttention.PresentRetain",
+          "shader": "dmmha-present.wgsl.jinja",
+          "derive": {
+            "appendRow": false,
+            "workgroupSize": "tunables.COPY_WORKGROUP_SIZE",
+            "itemsPerLane": "retainItemsPerLane",
+            "cacheScalar": "attnScalar"
+          },
+          "bindings": ["past_key", "past_value", "present_key", "present_value", "params_retain"],
+          "dispatch": {
+            "x": "min(ceilDiv(cacheWordsV4, tunables.COPY_WORKGROUP_SIZE * retainItemsPerLane), 65535)",
+            "y": "ceilDiv(ceilDiv(cacheWordsV4, tunables.COPY_WORKGROUP_SIZE * retainItemsPerLane), 65535)",
+            "z": 1
+          },
+          "viewAlias": [
+            { "input": "past_key", "output": "present_key" },
+            { "input": "past_value", "output": "present_value" }
+          ]
+        },
+        {
+          "id": "present_append",
+          "name": "DecoderMaskedMultiHeadAttention.PresentAppend",
+          "shader": "dmmha-present.wgsl.jinja",
+          "derive": {
+            "appendRow": true,
+            "workgroupSize": "tunables.COPY_WORKGROUP_SIZE",
+            "headSizeV4": "headDimV4",
+            "numHeads": "numHeads"
+          },
+          "bindings": ["key_row", "value_row", "past_sequence_length", "present_key", "present_value", "params_append"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize * hiddenSize / 4), (tunables.COPY_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize * hiddenSize / 4), (tunables.COPY_WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
+        },
+        {
+          "id": "split_attention",
+          "name": "DecoderMaskedMultiHeadAttention.SelfSplitKNoSg",
+          "shader": "attn-flash-decode-splitk.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "qNumHeads": "numHeads",
+            "kvNumHeads": "numHeads",
+            "headDim": "headSize",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits",
+            "usesF16": "tensorDtypes.queryT == \"float16\"",
+            "cacheSeqlens": true,
+            "cacheSeqlensIndex": "\"0u\"",
+            "hasKeyPadMask": "present.maskIndexT",
+            "maskFilterValue": "attrs.mask_filter_value",
+            "hasMask": "present.attentionBiasT",
+            "useSubgroups": false
+          },
+          "bindings": ["query", "key_cache", "value_cache", "seqlens_k", "key_pad_mask", "attn_mask", "partial_out", "partial_stats", "params_self_bias"],
+          "dispatch": { "x": "numSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        },
+        {
+          "id": "merge",
+          "name": "DecoderMaskedMultiHeadAttention.SelfMergeNoSg",
+          "shader": "attn-flash-decode-splitk-merge.wgsl.jinja",
+          "derive": {
+            "layout": "\"bhsd\"",
+            "splitQueries": false,
+            "scalar": "attnScalar",
+            "qNumHeads": "numHeads",
+            "headDimV4": "headDimV4",
+            "numSplits": "numSplits"
+          },
+          "bindings": ["partial_out_read", "partial_stats_read", "output"],
+          "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
+        }
+      ]
+    }
+  ]
+}

build/webgpu/metadata.json ADDED Viewed

	@@ -0,0 +1,48 @@

+{
+  "name": "com.microsoft.DecoderMaskedMultiHeadAttention",
+  "id": "_com_microsoft_decodermaskedmultiheadattention_webgpu_d93bf4a",
+  "version": 1,
+  "license": "Apache-2.0",
+  "backend": { "type": "webgpu" },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "attn-flash-decode-splitk-merge.wgsl.jinja": "7fi0SKKDfYjYG+l9OHNcodIeSFH5zXKMIJi9iv32rZo=",
+      "attn-flash-decode-splitk.wgsl.jinja": "BFL68/1x2lxrN/lTCgC5cy0bmEnkeKQksPtZKwhLE7A=",
+      "bench.json": "Kow9t33DDsCfO1hnX4CBJ3bKQkH284pj5sKvwirbabA=",
+      "dmmha-present.wgsl.jinja": "OCydRlEFGVeylOtabFeuAxrofZ/c+uqe5o1RgyWx9fU=",
+      "manifest.json": "OemoYbwPetpiqBxosVOlV2e20zGYfgzT4zv0y4HdcOc=",
+      "test.json": "2CVKz1jxKg77lGzgcrIBiv1FqPQNSORMuRdC2qRXNyw="
+    }
+  },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.1",
+    "variants": {
+      "cross_direct": ["attn-flash-decode-splitk.wgsl.jinja"],
+      "cross": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja"],
+      "cross_qk_direct": ["attn-flash-decode-splitk.wgsl.jinja"],
+      "cross_qk": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja"],
+      "cross_mask_direct": ["attn-flash-decode-splitk.wgsl.jinja"],
+      "cross_mask": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja"],
+      "cross_mask_qk_direct": ["attn-flash-decode-splitk.wgsl.jinja"],
+      "cross_mask_qk": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja"],
+      "cross_nosg_direct": ["attn-flash-decode-splitk.wgsl.jinja"],
+      "cross_nosg": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja"],
+      "cross_qk_nosg_direct": ["attn-flash-decode-splitk.wgsl.jinja"],
+      "cross_qk_nosg": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja"],
+      "cross_mask_nosg_direct": ["attn-flash-decode-splitk.wgsl.jinja"],
+      "cross_mask_nosg": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja"],
+      "cross_mask_qk_nosg_direct": ["attn-flash-decode-splitk.wgsl.jinja"],
+      "cross_mask_qk_nosg": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja"],
+      "self": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja", "dmmha-present.wgsl.jinja"],
+      "self_bias": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja", "dmmha-present.wgsl.jinja"],
+      "self_mask": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja", "dmmha-present.wgsl.jinja"],
+      "self_mask_bias": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja", "dmmha-present.wgsl.jinja"],
+      "self_nosg": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja", "dmmha-present.wgsl.jinja"],
+      "self_bias_nosg": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja", "dmmha-present.wgsl.jinja"],
+      "self_mask_nosg": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja", "dmmha-present.wgsl.jinja"],
+      "self_mask_bias_nosg": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja", "dmmha-present.wgsl.jinja"]
+    }
+  }
+}

build/webgpu/test.json ADDED Viewed

	@@ -0,0 +1,1580 @@

+{
+  "cases": [
+    {
+      "name": "cross_plain_b1_h4_kv16_d16",
+      "provenance": {
+        "notes": "Cross attention with no mask and no diagnostic row: the minimum contract, where every encoder key contributes."
+      },
+      "attrs": { "num_heads": 4 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 4, 16, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 4, 16, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [1, 1, 64], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "cross_plain_b2_h6_kv40_d32",
+      "provenance": {
+        "notes": "Cross attention over a key count that is not a multiple of the split tile, so the last partition absorbs the remainder."
+      },
+      "attrs": { "num_heads": 6 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 1, 192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [2, 6, 40, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [2, 6, 40, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [2, 1, 192], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "cross_mask_b1_h4_kv16_d16",
+      "provenance": {
+        "notes": "Cross attention with an alternating key-padding mask. A zero entry adds the filter value, so a masked key keeps a tiny but non-zero weight."
+      },
+      "attrs": { "num_heads": 4 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 4, 16, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 4, 16, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        },
+        "maskIndexT": { "dtype": "int32", "shape": [1, 16], "data": { "kind": "cycle", "values": [1, 0] } }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [1, 1, 64], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "cross_mask_all_padded_b1_h2_kv8_d16",
+      "provenance": {
+        "notes": "Every mask entry is zero. The additive filter leaves the row uniform, where a knockout mask would return zeros; this case separates the two readings."
+      },
+      "attrs": { "num_heads": 2 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 2, 8, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 2, 8, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        },
+        "maskIndexT": { "dtype": "int32", "shape": [1, 8], "data": { "kind": "constant", "value": 0 } }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [1, 1, 32], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "cross_mask_filter_value_b1_h4_kv12_d16",
+      "provenance": {
+        "notes": "A non-default mask filter value, small enough that masked keys still carry visible weight."
+      },
+      "attrs": { "num_heads": 4, "mask_filter_value": -4 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 4, 12, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 4, 12, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        },
+        "maskIndexT": { "dtype": "int32", "shape": [1, 12], "data": { "kind": "cycle", "values": [1, 1, 0] } }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [1, 1, 64], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "cross_qk_b1_h4_kv16_d16",
+      "provenance": {
+        "notes": "Cross attention emitting the pre-softmax score row. The values are the scaled dot products, so a post-softmax reading would show every row summing to one."
+      },
+      "attrs": { "num_heads": 4, "output_qk": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 4, 16, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 4, 16, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 1, 64], "tolerance": 0.00002 },
+        "qkT": { "dtype": "float32", "shape": [1, 4, 1, 16], "tolerance": 0.00002 }
+      }
+    },
+    {
+      "name": "cross_mask_qk_b1_h4_kv20_d16",
+      "provenance": {
+        "notes": "The score row is emitted after the mask is added, so the masked columns sit a filter value below their neighbours."
+      },
+      "attrs": { "num_heads": 4, "output_qk": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 4, 20, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 4, 20, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        },
+        "maskIndexT": { "dtype": "int32", "shape": [1, 20], "data": { "kind": "cycle", "values": [1, 0, 1] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 1, 64], "tolerance": 0.00002 },
+        "qkT": { "dtype": "float32", "shape": [1, 4, 1, 20], "tolerance": 0.00002 }
+      }
+    },
+    {
+      "name": "cross_qk_scale_b2_h3_kv24_d32",
+      "provenance": {
+        "notes": "An explicit scale overrides the default reciprocal square root, and the emitted score row shows it directly."
+      },
+      "attrs": { "num_heads": 3, "scale": 0.125, "output_qk": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 1, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [2, 3, 24, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [2, 3, 24, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [2, 1, 96], "tolerance": 0.00002 },
+        "qkT": { "dtype": "float32", "shape": [2, 3, 1, 24], "tolerance": 0.00002 }
+      }
+    },
+    {
+      "name": "cross_plain_b1_h2_kv1_d16",
+      "provenance": {
+        "notes": "A single encoder key: the softmax is degenerate and the output is that key's value row."
+      },
+      "attrs": { "num_heads": 2 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 2, 1, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 2, 1, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [1, 1, 32], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "cross_plain_b1_h1_kv1500_d64",
+      "provenance": {
+        "notes": "Whisper-scale cross attention geometry at one head, where the key count is far larger than the head dimension."
+      },
+      "attrs": { "num_heads": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 1500, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 1500, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [1, 1, 64], "tolerance": 0.00005 } }
+    },
+    {
+      "name": "cross_head_size_four_b1_h4_kv8_d4",
+      "provenance": {
+        "notes": "Cross-attention with 4 heads and head width 4, the minimum that fills one four-element vector, over an 8-token key/value cache."
+      },
+      "attrs": { "num_heads": 4 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 4, 8, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 4, 8, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [1, 1, 16], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "cross_plain_f16_b1_h4_kv16_d16",
+      "provenance": {
+        "notes": "The float16 twin of the plain cross route: operands and output narrow, while the score and online softmax stay float32."
+      },
+      "attrs": { "num_heads": 4 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [1, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [1, 4, 16, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [1, 4, 16, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float16", "shape": [1, 1, 64], "tolerance": 0.003 } }
+    },
+    {
+      "name": "cross_mask_qk_f16_b1_h4_kv20_d16",
+      "provenance": {
+        "notes": "Float16 operands with a float32 score row, the combination the schema allows by typing the diagnostic output independently."
+      },
+      "attrs": { "num_heads": 4, "output_qk": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [1, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [1, 4, 20, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [1, 4, 20, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        },
+        "maskIndexT": { "dtype": "int32", "shape": [1, 20], "data": { "kind": "cycle", "values": [1, 0, 1] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float16", "shape": [1, 1, 64], "tolerance": 0.003 },
+        "qkT": { "dtype": "float32", "shape": [1, 4, 1, 20], "tolerance": 0.003 }
+      }
+    },
+    {
+      "name": "cross_qk_f16_row_b1_h4_kv16_d16",
+      "provenance": { "notes": "A float16 score row, so the diagnostic output narrows with the operands." },
+      "attrs": { "num_heads": 4, "output_qk": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [1, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [1, 4, 16, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [1, 4, 16, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float16", "shape": [1, 1, 64], "tolerance": 0.003 },
+        "qkT": { "dtype": "float16", "shape": [1, 4, 1, 16], "tolerance": 0.002 }
+      }
+    },
+    {
+      "name": "self_plain_b1_h4_cap8_d16_past3",
+      "provenance": {
+        "notes": "Self attention appending token four into a cache of capacity eight. Only the first four rows are attended; the rest keep their past values."
+      },
+      "attrs": { "num_heads": 4, "past_present_share_buffer": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 4, 8, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.05, "cosStep": 0.13 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 4, 8, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.29, "offset": 0.7 }
+        },
+        "pastSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [3] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 1, 64], "tolerance": 0.00002 },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 4, 8, 16], "tolerance": 0.000001 },
+        "presentValueT": { "dtype": "float32", "shape": [1, 4, 8, 16], "tolerance": 0.000001 }
+      }
+    },
+    {
+      "name": "self_plain_b1_h4_cap8_d16_past0",
+      "provenance": {
+        "notes": "An empty cache: the appended token is the only key, so the output is its own value row."
+      },
+      "attrs": { "num_heads": 4, "past_present_share_buffer": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 4, 8, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.05, "cosStep": 0.13 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 4, 8, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.29, "offset": 0.7 }
+        },
+        "pastSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [0] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 1, 64], "tolerance": 0.00002 },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 4, 8, 16], "tolerance": 0.000001 },
+        "presentValueT": { "dtype": "float32", "shape": [1, 4, 8, 16], "tolerance": 0.000001 }
+      }
+    },
+    {
+      "name": "self_plain_b1_h4_cap8_d16_past6",
+      "provenance": {
+        "notes": "Self-attention with a shared past/present buffer sized for 8 tokens; 6 past tokens plus the appended token leave one row unused. 4 heads, head width 16."
+      },
+      "attrs": { "num_heads": 4, "past_present_share_buffer": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 4, 8, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.05, "cosStep": 0.13 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 4, 8, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.29, "offset": 0.7 }
+        },
+        "pastSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [6] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 1, 64], "tolerance": 0.00002 },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 4, 8, 16], "tolerance": 0.000001 },
+        "presentValueT": { "dtype": "float32", "shape": [1, 4, 8, 16], "tolerance": 0.000001 }
+      }
+    },
+    {
+      "name": "self_mask_b2_h4_cap16_d16_past9",
+      "provenance": {
+        "notes": "A key-padding mask over the live length only; the cache rows past it are outside the mask and outside the softmax."
+      },
+      "attrs": { "num_heads": 4, "past_present_share_buffer": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [2, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [2, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        },
+        "maskIndexT": { "dtype": "int32", "shape": [2, 10], "data": { "kind": "cycle", "values": [1, 1, 0] } },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 4, 16, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.05, "cosStep": 0.13 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 4, 16, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.29, "offset": 0.7 }
+        },
+        "pastSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [9] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [2, 1, 64], "tolerance": 0.00002 },
+        "presentKeyT": { "dtype": "float32", "shape": [2, 4, 16, 16], "tolerance": 0.000001 },
+        "presentValueT": { "dtype": "float32", "shape": [2, 4, 16, 16], "tolerance": 0.000001 }
+      }
+    },
+    {
+      "name": "self_bias_b1_h4_cap16_d16_past5",
+      "provenance": {
+        "notes": "An additive score bias broadcast across batch and heads, the shape upstream's own self-attention test uses."
+      },
+      "attrs": { "num_heads": 4, "past_present_share_buffer": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        },
+        "attentionBiasT": {
+          "dtype": "float32",
+          "shape": [1, 1, 1, 6],
+          "data": { "kind": "fillFloat32", "sinStep": 0.37, "cosStep": 0.19, "scale": 0.3 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 4, 16, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.05, "cosStep": 0.13 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 4, 16, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.29, "offset": 0.7 }
+        },
+        "pastSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [5] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 1, 64], "tolerance": 0.00002 },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 4, 16, 16], "tolerance": 0.000001 },
+        "presentValueT": { "dtype": "float32", "shape": [1, 4, 16, 16], "tolerance": 0.000001 }
+      }
+    },
+    {
+      "name": "self_bias_per_head_b2_h4_cap12_d16_past7",
+      "provenance": {
+        "notes": "A score bias with a real head dimension, so the head stride is non-zero and a broadcast reading would mix heads."
+      },
+      "attrs": { "num_heads": 4, "past_present_share_buffer": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [2, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [2, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        },
+        "attentionBiasT": {
+          "dtype": "float32",
+          "shape": [1, 4, 1, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.41, "scale": 0.3 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 4, 12, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.05, "cosStep": 0.13 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 4, 12, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.29, "offset": 0.7 }
+        },
+        "pastSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [7] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [2, 1, 64], "tolerance": 0.00002 },
+        "presentKeyT": { "dtype": "float32", "shape": [2, 4, 12, 16], "tolerance": 0.000001 },
+        "presentValueT": { "dtype": "float32", "shape": [2, 4, 12, 16], "tolerance": 0.000001 }
+      }
+    },
+    {
+      "name": "self_mask_bias_b2_h4_cap16_d16_past9",
+      "provenance": {
+        "notes": "Both additive terms at once. The bias is added before the mask, matching the reference's term order."
+      },
+      "attrs": { "num_heads": 4, "past_present_share_buffer": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [2, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [2, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        },
+        "maskIndexT": { "dtype": "int32", "shape": [2, 10], "data": { "kind": "cycle", "values": [1, 0] } },
+        "attentionBiasT": {
+          "dtype": "float32",
+          "shape": [2, 4, 1, 10],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.53, "scale": 0.3 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 4, 16, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.05, "cosStep": 0.13 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 4, 16, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.29, "offset": 0.7 }
+        },
+        "pastSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [9] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [2, 1, 64], "tolerance": 0.00002 },
+        "presentKeyT": { "dtype": "float32", "shape": [2, 4, 16, 16], "tolerance": 0.000001 },
+        "presentValueT": { "dtype": "float32", "shape": [2, 4, 16, 16], "tolerance": 0.000001 }
+      }
+    },
+    {
+      "name": "self_bias_batch_only_b2_h4_cap12_d16_past5",
+      "provenance": {
+        "notes": "A score bias with a batch dimension but a broadcast head dimension, the mixed stride pair."
+      },
+      "attrs": { "num_heads": 4, "past_present_share_buffer": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [2, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [2, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        },
+        "attentionBiasT": {
+          "dtype": "float32",
+          "shape": [2, 1, 1, 6],
+          "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.31, "scale": 0.3 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 4, 12, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.05, "cosStep": 0.13 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 4, 12, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.29, "offset": 0.7 }
+        },
+        "pastSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [5] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [2, 1, 64], "tolerance": 0.00002 },
+        "presentKeyT": { "dtype": "float32", "shape": [2, 4, 12, 16], "tolerance": 0.000001 },
+        "presentValueT": { "dtype": "float32", "shape": [2, 4, 12, 16], "tolerance": 0.000001 }
+      }
+    },
+    {
+      "name": "self_scale_b1_h4_cap8_d16_past3",
+      "provenance": {
+        "notes": "An explicit scale on the self route, overriding the reciprocal square root of the head size."
+      },
+      "attrs": { "num_heads": 4, "past_present_share_buffer": 1, "scale": 0.2 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 4, 8, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.05, "cosStep": 0.13 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 4, 8, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.29, "offset": 0.7 }
+        },
+        "pastSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [3] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 1, 64], "tolerance": 0.00002 },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 4, 8, 16], "tolerance": 0.000001 },
+        "presentValueT": { "dtype": "float32", "shape": [1, 4, 8, 16], "tolerance": 0.000001 }
+      }
+    },
+    {
+      "name": "self_plain_b1_h8_cap448_d64_past223",
+      "provenance": { "notes": "Whisper-base decoder self-attention geometry, with the cache half filled." },
+      "attrs": { "num_heads": 8, "past_present_share_buffer": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 8, 448, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.05, "cosStep": 0.13 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 8, 448, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.29, "offset": 0.7 }
+        },
+        "pastSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [223] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 1, 512], "tolerance": 0.00005 },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 8, 448, 64], "tolerance": 0.000001 },
+        "presentValueT": { "dtype": "float32", "shape": [1, 8, 448, 64], "tolerance": 0.000001 }
+      }
+    },
+    {
+      "name": "self_mask_all_padded_b1_h2_cap8_d16_past3",
+      "provenance": {
+        "notes": "Every live key is masked. The additive filter keeps the row a uniform average rather than collapsing it to zeros."
+      },
+      "attrs": { "num_heads": 2, "past_present_share_buffer": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        },
+        "maskIndexT": { "dtype": "int32", "shape": [1, 4], "data": { "kind": "constant", "value": 0 } },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 2, 8, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.05, "cosStep": 0.13 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 2, 8, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.29, "offset": 0.7 }
+        },
+        "pastSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [3] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 1, 32], "tolerance": 0.00002 },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 2, 8, 16], "tolerance": 0.000001 },
+        "presentValueT": { "dtype": "float32", "shape": [1, 2, 8, 16], "tolerance": 0.000001 }
+      }
+    },
+    {
+      "name": "self_retain_stage4_b2_h8_cap448_d64_past31",
+      "provenance": {
+        "notes": "Self-attention with a shared past/present buffer sized for 448 tokens across 2 batches; 31 past tokens are extended by the appended token. 8 heads, head width 64."
+      },
+      "attrs": { "num_heads": 8, "past_present_share_buffer": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 1, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [2, 1, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [2, 1, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 8, 448, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.05, "cosStep": 0.13 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 8, 448, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.29, "offset": 0.7 }
+        },
+        "pastSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [31] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [2, 1, 512], "tolerance": 0.00005 },
+        "presentKeyT": { "dtype": "float32", "shape": [2, 8, 448, 64], "tolerance": 0.000001 },
+        "presentValueT": { "dtype": "float32", "shape": [2, 8, 448, 64], "tolerance": 0.000001 }
+      }
+    },
+    {
+      "name": "self_retain_stage8_f16_b4_h8_cap448_d64_past223",
+      "provenance": {
+        "notes": "A cache large enough for the retain copy to stage the full eight words per lane, the deepest rung, in float16."
+      },
+      "attrs": { "num_heads": 8, "past_present_share_buffer": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [4, 1, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [4, 1, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [4, 1, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        },
+        "pastKeyT": {
+          "dtype": "float16",
+          "shape": [4, 8, 448, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.05, "cosStep": 0.13 }
+        },
+        "pastValueT": {
+          "dtype": "float16",
+          "shape": [4, 8, 448, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.29, "offset": 0.7 }
+        },
+        "pastSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [223] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float16", "shape": [4, 1, 512], "tolerance": 0.003 },
+        "presentKeyT": { "dtype": "float16", "shape": [4, 8, 448, 64], "tolerance": 0.001 },
+        "presentValueT": { "dtype": "float16", "shape": [4, 8, 448, 64], "tolerance": 0.001 }
+      }
+    },
+    {
+      "name": "self_plain_f16_b1_h4_cap8_d16_past3",
+      "provenance": {
+        "notes": "The float16 twin of the plain self route, including the float16 cache the append pass writes."
+      },
+      "attrs": { "num_heads": 4, "past_present_share_buffer": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [1, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [1, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [1, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        },
+        "pastKeyT": {
+          "dtype": "float16",
+          "shape": [1, 4, 8, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.05, "cosStep": 0.13 }
+        },
+        "pastValueT": {
+          "dtype": "float16",
+          "shape": [1, 4, 8, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.29, "offset": 0.7 }
+        },
+        "pastSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [3] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float16", "shape": [1, 1, 64], "tolerance": 0.003 },
+        "presentKeyT": { "dtype": "float16", "shape": [1, 4, 8, 16], "tolerance": 0.001 },
+        "presentValueT": { "dtype": "float16", "shape": [1, 4, 8, 16], "tolerance": 0.001 }
+      }
+    },
+    {
+      "name": "self_mask_bias_f16_b2_h4_cap16_d16_past9",
+      "provenance": {
+        "notes": "Float16 operands with both additive terms, so the bias narrows with the tensor type while the score stays float32."
+      },
+      "attrs": { "num_heads": 4, "past_present_share_buffer": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [2, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [2, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [2, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        },
+        "maskIndexT": { "dtype": "int32", "shape": [2, 10], "data": { "kind": "cycle", "values": [1, 0] } },
+        "attentionBiasT": {
+          "dtype": "float16",
+          "shape": [2, 4, 1, 10],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.53, "scale": 0.3 }
+        },
+        "pastKeyT": {
+          "dtype": "float16",
+          "shape": [2, 4, 16, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.05, "cosStep": 0.13 }
+        },
+        "pastValueT": {
+          "dtype": "float16",
+          "shape": [2, 4, 16, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.29, "offset": 0.7 }
+        },
+        "pastSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [9] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float16", "shape": [2, 1, 64], "tolerance": 0.003 },
+        "presentKeyT": { "dtype": "float16", "shape": [2, 4, 16, 16], "tolerance": 0.001 },
+        "presentValueT": { "dtype": "float16", "shape": [2, 4, 16, 16], "tolerance": 0.001 }
+      }
+    },
+    {
+      "name": "cross_single_split_boundary_kv255_d64_float32",
+      "provenance": {
+        "notes": "Adjacent key lengths straddle the default computed split count; direct and split/merge must preserve the same full attention row."
+      },
+      "attrs": { "num_heads": 6 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 384],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 6, 255, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 6, 255, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [1, 1, 384], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "cross_single_split_boundary_kv255_d64_float16",
+      "provenance": {
+        "notes": "Adjacent key lengths straddle the default computed split count; direct and split/merge must preserve the same full attention row."
+      },
+      "attrs": { "num_heads": 6 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [1, 1, 384],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [1, 6, 255, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [1, 6, 255, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float16", "shape": [1, 1, 384], "tolerance": 0.003 } }
+    },
+    {
+      "name": "cross_single_split_boundary_kv256_d64_float32",
+      "provenance": {
+        "notes": "Adjacent key lengths straddle the default computed split count; direct and split/merge must preserve the same full attention row."
+      },
+      "attrs": { "num_heads": 6 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 384],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 6, 256, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 6, 256, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [1, 1, 384], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "cross_single_split_boundary_kv256_d64_float16",
+      "provenance": {
+        "notes": "Adjacent key lengths straddle the default computed split count; direct and split/merge must preserve the same full attention row."
+      },
+      "attrs": { "num_heads": 6 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [1, 1, 384],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [1, 6, 256, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [1, 6, 256, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float16", "shape": [1, 1, 384], "tolerance": 0.003 } }
+    },
+    {
+      "name": "cross_single_split_boundary_kv257_d64_float32",
+      "provenance": {
+        "notes": "Adjacent key lengths straddle the default computed split count; direct and split/merge must preserve the same full attention row."
+      },
+      "attrs": { "num_heads": 6 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 384],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 6, 257, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 6, 257, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [1, 1, 384], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "cross_single_split_boundary_kv257_d64_float16",
+      "provenance": {
+        "notes": "Adjacent key lengths straddle the default computed split count; direct and split/merge must preserve the same full attention row."
+      },
+      "attrs": { "num_heads": 6 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [1, 1, 384],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [1, 6, 257, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [1, 6, 257, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float16", "shape": [1, 1, 384], "tolerance": 0.003 } }
+    },
+    {
+      "name": "cross_single_split_mask_qk_kv64_d64_float32",
+      "provenance": {
+        "notes": "Single-partition cross attention emits the scaled, finitely masked pre-softmax QK row while directly normalizing its output."
+      },
+      "attrs": { "num_heads": 6, "output_qk": 1, "mask_filter_value": -4 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 384],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 6, 64, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 6, 64, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        },
+        "maskIndexT": { "dtype": "int32", "shape": [1, 64], "data": { "kind": "cycle", "values": [1, 0, 1] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 1, 384], "tolerance": 0.00002 },
+        "qkT": { "dtype": "float32", "shape": [1, 6, 1, 64], "tolerance": 0.00002 }
+      }
+    },
+    {
+      "name": "cross_single_split_mask_qk_kv64_d64_float16",
+      "provenance": {
+        "notes": "Single-partition cross attention emits the scaled, finitely masked pre-softmax QK row while directly normalizing its output."
+      },
+      "attrs": { "num_heads": 6, "output_qk": 1, "mask_filter_value": -4 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [1, 1, 384],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [1, 6, 64, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [1, 6, 64, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        },
+        "maskIndexT": { "dtype": "int32", "shape": [1, 64], "data": { "kind": "cycle", "values": [1, 0, 1] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float16", "shape": [1, 1, 384], "tolerance": 0.003 },
+        "qkT": { "dtype": "float16", "shape": [1, 6, 1, 64], "tolerance": 0.003 }
+      }
+    },
+    {
+      "name": "cross_multi_split_mask_qk_kv257_d64_float32",
+      "provenance": {
+        "notes": "Cross-attention with 6 heads and head width 64, a 257-length key/value cache (256 plus one), a padding mask, and a custom mask_filter_value of -4. The diagnostic qkT output is requested alongside outputT."
+      },
+      "attrs": { "num_heads": 6, "output_qk": 1, "mask_filter_value": -4 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 384],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 6, 257, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 6, 257, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        },
+        "maskIndexT": { "dtype": "int32", "shape": [1, 257], "data": { "kind": "cycle", "values": [1, 0, 1] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 1, 384], "tolerance": 0.00002 },
+        "qkT": { "dtype": "float32", "shape": [1, 6, 1, 257], "tolerance": 0.00002 }
+      }
+    },
+    {
+      "name": "cross_multi_split_mask_qk_kv257_d64_float16",
+      "provenance": {
+        "notes": "Cross-attention with 6 heads and head width 64, a 257-length key/value cache (256 plus one), a padding mask, and a custom mask_filter_value of -4, with float16 query/key/value/output. The diagnostic qkT output is requested alongside outputT."
+      },
+      "attrs": { "num_heads": 6, "output_qk": 1, "mask_filter_value": -4 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [1, 1, 384],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [1, 6, 257, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [1, 6, 257, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        },
+        "maskIndexT": { "dtype": "int32", "shape": [1, 257], "data": { "kind": "cycle", "values": [1, 0, 1] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float16", "shape": [1, 1, 384], "tolerance": 0.003 },
+        "qkT": { "dtype": "float16", "shape": [1, 6, 1, 257], "tolerance": 0.003 }
+      }
+    },
+    {
+      "name": "qk_policy_b1_h4_kv256_d512_f32",
+      "provenance": {
+        "notes": "QK policy boundary or head-width fixture, checked against the CPU reference at the f32 attention tolerance."
+      },
+      "attrs": { "num_heads": 4 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 2048],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 4, 256, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 4, 256, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [1, 1, 2048], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "qk_policy_b1_h2_kv256_d1024_f32",
+      "provenance": {
+        "notes": "QK policy boundary or head-width fixture, checked against the CPU reference at the f32 attention tolerance."
+      },
+      "attrs": { "num_heads": 2 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 2048],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 2, 256, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 2, 256, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [1, 1, 2048], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "qk_policy_b1_h1_kv256_d2048_f32",
+      "provenance": {
+        "notes": "QK policy boundary or head-width fixture, checked against the CPU reference at the f32 attention tolerance."
+      },
+      "attrs": { "num_heads": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 2048],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 256, 2048],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 256, 2048],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [1, 1, 2048], "tolerance": 0.00002 } },
+      "requires": {
+        "limits": {
+          "maxComputeInvocationsPerWorkgroup": 512,
+          "maxComputeWorkgroupSizeX": 512,
+          "maxComputeWorkgroupStorageSize": 22528
+        }
+      }
+    },
+    {
+      "name": "qk_policy_b1_h15_kv257_d128_f32",
+      "provenance": {
+        "notes": "QK policy boundary or head-width fixture, checked against the CPU reference at the f32 attention tolerance."
+      },
+      "attrs": { "num_heads": 15 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 1920],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 15, 257, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 15, 257, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [1, 1, 1920], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "qk_policy_b1_h16_kv257_d128_f32",
+      "provenance": {
+        "notes": "QK policy boundary or head-width fixture, checked against the CPU reference at the f32 attention tolerance."
+      },
+      "attrs": { "num_heads": 16 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 2048],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 16, 257, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 16, 257, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [1, 1, 2048], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "qk_policy_b1_h17_kv257_d128_f32",
+      "provenance": {
+        "notes": "QK policy boundary or head-width fixture, checked against the CPU reference at the f32 attention tolerance."
+      },
+      "attrs": { "num_heads": 17 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 2176],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 17, 257, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 17, 257, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [1, 1, 2176], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "qk_policy_b1_h6_kv256_d32_f32",
+      "provenance": {
+        "notes": "QK policy boundary or head-width fixture, checked against the CPU reference at the f32 attention tolerance."
+      },
+      "attrs": { "num_heads": 6 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 6, 256, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 6, 256, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [1, 1, 192], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "qk_policy_b1_h6_kv256_d256_f32",
+      "provenance": {
+        "notes": "QK policy boundary or head-width fixture, checked against the CPU reference at the f32 attention tolerance."
+      },
+      "attrs": { "num_heads": 6 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 1536],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 6, 256, 256],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 6, 256, 256],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [1, 1, 1536], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "qk_policy_b1_h6_kv1500_d64_f32",
+      "provenance": {
+        "notes": "QK policy boundary or head-width fixture, checked against the CPU reference at the f32 attention tolerance."
+      },
+      "attrs": { "num_heads": 6 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 384],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 6, 1500, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 6, 1500, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [1, 1, 384], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "qk_policy_b4_h6_kv1500_d128_f32",
+      "provenance": {
+        "notes": "QK policy boundary or head-width fixture, checked against the CPU reference at the f32 attention tolerance."
+      },
+      "attrs": { "num_heads": 6 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [4, 1, 768],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [4, 6, 1500, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [4, 6, 1500, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [4, 1, 768], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "qk_policy_f32_query_f16_scores_kv64_d64",
+      "provenance": {
+        "notes": "Mixed-dtype diagnostic scores exercise the f32 QK reduction at the mixed-QK half-output tolerance, on both sides of the split boundary."
+      },
+      "attrs": { "num_heads": 6, "output_qk": 1, "mask_filter_value": -4 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 384],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 6, 64, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 6, 64, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        },
+        "maskIndexT": { "dtype": "int32", "shape": [1, 64], "data": { "kind": "cycle", "values": [1, 0, 1] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 1, 384], "tolerance": 0.00002 },
+        "qkT": { "dtype": "float16", "shape": [1, 6, 1, 64], "tolerance": 0.002 }
+      }
+    },
+    {
+      "name": "qk_policy_f32_query_f16_scores_kv257_d64",
+      "provenance": {
+        "notes": "Mixed-dtype diagnostic scores exercise the f32 QK reduction at the mixed-QK half-output tolerance, on both sides of the split boundary."
+      },
+      "attrs": { "num_heads": 6, "output_qk": 1, "mask_filter_value": -4 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 384],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 6, 257, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 6, 257, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "offset": 0.7 }
+        },
+        "maskIndexT": { "dtype": "int32", "shape": [1, 257], "data": { "kind": "cycle", "values": [1, 0, 1] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 1, 384], "tolerance": 0.00002 },
+        "qkT": { "dtype": "float16", "shape": [1, 6, 1, 257], "tolerance": 0.002 }
+      }
+    }
+  ]
+}