Xenova HF Staff commited on 8 days ago

Commit

43f449f

verified ·

1 Parent(s): 2157e54

sync 6fdf6301e2bb

Browse files

Files changed (21) hide show

README.md +5 -2
build/webgpu/attention-rank4-tiled.wgsl.jinja +1 -25
build/webgpu/attn-flash-decode-splitk-merge.wgsl.jinja +2 -9
build/webgpu/attn-flash-decode-splitk.wgsl.jinja +14 -115
build/webgpu/attn-flash-online.wgsl.jinja +8 -44
build/webgpu/attn-flash-prefill-cluster.wgsl.jinja +41 -194
build/webgpu/attn-flash-q32-broadcast.wgsl.jinja +46 -11
build/webgpu/attn-materialized-apply-f32.wgsl.jinja +3 -2
build/webgpu/attn-materialized-rowstats-combine-f32.wgsl.jinja +9 -5
build/webgpu/attn-materialized-score-f32.wgsl.jinja +6 -85
build/webgpu/attn-materialized-sgmat-f32.wgsl.jinja +28 -54
build/webgpu/attn-materialized-softmax-f32.wgsl.jinja +6 -41
build/webgpu/attn-online-scalar.wgsl.jinja +39 -28
build/webgpu/attn-small-head-parallel.wgsl.jinja +3 -9
build/webgpu/attn-small-head-value.wgsl.jinja +198 -0
build/webgpu/bench.json +0 -0
build/webgpu/manifest.json +224 -362
build/webgpu/metadata.json +24 -21
build/webgpu/mha-small-seq-blocked.wgsl.jinja +5 -4
build/webgpu/mha-small-seq.wgsl.jinja +26 -6
build/webgpu/test.json +0 -0

README.md CHANGED Viewed

@@ -52,6 +52,8 @@ Attributes and default values (overridable per request):
 One implementation is selected per call from the device capabilities, the request shapes and the dtypes; these notes say what each one covers.
 - `qkv_bias_small_seq_blocked` — Blocked whole-head attention for short key sequences with bias and optional causal masking. A workgroup stages one head's K and V once, then assigns each query in a small block to a group of key lanes whose register-resident online-softmax partials are merged in one pass.
 - `qkv_no_bias_small_seq_blocked` — Blocked whole-head attention for short key sequences without bias and with optional causal masking. A workgroup stages one head's K and V once, then assigns each query in a small block to a group of key lanes whose register-resident online-softmax partials are merged in one pass.
 - `qkv_no_bias_small_seq` — Whole-head attention for short bidirectional float32 requests without bias. One workgroup stages a head's complete key and value planes, and each participating invocation owns one query row through scoring, normalization, and context accumulation.
@@ -100,7 +102,7 @@ Some implementation variants require `subgroup-matrix`, `shader-f16`, and `subgr
 - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
-- [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 - [`attention-rank4-tiled.wgsl.jinja`](build/webgpu/attention-rank4-tiled.wgsl.jinja)
 - [`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)
@@ -114,13 +116,14 @@ Some implementation variants require `subgroup-matrix`, `shader-f16`, and `subgr
 - [`attn-materialized-softmax-f32.wgsl.jinja`](build/webgpu/attn-materialized-softmax-f32.wgsl.jinja)
 - [`attn-online-scalar.wgsl.jinja`](build/webgpu/attn-online-scalar.wgsl.jinja)
 - [`attn-small-head-parallel.wgsl.jinja`](build/webgpu/attn-small-head-parallel.wgsl.jinja)
 - [`mha-small-seq-blocked.wgsl.jinja`](build/webgpu/mha-small-seq-blocked.wgsl.jinja)
 - [`mha-small-seq.wgsl.jinja`](build/webgpu/mha-small-seq.wgsl.jinja)
 ## Use with `@huggingface/kernels`
 ```sh
-npm install --save-exact @huggingface/kernels@0.0.1-preview.2
 ```
 Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.

 One implementation is selected per call from the device capabilities, the request shapes and the dtypes; these notes say what each one covers.
+- `qkv_no_bias_small_head_value_subgroups` — Cooperatively reduces values across key lanes and reuses keys/values for two queries when device workgroup storage fits. Uses portable-width subgroup collectives with f32 scores, probabilities and accumulation. Short key sequences use the serial value path.
+- `qkv_no_bias_small_head_value_tree` — Cooperatively reduces values across key lanes and reuses keys/values for two queries when device workgroup storage fits. Uses shared-memory tree reductions with f32 scores, probabilities and accumulation. Short key sequences use the serial value path.
 - `qkv_bias_small_seq_blocked` — Blocked whole-head attention for short key sequences with bias and optional causal masking. A workgroup stages one head's K and V once, then assigns each query in a small block to a group of key lanes whose register-resident online-softmax partials are merged in one pass.
 - `qkv_no_bias_small_seq_blocked` — Blocked whole-head attention for short key sequences without bias and with optional causal masking. A workgroup stages one head's K and V once, then assigns each query in a small block to a group of key lanes whose register-resident online-softmax partials are merged in one pass.
 - `qkv_no_bias_small_seq` — Whole-head attention for short bidirectional float32 requests without bias. One workgroup stages a head's complete key and value planes, and each participating invocation owns one query row through scoring, normalization, and context accumulation.
 - [`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
 - [`attention-rank4-tiled.wgsl.jinja`](build/webgpu/attention-rank4-tiled.wgsl.jinja)
 - [`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)
 - [`attn-materialized-softmax-f32.wgsl.jinja`](build/webgpu/attn-materialized-softmax-f32.wgsl.jinja)
 - [`attn-online-scalar.wgsl.jinja`](build/webgpu/attn-online-scalar.wgsl.jinja)
 - [`attn-small-head-parallel.wgsl.jinja`](build/webgpu/attn-small-head-parallel.wgsl.jinja)
+- [`attn-small-head-value.wgsl.jinja`](build/webgpu/attn-small-head-value.wgsl.jinja)
 - [`mha-small-seq-blocked.wgsl.jinja`](build/webgpu/mha-small-seq-blocked.wgsl.jinja)
 - [`mha-small-seq.wgsl.jinja`](build/webgpu/mha-small-seq.wgsl.jinja)
 ## Use with `@huggingface/kernels`
 ```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.3
 ```
 Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.

build/webgpu/attention-rank4-tiled.wgsl.jinja CHANGED Viewed

@@ -1,4 +1,4 @@
-{% set layout = layout if layout is defined else "bnsh" %}
 {{ env.wgsl.resourceDeclarations }}
 // Thread-per-query online-softmax flash fallback for rank-4 (BNSH) attention with
@@ -16,13 +16,11 @@ const BLOCK_M: u32 = {{ blockM }}u;
 var<workgroup> acc: array<f32, {{ vHeadCap * blockM }}>;
 {% set ATTN_SCALE_DIM = "params.headSize" %}
-{% if ATTN_SCALE_DIM is not defined %}{% set ATTN_SCALE_DIM = "HEAD_DIM" %}{% endif %}
 fn scale_value() -> f32 {
   if (params.scale != 0.0) { return params.scale; }
   return inverseSqrt(f32({{ ATTN_SCALE_DIM }}));
 }
 fn kv_head(q_head: u32) -> u32 {
   return q_head / (params.qHeads / params.kvHeads);
 }
@@ -51,14 +49,8 @@ fn main(
   if (qs >= params.qSeq) { return; }
   let kh = kv_head(qh);
-{% if layout == "bsh" %}
   // Token-major packed QKV: row = (batch*seq + token)*HIDDEN + head*head_dim.
   let qBase = (batch * params.qSeq + qs) * params.qHidden + qh * params.headSize;
-{% else %}
-  let qBase = ((batch * params.qHeads + qh) * params.qSeq + qs) * params.headSize;
-  let kBase = (batch * params.kvHeads + kh) * params.kvSeq;
-  let vBase = (batch * params.kvHeads + kh) * params.kvSeq;
-{% endif %}
   let scale = scale_value();
   for (var d: u32 = 0u; d < params.vHeadSize; d = d + 1u) {
@@ -76,21 +68,13 @@ fn main(
   for (var ks: u32 = 0u; ks < maxKj; ks = ks + 1u) {
     var masked = false;
-{% if hasMask and maskIsBool %}
-    let mIdxB = batch * params.maskBatchStride + qh * params.maskHeadStride + qs * params.maskSeqStride + ks;
-    masked = attn_mask[mIdxB] == 0u;
-{% endif %}
     // A rejected bool-mask key has zero softmax mass. Skipping it is safe here:
     // this thread-per-query kernel has no barriers inside the key loop.
     if (masked) { continue; }
     // Adjacent query threads load the same K/V row for this key.
     var score: f32 = -3.4028234663852886e38;
     if (!masked) {
-{% if layout == "bsh" %}
       let kRow = (batch * params.kvSeq + ks) * params.kvHidden + kh * params.headSize;
-{% else %}
-      let kRow = (kBase + ks) * params.headSize;
-{% endif %}
       var dot: f32 = 0.0;
       for (var d: u32 = 0u; d < params.headSize; d = d + 1u) {
         dot = dot + f32(q[qBase + d]) * f32(k[kRow + d]);
@@ -110,22 +94,14 @@ fn main(
     let weight = exp(score - next_max);
     running_max = next_max;
     running_denom = running_denom * prev_scale + weight;
-{% if layout == "bsh" %}
     let vRow = (batch * params.kvSeq + ks) * params.vHidden + kh * params.vHeadSize;
-{% else %}
-    let vRow = (vBase + ks) * params.vHeadSize;
-{% endif %}
     for (var d: u32 = 0u; d < params.vHeadSize; d = d + 1u) {
       acc[d * BLOCK_M + tid] = acc[d * BLOCK_M + tid] * prev_scale + weight * f32(v[vRow + d]);
     }
   }
   let inv_denom = select(0.0, 1.0 / running_denom, running_denom > 0.0);
-{% if layout == "bsh" %}
   let yBase = (batch * params.qSeq + qs) * (params.qHeads * params.vHeadSize) + qh * params.vHeadSize;
-{% else %}
-  let yBase = ((batch * params.qHeads + qh) * params.qSeq + qs) * params.vHeadSize;
-{% endif %}
   for (var d: u32 = 0u; d < params.vHeadSize; d = d + 1u) {
     y[yBase + d] = {{ scalar }}(acc[d * BLOCK_M + tid] * inv_denom);
   }

+{% set maskIsBool = false %}
 {{ env.wgsl.resourceDeclarations }}
 // Thread-per-query online-softmax flash fallback for rank-4 (BNSH) attention with
 var<workgroup> acc: array<f32, {{ vHeadCap * blockM }}>;
 {% set ATTN_SCALE_DIM = "params.headSize" %}
 fn scale_value() -> f32 {
   if (params.scale != 0.0) { return params.scale; }
   return inverseSqrt(f32({{ ATTN_SCALE_DIM }}));
 }
 fn kv_head(q_head: u32) -> u32 {
   return q_head / (params.qHeads / params.kvHeads);
 }
   if (qs >= params.qSeq) { return; }
   let kh = kv_head(qh);
   // Token-major packed QKV: row = (batch*seq + token)*HIDDEN + head*head_dim.
   let qBase = (batch * params.qSeq + qs) * params.qHidden + qh * params.headSize;
   let scale = scale_value();
   for (var d: u32 = 0u; d < params.vHeadSize; d = d + 1u) {
   for (var ks: u32 = 0u; ks < maxKj; ks = ks + 1u) {
     var masked = false;
     // A rejected bool-mask key has zero softmax mass. Skipping it is safe here:
     // this thread-per-query kernel has no barriers inside the key loop.
     if (masked) { continue; }
     // Adjacent query threads load the same K/V row for this key.
     var score: f32 = -3.4028234663852886e38;
     if (!masked) {
       let kRow = (batch * params.kvSeq + ks) * params.kvHidden + kh * params.headSize;
       var dot: f32 = 0.0;
       for (var d: u32 = 0u; d < params.headSize; d = d + 1u) {
         dot = dot + f32(q[qBase + d]) * f32(k[kRow + d]);
     let weight = exp(score - next_max);
     running_max = next_max;
     running_denom = running_denom * prev_scale + weight;
     let vRow = (batch * params.kvSeq + ks) * params.vHidden + kh * params.vHeadSize;
     for (var d: u32 = 0u; d < params.vHeadSize; d = d + 1u) {
       acc[d * BLOCK_M + tid] = acc[d * BLOCK_M + tid] * prev_scale + weight * f32(v[vRow + d]);
     }
   }
   let inv_denom = select(0.0, 1.0 / running_denom, running_denom > 0.0);
   let yBase = (batch * params.qSeq + qs) * (params.qHeads * params.vHeadSize) + qh * params.vHeadSize;
   for (var d: u32 = 0u; d < params.vHeadSize; d = d + 1u) {
     y[yBase + d] = {{ scalar }}(acc[d * BLOCK_M + tid] * inv_denom);
   }

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

@@ -1,9 +1,6 @@
-{% if scalar == "f16" %}
-enable f16;
-{% endif %}
 {% if splitQueries is not defined %}{% set splitQueries = false %}{% endif %}
-{% if hasGate is not defined %}{% set hasGate = false %}{% endif %}
 {% set qSeq = qSeq | default(0) %}
 {{ env.wgsl.resourceDeclarations }}
 // Split-K flash decode, pass 2 of 2. Combines the per-split un-normalized online
@@ -49,6 +46,7 @@ fn shifted_value(value: f32, maxValue: f32) -> f32 {
   let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
   return select(value - maxValue, 0.0, equalFiniteMax);
 }
 fn exp_shift(value: f32, maxValue: f32) -> f32 {
   return exp(shifted_value(value, maxValue));
 }
@@ -104,11 +102,6 @@ fn main(
   // V bias row base: 2 * qHidden (skip the packed Q and K bias blocks) + this head.
   let vBiasBase = 2u * Q_HIDDEN + h * HEAD_DIM + d4 * 4u;
   outValue = outValue + vec4<f32>(bias[vBiasBase], bias[vBiasBase + 1u], bias[vBiasBase + 2u], bias[vBiasBase + 3u]);
-{% endif %}
-{% if hasGate %}
-  // The gated route multiplies the normalized attention output elementwise by its gate.
-  let gateV = vec4<f32>(gate[qBaseV4 + d4]);
-  outValue = outValue * (vec4<f32>(1.0) / (vec4<f32>(1.0) + exp(-gateV)));
 {% endif %}
   output[qBaseV4 + d4] = vec4<{{ scalar }}>(outValue);
 }

 {% if splitQueries is not defined %}{% set splitQueries = false %}{% endif %}
 {% set qSeq = qSeq | default(0) %}
+{% set qHidden = qHidden | default(0) %}
 {{ env.wgsl.resourceDeclarations }}
 // Split-K flash decode, pass 2 of 2. Combines the per-split un-normalized online
   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));
 }
   // V bias row base: 2 * qHidden (skip the packed Q and K bias blocks) + this head.
   let vBiasBase = 2u * Q_HIDDEN + h * HEAD_DIM + d4 * 4u;
   outValue = outValue + vec4<f32>(bias[vBiasBase], bias[vBiasBase + 1u], bias[vBiasBase + 2u], bias[vBiasBase + 3u]);
 {% endif %}
   output[qBaseV4 + d4] = vec4<{{ scalar }}>(outValue);
 }

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

@@ -1,11 +1,6 @@
 {% if useSubgroups is not defined %}{% set useSubgroups = true %}{% endif %}
 {% if splitQueries is not defined %}{% set splitQueries = false %}{% endif %}
-{% if quantizedCache is not defined %}{% set quantizedCache = false %}{% endif %}
-{% if cacheSeqlens is not defined %}{% set cacheSeqlens = false %}{% endif %}
 {% if hasMask is not defined %}{% set hasMask = false %}{% endif %}
-{% if maskIsBool is not defined %}{% set maskIsBool = false %}{% endif %}
-{% set scale = scale | default("0.0") %}
-{% if hasRotary is not defined %}{% set hasRotary = false %}{% endif %}
 {% set qSeq = qSeq | default(0) %}
 {% set splitKWorkgroupSize = workgroupSizeSpec if workgroupSizeSpec is defined else tunables.WORKGROUP_SIZE %}
 {% if useSubgroups %}
@@ -13,7 +8,8 @@ enable subgroups;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
-// Split-K flash attention, pass 1 of 2; the merge pass follows. This geometry
 // handles decode and short-query, long-context prefill inputs.
 //
 // The non-split flash decode launches only `batch * numHeads` workgroups, each
@@ -56,6 +52,7 @@ fn shifted_value(value: f32, maxValue: f32) -> f32 {
   let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
   return select(value - maxValue, 0.0, equalFiniteMax);
 }
 fn exp_shift(value: f32, maxValue: f32) -> f32 {
   return exp(shifted_value(value, maxValue));
 }
@@ -63,7 +60,7 @@ fn exp_shift(value: f32, maxValue: f32) -> f32 {
 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) %}
 {% set jGroups = (splitKWorkgroupSize / headDimV4)|int %}
 {% set jSplitV = (splitKWorkgroupSize % headDimV4 == 0) and (jGroups >= 2) %}
 {% if coopQk %}
@@ -132,41 +129,9 @@ fn combine_partials(m: f32, d: f32, lidx: u32, sgSize: u32) -> vec2<f32> {
   return combinedMD;
 }
 {% else %}
-{% set mdStreamed = mdStreams is defined %}
-{% set mdStreams = mdStreams if mdStreams is defined else 1 %}
-{% set mdExtent = "WG" if mdStreams == 1 else "WG * " ~ mdStreams ~ "u" %}
 var<workgroup> partialM: array<f32, {{ mdExtent }}>;
 var<workgroup> partialD: array<f32, {{ mdExtent }}>;
-{% if mdStreamed %}
-// In-place fold of {{ mdStreams }} streams. Input partials occupy
-// partialM/partialD; stream s returns its merged pair in slot s * WG.
-fn combine_partials_streams(lidx: u32) {
-  workgroupBarrier();
-  var stride = WG / 2u;
-  loop {
-    if (stride == 0u) {
-      break;
-    }
-    if (lidx < stride) {
-{% for s in range(mdStreams) %}
-      {
-        let slot = {{ s }}u * WG + lidx;
-        let m1 = partialM[slot];
-        let d1 = partialD[slot];
-        let m2 = partialM[slot + stride];
-        let d2 = partialD[slot + stride];
-        let mNew = max(m1, m2);
-        partialD[slot] = d1 * exp_shift(m1, mNew) + d2 * exp_shift(m2, mNew);
-        partialM[slot] = mNew;
-      }
-{% endfor %}
-    }
-    workgroupBarrier();
-    stride = stride / 2u;
-  }
-}
-{% else %}
 fn combine_partials(m: f32, d: f32, lidx: u32) -> vec2<f32> {
   partialM[lidx] = m;
@@ -196,54 +161,12 @@ fn combine_partials(m: f32, d: f32, lidx: u32) -> vec2<f32> {
   return merged;
 }
 {% endif %}
-{% endif %}
-{% if ATTN_SCALE_DIM is not defined %}{% set ATTN_SCALE_DIM = "HEAD_DIM" %}{% endif %}
-fn scale_value() -> f32 {
   if (params.scale != 0.0) { return params.scale; }
   return inverseSqrt(f32({{ ATTN_SCALE_DIM }}));
 }
-{% if quantizedCache %}
-{% macro emit_quant_scale4(kind, scaleBuffer) %}
-fn {{ kind }}scale4(d4: u32, hk: u32) -> vec4<f32> {
-  if (params.perChannel == 0u) {
-    return vec4<f32>({{ scaleBuffer }}[0]);
-  }
-  let base = hk * HEAD_DIM + d4 * 4u;
-  return vec4<f32>(
-    {{ scaleBuffer }}[base],
-    {{ scaleBuffer }}[base + 1u],
-    {{ scaleBuffer }}[base + 2u],
-    {{ scaleBuffer }}[base + 3u]
-  );
-}
-{%- endmacro %}
-{%- macro emit_quant_load4(format, kind, buffer, scaleBuffer) %}
-{{ emit_quant_scale4(kind, scaleBuffer) }}
-fn load_{{ kind }}4(indexV4: u32, d4: u32, hk: u32) -> vec4<f32> {
-{%- if format == "int8" %}
-  return vec4<f32>({{ buffer }}[indexV4]) * {{ kind }}scale4(d4, hk);
-{%- else %}
-  // Two elements cover this vec4: each carries two +8-biased nibbles, low first.
-  let rowBase = indexV4 - d4;
-  let lo = {{ buffer }}[rowBase + d4 * 2u];
-  let hi = {{ buffer }}[rowBase + d4 * 2u + 1u];
-  let nibbles = vec4<i32>(
-    i32(lo & 0xFu), i32((lo >> 4u) & 0xFu),
-    i32(hi & 0xFu), i32((hi >> 4u) & 0xFu)
-  );
-  let signed = nibbles - vec4<i32>(8);
-  return vec4<f32>(signed) * {{ kind }}scale4(d4, hk);
-{%- endif %}
-}
-{%- endmacro %}
-{{ emit_quant_load4("int8", "key", "key", "k_scale") }}
-{{ emit_quant_load4("int8", "value", "value", "v_scale") }}
-{% else %}
 fn load_key4(indexV4: u32) -> vec4<f32> {
   return vec4<f32>(key[indexV4]);
 }
@@ -251,16 +174,17 @@ fn load_key4(indexV4: u32) -> vec4<f32> {
 fn load_value4(indexV4: u32) -> vec4<f32> {
   return vec4<f32>(value[indexV4]);
 }
-{% endif %}
 {% if hasBias %}
 // Packed [Q; K; V] bias rows (token-independent). The Q bias folds into the
 // query row before the Q.K dots; the K bias adds a constant to every key score
 // that softmax cancels, so it is skipped; the V bias is token-independent and
 // is applied once in the merge pass after the final normalize.
 fn load_bias4(base: u32, d4: u32) -> vec4<f32> {
   let offset = base + d4 * 4u;
-  return vec4<f32>(bias[offset], bias[offset + 1u], bias[offset + 2u], bias[offset + 3u]);
 }
 {% endif %}
@@ -290,13 +214,7 @@ fn main(
   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[b]) + 1u);
-{% else %}
   let kvSeq = cacheSeq;
-{% endif %}
   // Query row (decode uses token zero; short-query prefill folds the token into wg.x).
 {% if splitQueries %}
@@ -317,7 +235,6 @@ fn main(
     splitEnd = kvSeq;
   }
-{% set hasBias = hasBias is defined and hasBias %}
   for (var d4 = tid; d4 < HEAD_DIM_V4; d4 = d4 + WG) {
     var qv = vec4<f32>(query[qBaseV4 + d4]);
 {% if hasBias %}
@@ -357,7 +274,7 @@ fn main(
       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{% if quantizedCache %}, d4, hKv{% endif %}));
         }
       }
       let sj = subgroupAdd(accS);
@@ -376,7 +293,7 @@ fn main(
       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{% if quantizedCache %}, d4, hKv{% endif %}));
       }
       score = acc * scale;
       m = score;
@@ -391,21 +308,10 @@ fn main(
       let maskQuery = 0u;
 {% endif %}
       let maskIndex = b * params.maskBatchStride + h * params.maskHeadStride + maskQuery * params.maskSeqStride + kj;
-{% if maskIsBool %}
-      // A rejected bool-mask key contributes no probability mass. The merge
-      // pass already maps a zero global denominator to an all-zero output row.
-      if (attn_mask[maskIndex] == 0u) {
-        keyAllowed = false;
-        score = -FLT_MAX;
-        dPart = 0.0;
-      }
-{% else %}
       score = score + f32(attn_mask[maskIndex]);
-{% endif %}
       m = 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.
@@ -421,7 +327,6 @@ fn main(
     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
@@ -434,10 +339,7 @@ fn main(
     loop {
       if (jj >= tileCount) { break; }
       vacc = vacc + probs[jj] * load_value4(
-        kvBaseV4 + (kjBase + jj) * kvTokenStrideV4 + d4v{% if quantizedCache %},
-        d4v,
-        hKv{% endif %}
-      );
       jj = jj + J_GROUPS;
     }
     vacc_sh[tid] = vacc;
@@ -455,10 +357,7 @@ fn main(
       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{% if quantizedCache %},
-          d4,
-          hKv{% endif %}
-        );
       }
       running_out[d4] = running_out[d4] * correction + vSum;
     }

 {% if useSubgroups is not defined %}{% set useSubgroups = true %}{% endif %}
 {% if splitQueries is not defined %}{% set splitQueries = false %}{% endif %}
 {% if hasMask is not defined %}{% set hasMask = false %}{% endif %}
 {% set qSeq = qSeq | default(0) %}
 {% set splitKWorkgroupSize = workgroupSizeSpec if workgroupSizeSpec is defined else tunables.WORKGROUP_SIZE %}
 {% if useSubgroups %}
 {% 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
   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 %}
   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;
   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]);
 }
 {% if hasBias %}
 // Packed [Q; K; V] bias rows (token-independent). The Q bias folds into the
 // query row before the Q.K dots; the K bias adds a constant to every key score
 // that softmax cancels, so it is skipped; the V bias is token-independent and
 // is applied once in the merge pass after the final normalize.
+{% set BW = "" %}
+{% set BC = "" %}
 fn load_bias4(base: u32, d4: u32) -> vec4<f32> {
   let offset = base + d4 * 4u;
+  return vec4<f32>({{ BW }}bias[offset]{{ BC }}, {{ BW }}bias[offset + 1u]{{ BC }}, {{ BW }}bias[offset + 2u]{{ BC }}, {{ BW }}bias[offset + 3u]{{ BC }});
 }
 {% endif %}
   let tid = lid.x;
   let hKv = h / (Q_HEADS / KV_HEADS);
   let cacheSeq = params.kvSeq;
   let kvSeq = cacheSeq;
   // Query row (decode uses token zero; short-query prefill folds the token into wg.x).
 {% if splitQueries %}
     splitEnd = kvSeq;
   }
   for (var d4 = tid; d4 < HEAD_DIM_V4; d4 = d4 + WG) {
     var qv = vec4<f32>(query[qBaseV4 + d4]);
 {% if hasBias %}
       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);
       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;
       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 %}    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.
     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
     loop {
       if (jj >= tileCount) { break; }
       vacc = vacc + probs[jj] * load_value4(
+        kvBaseV4 + (kjBase + jj) * kvTokenStrideV4 + d4v      );
       jj = jj + J_GROUPS;
     }
     vacc_sh[tid] = vacc;
       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;
     }

build/webgpu/attn-flash-online.wgsl.jinja CHANGED Viewed

@@ -28,8 +28,8 @@ const Q_HIDDEN_V4: u32 = {{ qHiddenV4 }}u;
 const KV_HIDDEN_V4: u32 = {{ kvHiddenV4 }}u;
 const Q_HEADS: u32 = {{ qNumHeads }}u;
 const KV_HEADS: u32 = {{ kvNumHeads }}u;
-{% set qHeads = "params.qHeads" if headsFromParams else "Q_HEADS" %}
-{% set kvHeads = "params.kvHeads" if headsFromParams else "KV_HEADS" %}
 const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
 // FLT_MAX, not -inf, as the online (m, d) accumulator init: merges must keep
 // `m - m` finite so an empty lane / all--inf row contributes the exact
@@ -50,6 +50,7 @@ fn shifted_value(value: f32, maxValue: f32) -> f32 {
   let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
   return select(value - maxValue, 0.0, equalFiniteMax);
 }
 fn exp_shift(value: f32, maxValue: f32) -> f32 {
   return exp(shifted_value(value, maxValue));
 }
@@ -63,7 +64,6 @@ var<workgroup> probs: array<f32, WG>;
 // Both the subgroup and portable barrier-tree engines return the same merged
 // pair to every invocation. Repeated merges require a workgroup barrier between
 // calls before their shared partial storage is reused.
-{% set combineSubgroups = combineSubgroups is defined and combineSubgroups %}
 {% if combineSubgroups %}
 // Cross-subgroup merge that assumes nothing about which invocations share a
 // subgroup or how many subgroups there are: each subgroup's elected lane
@@ -116,41 +116,9 @@ fn combine_partials(m: f32, d: f32, lidx: u32, sgSize: u32) -> vec2<f32> {
   return combinedMD;
 }
 {% else %}
-{% set mdStreamed = mdStreams is defined %}
-{% set mdStreams = mdStreams if mdStreams is defined else 1 %}
-{% set mdExtent = "WG" if mdStreams == 1 else "WG * " ~ mdStreams ~ "u" %}
 var<workgroup> partialM: array<f32, {{ mdExtent }}>;
 var<workgroup> partialD: array<f32, {{ mdExtent }}>;
-{% if mdStreamed %}
-// In-place fold of {{ mdStreams }} streams. Input partials occupy
-// partialM/partialD; stream s returns its merged pair in slot s * WG.
-fn combine_partials_streams(lidx: u32) {
-  workgroupBarrier();
-  var stride = WG / 2u;
-  loop {
-    if (stride == 0u) {
-      break;
-    }
-    if (lidx < stride) {
-{% for s in range(mdStreams) %}
-      {
-        let slot = {{ s }}u * WG + lidx;
-        let m1 = partialM[slot];
-        let d1 = partialD[slot];
-        let m2 = partialM[slot + stride];
-        let d2 = partialD[slot + stride];
-        let mNew = max(m1, m2);
-        partialD[slot] = d1 * exp_shift(m1, mNew) + d2 * exp_shift(m2, mNew);
-        partialM[slot] = mNew;
-      }
-{% endfor %}
-    }
-    workgroupBarrier();
-    stride = stride / 2u;
-  }
-}
-{% else %}
 fn combine_partials(m: f32, d: f32, lidx: u32) -> vec2<f32> {
   partialM[lidx] = m;
@@ -180,22 +148,20 @@ fn combine_partials(m: f32, d: f32, lidx: u32) -> vec2<f32> {
   return merged;
 }
 {% endif %}
-{% endif %}
 // An explicit-zero specialization bakes the scale as 0. Otherwise,
 // params.scale == 0 encodes an omitted scale and selects 1/sqrt(headDim).
-{% if ATTN_SCALE_DIM is not defined %}{% set ATTN_SCALE_DIM = "HEAD_DIM" %}{% endif %}
-fn scale_value() -> f32 {
   if (params.scale != 0.0) { return params.scale; }
   return inverseSqrt(f32({{ ATTN_SCALE_DIM }}));
 }
 {% if hasBias %}
 fn load_bias4(base: u32, d4: u32) -> vec4<f32> {
   let offset = base + d4 * 4u;
-  return vec4<f32>(bias[offset], bias[offset + 1u], bias[offset + 2u], bias[offset + 3u]);
 }
 {% endif %}
@@ -221,7 +187,6 @@ fn main(
   let kvTokenStrideV4 = KV_HIDDEN_V4;
   // Cooperative vec4 Q-row load; init the output accumulator.
-{% set hasBias = hasBias is defined and hasBias %}
   for (var d4 = tid; d4 < HEAD_DIM_V4; d4 = d4 + WG) {
     var qv = vec4<f32>(query[qBaseV4 + d4]);
 {% if hasBias %}
@@ -297,7 +262,6 @@ fn main(
     probs[tid] = prob;
     workgroupBarrier();
     // running_out[d4] is owned by the same thread (tid ≡ d4 mod WG) across all
     // tiles, so the rescale-and-accumulate below is race-free; vec4 V loads.
     let tileCount = min(WG, keyBound - kjBase);

 const KV_HIDDEN_V4: u32 = {{ kvHiddenV4 }}u;
 const Q_HEADS: u32 = {{ qNumHeads }}u;
 const KV_HEADS: u32 = {{ kvNumHeads }}u;
+{% set qHeads = "Q_HEADS" %}
+{% set kvHeads = "KV_HEADS" %}
 const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
 // FLT_MAX, not -inf, as the online (m, d) accumulator init: merges must keep
 // `m - m` finite so an empty lane / all--inf row contributes the exact
   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));
 }
 // 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.
 {% 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
   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;
   return merged;
 }
 {% endif %}
 // An explicit-zero specialization bakes the scale as 0. Otherwise,
 // params.scale == 0 encodes an omitted scale and selects 1/sqrt(headDim).
+{% set ATTN_SCALE_DIM = "HEAD_DIM" %}fn scale_value() -> f32 {
   if (params.scale != 0.0) { return params.scale; }
   return inverseSqrt(f32({{ ATTN_SCALE_DIM }}));
 }
 {% if hasBias %}
+{% set BW = "" %}
+{% set BC = "" %}
 fn load_bias4(base: u32, d4: u32) -> vec4<f32> {
   let offset = base + d4 * 4u;
+  return vec4<f32>({{ BW }}bias[offset]{{ BC }}, {{ BW }}bias[offset + 1u]{{ BC }}, {{ BW }}bias[offset + 2u]{{ BC }}, {{ BW }}bias[offset + 3u]{{ BC }});
 }
 {% endif %}
   let kvTokenStrideV4 = KV_HIDDEN_V4;
   // Cooperative vec4 Q-row load; init the output accumulator.
   for (var d4 = tid; d4 < HEAD_DIM_V4; d4 = d4 + WG) {
     var qv = vec4<f32>(query[qBaseV4 + d4]);
 {% if hasBias %}
     probs[tid] = prob;
     workgroupBarrier();
     // running_out[d4] is owned by the same thread (tid ≡ d4 mod WG) across all
     // tiles, so the rescale-and-accumulate below is race-free; vec4 V loads.
     let tileCount = min(WG, keyBound - kjBase);

build/webgpu/attn-flash-prefill-cluster.wgsl.jinja CHANGED Viewed

@@ -1,32 +1,27 @@
-{% set sourceProfile = sourceProfile if sourceProfile is defined else 0 %}
-{% set scaling = scaling | default("0.0") %}
-{% set qkvStrideV4 = qkvStrideV4 | default(0) %}
-{% set QSEQ = "params.seq_len" if sourceProfile == 1 else "params.qSeq" %}
-{% set KVSEQ = "(params.past_len + params.seq_len)" if sourceProfile == 1 else "params.kvSeq" %}
-{% set IS_CAUSAL = "1u" if sourceProfile == 1 else "params.isCausal" %}
-{% set Q_STRIDE = "QKV_STRIDE_V4" if sourceProfile == 1 else "Q_HIDDEN_V4" %}
-{% set QUERY = "qkv" if sourceProfile == 1 else "query" %}
-{% set KEY = "cache_keys" if sourceProfile == 1 else "key" %}
-{% set VALUE = "cache_values" if sourceProfile == 1 else "value" %}
-{% set OUTPUT = "attn_out" if sourceProfile == 1 else "output" %}
-{% if sourceProfile == 1 %}{% set ATTN_SCALE_OVERRIDE = scaling %}{% endif %}
 {% if useSubgroups is not defined %}{% set useSubgroups = true %}{% endif %}
 {% if batchNoSgReduction is not defined %}{% set batchNoSgReduction = false %}{% endif %}
 {% if hasMask is not defined %}{% set hasMask = false %}{% endif %}
-{% if maskIsKeyKeep is not defined %}{% set maskIsKeyKeep = false %}{% endif %}
-{% if maskIsBool is not defined %}{% set maskIsBool = false %}{% endif %}
-{% if hasSoftcap is not defined %}{% set hasSoftcap = false %}{% endif %}
-{% if hasHeadSink is not defined %}{% set hasHeadSink = false %}{% endif %}
 {% set Q_HIDDEN = qHidden | default(0) %}
-{% if quantCacheFormat is not defined %}{% set quantCacheFormat = "" %}{% endif %}
-{% if useSeqlens is not defined %}{% set useSeqlens = false %}{% endif %}
 // A windowed cache binds a fixed CAPACITY but keeps only the most recent
 // min(total, capacity) rows resident. params.kvSeq then names the physical row
 // count, which is still the right batch stride but the wrong attention bound, so
 // the bounds read kvActive instead. Modes without seqlens use `params.kvSeq`
 // in both roles.
-{% set KVA = "kvActive" if useSeqlens else KVSEQ %}
-{% macro score_expr(part) %}{% if hasSoftcap %}params.softcap * tanh(clamp(({{ part }} * SCALE) / params.softcap, -30.0, 30.0)){% else %}{{ part }} * SCALE{% endif %}{% endmacro %}
 {% if useSubgroups %}
 enable subgroups;
 {% endif %}
@@ -36,15 +31,30 @@ enable subgroups;
 // workgroup storage; values widen to f32 when read. Other specializations stage
 // them as f32. Score and weighted-value accumulation remain in f32 throughout.
 {% set STAGE_MASK = hasMask and useSubgroups and stageMask is defined and stageMask %}
-{% set MASK_IS_INT = hasMask and (maskIsKeyKeep or maskIsBool) %}
-{% set MASK_TILE_TYPE = "u32" if MASK_IS_INT else "f32" %}
-{% set MASK_TILE_LOAD = "attn_mask[maskIndex]" if MASK_IS_INT else "f32(attn_mask[maskIndex])" %}
-{% set MASK_TILE_ZERO = "0u" if MASK_IS_INT else "0.0" %}
-{% set MASK_ELEMENT = "maskValue" if STAGE_MASK else "attn_mask[maskIndex]" %}
 {% set MASK_ADDITIVE = "maskValue" if STAGE_MASK else "f32(attn_mask[maskIndex])" %}
 {% set SLICE_COUNT = ((headDimV4 / LPQ) | int) %}
 {% set QPL = (2 if (not usesF16 and SLICE_COUNT <= 4 and TILE_K <= 16 and TILE_Q % 2 == 0) else 1) if useSubgroups else 1 %}
 {% macro qn(name, qi) %}{{ name }}{% if qi > 0 %}_q{{ qi }}{% endif %}{% endmacro %}
 // Tiled flash prefill attention with configurable-width query clusters for
 // token-major [batch, seq, heads*headDim] attention ops. Each workgroup covers
@@ -64,8 +74,7 @@ enable subgroups;
 const HEAD_DIM: u32 = {{ headDim }}u;
 const HEAD_DIM_V4: u32 = {{ headDimV4 }}u;
 const Q_HIDDEN_V4: u32 = {{ qHiddenV4 }}u;   // bsh token-major hidden stride
-{% if sourceProfile == 1 %}const QKV_STRIDE_V4: u32 = {{ qkvStrideV4 }}u; // packed [Q; K; V] input row stride
-{% endif %}const KV_HIDDEN_V4: u32 = {{ kvHiddenV4 }}u;
 const Q_HEADS: u32 = {{ qNumHeads }}u;
 const KV_HEADS: u32 = {{ kvNumHeads }}u;
 const TILE_Q: u32 = {{ TILE_Q }}u;
@@ -79,24 +88,10 @@ const WG: u32 = (TILE_Q / QPL) * LPQ;
 {% else %}
 const WG: u32 = TILE_Q * LPQ;
 {% endif %}
-{% if MASK_IS_INT %}
-// Key-keep masks in contrib attention use a finite low logit for a rejected
-// key. Logical ONNX bool masks use the exclusion sentinel instead, so a fully
-// masked row has zero mass. Keeping one declaration shape for both mask modes
-// lets both mask modes share the score loop below.
 const NEG_INF: f32 = -3.4028234663852886e38;
-const MASK_NEG: f32 = {{ "-1e38" if maskIsKeyKeep else "-3.4028234663852886e38" }};
-{% else %}
-const NEG_INF: f32 = -3.4028234663852886e38;
-{% endif %}
 var<workgroup> k_tile: array<vec4<{{ ST }}>, TILE_K * (HEAD_DIM / 4u)>;
 var<workgroup> v_tile: array<vec4<{{ ST }}>, TILE_K * (HEAD_DIM / 4u)>;
-{% if STAGE_MASK %}
-// Each LPQ cluster consumes one mask value per (query,key), so stage the
-// TILE_Q x TILE_K mask tile once instead of issuing LPQ duplicate global loads.
-var<workgroup> mask_tile: array<{{ MASK_TILE_TYPE }}, TILE_Q * TILE_K>;
-{% endif %}
 {% if not useSubgroups %}
 {% if batchNoSgReduction %}
 // No-subgroups cluster reduction scratch for a whole K tile. Staging every
@@ -109,62 +104,17 @@ var<workgroup> red: array<f32, WG>;
 {% endif %}
 {% endif %}
-{% if ATTN_SCALE_DIM is not defined %}{% set ATTN_SCALE_DIM = "HEAD_DIM" %}{% endif %}
-fn scale_value() -> f32 {
-{% if ATTN_SCALE_OVERRIDE is defined %}
-  return {{ ATTN_SCALE_OVERRIDE }};
-{% else %}
   if (params.scale != 0.0) { return params.scale; }
   return inverseSqrt(f32({{ ATTN_SCALE_DIM }}));
-{% endif %}
-}
-{% if quantCacheFormat %}
-// A quantized cache is dequantized once per key into the staged tile, then read
-// by all TILE_Q queries in the workgroup. The unpack cost is amortized over the
-// tile height instead of paid once per (query, key).
-{% macro emit_quant_scale4(kind, scaleBuffer) %}
-fn {{ kind }}scale4(d4: u32, hk: u32) -> vec4<f32> {
-  if (params.perChannel == 0u) {
-    return vec4<f32>({{ scaleBuffer }}[0]);
-  }
-  let base = hk * HEAD_DIM + d4 * 4u;
-  return vec4<f32>(
-    {{ scaleBuffer }}[base],
-    {{ scaleBuffer }}[base + 1u],
-    {{ scaleBuffer }}[base + 2u],
-    {{ scaleBuffer }}[base + 3u]
-  );
-}
-{%- endmacro %}
-{%- macro emit_quant_load4(format, kind, buffer, scaleBuffer) %}
-{{ emit_quant_scale4(kind, scaleBuffer) }}
-fn load_{{ kind }}4(indexV4: u32, d4: u32, hk: u32) -> vec4<f32> {
-{%- if format == "int8" %}
-  return vec4<f32>({{ buffer }}[indexV4]) * {{ kind }}scale4(d4, hk);
-{%- else %}
-  // Two elements cover this vec4: each carries two +8-biased nibbles, low first.
-  let rowBase = indexV4 - d4;
-  let lo = {{ buffer }}[rowBase + d4 * 2u];
-  let hi = {{ buffer }}[rowBase + d4 * 2u + 1u];
-  let nibbles = vec4<i32>(
-    i32(lo & 0xFu), i32((lo >> 4u) & 0xFu),
-    i32(hi & 0xFu), i32((hi >> 4u) & 0xFu)
-  );
-  let signed = nibbles - vec4<i32>(8);
-  return vec4<f32>(signed) * {{ kind }}scale4(d4, hk);
-{%- endif %}
 }
-{%- endmacro %}
-{{ emit_quant_load4(quantCacheFormat, "key", KEY, "k_scale") }}
-{{ emit_quant_load4(quantCacheFormat, "value", VALUE, "v_scale") }}
-{% endif %}
 {% if hasBias %}
 fn load_bias4(base: u32, d4: u32) -> vec4<f32> {
   let offset = base + d4 * 4u;
-  return vec4<f32>(bias[offset], bias[offset + 1u], bias[offset + 2u], bias[offset + 3u]);
 }
 {% endif %}
@@ -212,17 +162,6 @@ fn main(
   var {{ qn("m", qi) }}: f32 = NEG_INF;
   var {{ qn("l", qi) }}: f32 = 0.0;
 {% endfor %}
-{% if useSeqlens %}
-  // Resident rows of a windowed cache: the survivors were shifted down to [0, kvActive),
-  // so query/key DISTANCE is unchanged and every bound below reads as if the cache were
-  // exactly kvActive long. Rotary is excluded from this path (it would need the absolute
-  // position, not the cache-relative one), so pastLenForRope keeps the physical length.
-  // Clamping below by the query count matches the cache-update passes and the scalar
-  // path: a right-padded batch (seqlens_k[b]+1 < qSeq) still appends its whole chunk,
-  // so its queries score against all of it.
-  let kvActive = min({{ KVSEQ }}, max({{ QSEQ }}, u32(seqlens_k[b]) + 1u));
-{% endif %}
   // Causal ceiling per query; the key loop runs over the workgroup's union range
   // (uniform trip count), masking out-of-range (query, key) pairs.
   // Upper-left causal and/or sliding-window bounds. Query qIdx sits at
@@ -260,33 +199,13 @@ fn main(
       let kj = kStart + slot;
       if (kj < wgEnd) {
         let base4 = kvBatch4 + kj * KV_HIDDEN_V4 + d4;
-{% if quantCacheFormat %}
-        k_tile[i] = vec4<{{ ST }}>(load_key4(base4, d4, hKv));
-        v_tile[i] = vec4<{{ ST }}>(load_value4(base4, d4, hKv));
-{% else %}
         k_tile[i] = vec4<{{ ST }}>({{ KEY }}[base4]);
         v_tile[i] = vec4<{{ ST }}>({{ VALUE }}[base4]);
-{% endif %}
       } else {
         k_tile[i] = vec4<{{ ST }}>(0.0);
         v_tile[i] = vec4<{{ ST }}>(0.0);
       }
     }
-{% if STAGE_MASK %}
-    // The K/V load barrier also publishes this compact mask tile.
-    for (var i: u32 = tid; i < TILE_Q * TILE_K; i = i + WG) {
-      let qSlot = i / TILE_K;
-      let kSlot = i % TILE_K;
-      let maskQ = min(wg.x * TILE_Q + qSlot, {{ QSEQ }} - 1u);
-      let maskK = kStart + kSlot;
-      if (maskK < wgEnd) {
-        let maskIndex = b * params.maskBatchStride + h * params.maskHeadStride + maskQ * params.maskSeqStride + maskK;
-        mask_tile[i] = {{ MASK_TILE_LOAD }};
-      } else {
-        mask_tile[i] = {{ MASK_TILE_ZERO }};
-      }
-    }
-{% endif %}
     workgroupBarrier();
     // TILE_K is a small shader constant; the loop updates the named q/o slices in place.
 {% if not useSubgroups and batchNoSgReduction %}
@@ -329,15 +248,7 @@ fn main(
       if (kj >= minKj && kj < maxKj) {
 {% if hasMask %}
         let maskIndex = b * params.maskBatchStride + h * params.maskHeadStride + qClamped * params.maskSeqStride + kj;
-{% if maskIsBool %}
-        if (attn_mask[maskIndex] != 0u) {
-          s[kk] = {{ score_expr("part") }};
-        } else {
-          s[kk] = MASK_NEG;
-        }
-{% else %}
         s[kk] = {{ score_expr("part") }} + f32(attn_mask[maskIndex]);
-{% endif %}
 {% else %}
         s[kk] = {{ score_expr("part") }};
 {% endif %}
@@ -374,29 +285,11 @@ fn main(
         let sc = partv[{{ qi }}];
       if (kj >= {{ qn("minKj", qi) }} && kj < {{ qn("maxKj", qi) }}) {
   {% if hasMask %}
-  {% if STAGE_MASK %}
-          let maskValue = mask_tile[(qSub + {{ qi }}u) * TILE_K + kk];
-  {% else %}
           // Broadcast-strided mask address (a 0 stride collapses that axis; rank-2
           // [q, k] masks set batch/head strides to 0). {{ qn("qClamped", qi) }} keeps the seq index
           // in-bounds for padding queries in the last tile (their output is dropped).
           let maskIndex = b * params.maskBatchStride + h * params.maskHeadStride + {{ qn("qClamped", qi) }} * params.maskSeqStride + kj;
-  {% endif %}
-  {% if maskIsKeyKeep %}
-          // A broadcast key mask uses 1 for a retained key and 0 for padding.
-          {{ qn("s", qi) }}[kk] = {{ score_expr("sc") }} + (1.0 - f32({{ MASK_ELEMENT }})) * MASK_NEG;
-  {% elif maskIsBool %}
-          // Logical bool: a rejected key contributes no softmax mass. Leaving
-          // the initialized NEG_INF sentinel in place makes a fully masked row
-          // land on the zero-denominator output guard below.
-          if ({{ MASK_ELEMENT }} != 0u) {
-            {{ qn("s", qi) }}[kk] = {{ score_expr("sc") }};
-          } else {
-            {{ qn("s", qi) }}[kk] = MASK_NEG;
-          }
-  {% else %}
           {{ qn("s", qi) }}[kk] = {{ score_expr("sc") }} + {{ MASK_ADDITIVE }};
-  {% endif %}
   {% else %}
           {{ qn("s", qi) }}[kk] = {{ score_expr("sc") }};
   {% endif %}
@@ -434,29 +327,11 @@ fn main(
 {% endif %}
       if (kj >= minKj && kj < maxKj) {
 {% if hasMask %}
-{% if STAGE_MASK %}
-        let maskValue = mask_tile[qSub * TILE_K + kk];
-{% else %}
         // Broadcast-strided mask address (a 0 stride collapses that axis; rank-2
         // [q, k] masks set batch/head strides to 0). qClamped keeps the seq index
         // in-bounds for padding queries in the last tile (their output is dropped).
         let maskIndex = b * params.maskBatchStride + h * params.maskHeadStride + qClamped * params.maskSeqStride + kj;
-{% endif %}
-{% if maskIsKeyKeep %}
-        // A broadcast key mask uses 1 for a retained key and 0 for padding.
-        s[kk] = {{ score_expr("part") }} + (1.0 - f32({{ MASK_ELEMENT }})) * MASK_NEG;
-{% elif maskIsBool %}
-        // Logical bool: a rejected key contributes no softmax mass. Leaving
-        // the initialized NEG_INF sentinel in place makes a fully masked row
-        // land on the zero-denominator output guard below.
-        if ({{ MASK_ELEMENT }} != 0u) {
-          s[kk] = {{ score_expr("part") }};
-        } else {
-          s[kk] = MASK_NEG;
-        }
-{% else %}
         s[kk] = {{ score_expr("part") }} + {{ MASK_ADDITIVE }};
-{% endif %}
 {% else %}
         s[kk] = {{ score_expr("part") }};
 {% endif %}
@@ -464,24 +339,7 @@ fn main(
     }
 {% endif %}
-    // Per-thread online softmax over the tile. s[kk] is reused to hold the
-    // exponentiated probabilities for the PV accumulation below.
-{% for qi in range(QPL) %}
-    var {{ qn("tileMax", qi) }}: f32 = {{ qn("s", qi) }}[0];
-    for (var kk: u32 = 1u; kk < TILE_K; kk = kk + 1u) {
-      {{ qn("tileMax", qi) }} = max({{ qn("tileMax", qi) }}, {{ qn("s", qi) }}[kk]);
-    }
-    let {{ qn("newMax", qi) }} = max({{ qn("m", qi) }}, {{ qn("tileMax", qi) }});
-    let {{ qn("corr", qi) }} = select(exp({{ qn("m", qi) }} - {{ qn("newMax", qi) }}), 0.0, {{ qn("m", qi) }} == NEG_INF);
-    var {{ qn("pSum", qi) }}: f32 = 0.0;
-    for (var kk: u32 = 0u; kk < TILE_K; kk = kk + 1u) {
-      let pk = select(0.0, exp({{ qn("s", qi) }}[kk] - {{ qn("newMax", qi) }}), {{ qn("s", qi) }}[kk] != NEG_INF);
-      {{ qn("s", qi) }}[kk] = pk;
-      {{ qn("pSum", qi) }} = {{ qn("pSum", qi) }} + pk;
-    }
-    {{ qn("l", qi) }} = {{ qn("l", qi) }} * {{ qn("corr", qi) }} + {{ qn("pSum", qi) }};
-    {{ qn("m", qi) }} = {{ qn("newMax", qi) }};
-{% endfor %}
     // A boundary tile can address V rows outside a query's attended range, and
     // a dynamic-row producer may leave those rows unwritten. Because 0 * NaN is
     // NaN, the guarded loop selects the V operand away for range-excluded keys.
@@ -540,22 +398,11 @@ fn main(
 {% for qi in range(QPL) %}
   if ({{ qn("qValid", qi) }}) {
     let outBase4 = (b * {{ QSEQ }} + {{ qn("qIdx", qi) }}) * Q_HIDDEN_V4 + h * HEAD_DIM_V4 + lane8 * SLICE;
-{% if hasHeadSink %}
-    // The head sink is a learned logit that competes with the keys but carries
-    // no value, so it enters the denominator only and the weighted sum above is
-    // untouched. Renormalizing against max({{ qn("m", qi) }}, sink) keeps the exponentials in
-    // range when the sink dominates a fully-masked row.
-    let sink = f32(head_sink[h]);
-    let finalM = max({{ qn("m", qi) }}, sink);
-    let accScale = exp({{ qn("m", qi) }} - finalM);
-    let inv = accScale / (exp(sink - finalM) + {{ qn("l", qi) }} * accScale);
-{% else %}
     // {{ qn("l", qi) }} == 0 means this query had no probability-bearing key: either its
     // causal/window range is empty or its logical bool mask rejects every key.
     // Emit 0 rather than 0/0. Other attention paths enforce the same empty-row
     // contract by selecting on positive global mass.
     let inv = select(0.0, 1.0 / {{ qn("l", qi) }}, {{ qn("l", qi) }} > 0.0);
-{% endif %}
 {% for c in range(SLICE_COUNT) %}
     {{ OUTPUT }}[outBase4 + {{ c }}u] = vec4<{{ scalar }}>({{ attention_value(c, qi) }});
 {% endfor %}

+{% set QSEQ = "params.qSeq" %}
+{% set KVSEQ = "params.kvSeq" %}
+{% set IS_CAUSAL = "params.isCausal" %}
+{% set Q_STRIDE = "Q_HIDDEN_V4" %}
+{% set QUERY = "query" %}
+{% set KEY = "key" %}
+{% set VALUE = "value" %}
+{% set OUTPUT = "output" %}
 {% if useSubgroups is not defined %}{% set useSubgroups = true %}{% endif %}
 {% if batchNoSgReduction is not defined %}{% set batchNoSgReduction = false %}{% endif %}
 {% if hasMask is not defined %}{% set hasMask = false %}{% endif %}
 {% set Q_HIDDEN = qHidden | default(0) %}
 // A windowed cache binds a fixed CAPACITY but keeps only the most recent
 // min(total, capacity) rows resident. params.kvSeq then names the physical row
 // count, which is still the right batch stride but the wrong attention bound, so
 // the bounds read kvActive instead. Modes without seqlens use `params.kvSeq`
 // in both roles.
+{% if hasKeyLimit is not defined %}{% set hasKeyLimit = false %}{% endif %}
+{% if hasKeyLimit %}
+{% set KVA = "select(params.kvSeq, min(params.keyLimit, params.kvSeq), params.keyLimit > 0u)" %}
+{% else %}
+{% set KVA = KVSEQ %}
+{% endif %}
+{% macro score_expr(part) %}{{ part }} * SCALE{% endmacro %}
 {% if useSubgroups %}
 enable subgroups;
 {% endif %}
 // workgroup storage; values widen to f32 when read. Other specializations stage
 // them as f32. Score and weighted-value accumulation remain in f32 throughout.
 {% set STAGE_MASK = hasMask and useSubgroups and stageMask is defined and stageMask %}
 {% set MASK_ADDITIVE = "maskValue" if STAGE_MASK else "f32(attn_mask[maskIndex])" %}
 {% set SLICE_COUNT = ((headDimV4 / LPQ) | int) %}
 {% set QPL = (2 if (not usesF16 and SLICE_COUNT <= 4 and TILE_K <= 16 and TILE_Q % 2 == 0) else 1) if useSubgroups else 1 %}
 {% macro qn(name, qi) %}{{ name }}{% if qi > 0 %}_q{{ qi }}{% endif %}{% endmacro %}
+{% macro emit_tile_softmax() %}
+    // Per-thread online softmax over the tile. s[kk] is reused to hold the
+    // exponentiated probabilities for the PV accumulation below.
+{% for qi in range(QPL) %}
+    var {{ qn("tileMax", qi) }}: f32 = {{ qn("s", qi) }}[0];
+    for (var kk: u32 = 1u; kk < TILE_K; kk = kk + 1u) {
+      {{ qn("tileMax", qi) }} = max({{ qn("tileMax", qi) }}, {{ qn("s", qi) }}[kk]);
+    }
+    let {{ qn("newMax", qi) }} = max({{ qn("m", qi) }}, {{ qn("tileMax", qi) }});
+    let {{ qn("corr", qi) }} = select(exp({{ qn("m", qi) }} - {{ qn("newMax", qi) }}), 0.0, {{ qn("m", qi) }} == NEG_INF);
+    var {{ qn("pSum", qi) }}: f32 = 0.0;
+    for (var kk: u32 = 0u; kk < TILE_K; kk = kk + 1u) {
+      let pk = select(0.0, exp({{ qn("s", qi) }}[kk] - {{ qn("newMax", qi) }}), {{ qn("s", qi) }}[kk] != NEG_INF);
+      {{ qn("s", qi) }}[kk] = pk;
+      {{ qn("pSum", qi) }} = {{ qn("pSum", qi) }} + pk;
+    }
+    {{ qn("l", qi) }} = {{ qn("l", qi) }} * {{ qn("corr", qi) }} + {{ qn("pSum", qi) }};
+    {{ qn("m", qi) }} = {{ qn("newMax", qi) }};
+{% endfor %}
+{% endmacro %}
 // Tiled flash prefill attention with configurable-width query clusters for
 // token-major [batch, seq, heads*headDim] attention ops. Each workgroup covers
 const HEAD_DIM: u32 = {{ headDim }}u;
 const HEAD_DIM_V4: u32 = {{ headDimV4 }}u;
 const Q_HIDDEN_V4: u32 = {{ qHiddenV4 }}u;   // bsh token-major hidden stride
+const KV_HIDDEN_V4: u32 = {{ kvHiddenV4 }}u;
 const Q_HEADS: u32 = {{ qNumHeads }}u;
 const KV_HEADS: u32 = {{ kvNumHeads }}u;
 const TILE_Q: u32 = {{ TILE_Q }}u;
 {% else %}
 const WG: u32 = TILE_Q * LPQ;
 {% endif %}
 const NEG_INF: f32 = -3.4028234663852886e38;
 var<workgroup> k_tile: array<vec4<{{ ST }}>, TILE_K * (HEAD_DIM / 4u)>;
 var<workgroup> v_tile: array<vec4<{{ ST }}>, TILE_K * (HEAD_DIM / 4u)>;
 {% if not useSubgroups %}
 {% if batchNoSgReduction %}
 // No-subgroups cluster reduction scratch for a whole K tile. Staging every
 {% endif %}
 {% 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 }}));
 }
 {% if hasBias %}
+{% set BW = "" %}
+{% set BC = "" %}
 fn load_bias4(base: u32, d4: u32) -> vec4<f32> {
   let offset = base + d4 * 4u;
+  return vec4<f32>({{ BW }}bias[offset]{{ BC }}, {{ BW }}bias[offset + 1u]{{ BC }}, {{ BW }}bias[offset + 2u]{{ BC }}, {{ BW }}bias[offset + 3u]{{ BC }});
 }
 {% endif %}
   var {{ qn("m", qi) }}: f32 = NEG_INF;
   var {{ qn("l", qi) }}: f32 = 0.0;
 {% endfor %}
   // Causal ceiling per query; the key loop runs over the workgroup's union range
   // (uniform trip count), masking out-of-range (query, key) pairs.
   // Upper-left causal and/or sliding-window bounds. Query qIdx sits at
       let kj = kStart + slot;
       if (kj < wgEnd) {
         let base4 = kvBatch4 + kj * KV_HIDDEN_V4 + d4;
         k_tile[i] = vec4<{{ ST }}>({{ KEY }}[base4]);
         v_tile[i] = vec4<{{ ST }}>({{ VALUE }}[base4]);
       } else {
         k_tile[i] = vec4<{{ ST }}>(0.0);
         v_tile[i] = vec4<{{ ST }}>(0.0);
       }
     }
     workgroupBarrier();
     // TILE_K is a small shader constant; the loop updates the named q/o slices in place.
 {% if not useSubgroups and batchNoSgReduction %}
       if (kj >= minKj && kj < maxKj) {
 {% if hasMask %}
         let maskIndex = b * params.maskBatchStride + h * params.maskHeadStride + qClamped * params.maskSeqStride + kj;
         s[kk] = {{ score_expr("part") }} + f32(attn_mask[maskIndex]);
 {% else %}
         s[kk] = {{ score_expr("part") }};
 {% endif %}
         let sc = partv[{{ qi }}];
       if (kj >= {{ qn("minKj", qi) }} && kj < {{ qn("maxKj", qi) }}) {
   {% if hasMask %}
           // Broadcast-strided mask address (a 0 stride collapses that axis; rank-2
           // [q, k] masks set batch/head strides to 0). {{ qn("qClamped", qi) }} keeps the seq index
           // in-bounds for padding queries in the last tile (their output is dropped).
           let maskIndex = b * params.maskBatchStride + h * params.maskHeadStride + {{ qn("qClamped", qi) }} * params.maskSeqStride + kj;
           {{ qn("s", qi) }}[kk] = {{ score_expr("sc") }} + {{ MASK_ADDITIVE }};
   {% else %}
           {{ qn("s", qi) }}[kk] = {{ score_expr("sc") }};
   {% endif %}
 {% endif %}
       if (kj >= minKj && kj < maxKj) {
 {% if hasMask %}
         // Broadcast-strided mask address (a 0 stride collapses that axis; rank-2
         // [q, k] masks set batch/head strides to 0). qClamped keeps the seq index
         // in-bounds for padding queries in the last tile (their output is dropped).
         let maskIndex = b * params.maskBatchStride + h * params.maskHeadStride + qClamped * params.maskSeqStride + kj;
         s[kk] = {{ score_expr("part") }} + {{ MASK_ADDITIVE }};
 {% else %}
         s[kk] = {{ score_expr("part") }};
 {% endif %}
     }
 {% endif %}
+{{ emit_tile_softmax() }}
     // A boundary tile can address V rows outside a query's attended range, and
     // a dynamic-row producer may leave those rows unwritten. Because 0 * NaN is
     // NaN, the guarded loop selects the V operand away for range-excluded keys.
 {% for qi in range(QPL) %}
   if ({{ qn("qValid", qi) }}) {
     let outBase4 = (b * {{ QSEQ }} + {{ qn("qIdx", qi) }}) * Q_HIDDEN_V4 + h * HEAD_DIM_V4 + lane8 * SLICE;
     // {{ qn("l", qi) }} == 0 means this query had no probability-bearing key: either its
     // causal/window range is empty or its logical bool mask rejects every key.
     // Emit 0 rather than 0/0. Other attention paths enforce the same empty-row
     // contract by selecting on positive global mass.
     let inv = select(0.0, 1.0 / {{ qn("l", qi) }}, {{ qn("l", qi) }} > 0.0);
 {% for c in range(SLICE_COUNT) %}
     {{ OUTPUT }}[outBase4 + {{ c }}u] = vec4<{{ scalar }}>({{ attention_value(c, qi) }});
 {% endfor %}

build/webgpu/attn-flash-q32-broadcast.wgsl.jinja CHANGED Viewed

@@ -1,4 +1,5 @@
-// Register-resident flash prefill uses one 32-lane subgroup per workgroup. Each
 // lane owns a query row and keeps its Q slice and f32 output in registers. Lanes
 // cooperatively load K/V and broadcast them with subgroupShuffle, so each query
 // computes q·k and p·v without cross-lane reductions or workgroup storage.
@@ -10,11 +11,16 @@
 {% set ST = "f16" if usesF16 else "f32" %}
 {% set components = ["x", "y", "z", "w"] %}
 {% set USE_SUBGROUPS = useSubgroups if useSubgroups is defined else true %}
 {% set Q_STEP = qStep if qStep is defined else 32 %}
 {% set CAUSAL = false if (hasCausal is defined and not hasCausal) else true %}
 {% if USE_SUBGROUPS %}
 enable subgroups;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
 const HEAD_DIM_V4: u32 = {{ headDimV4 }}u;
@@ -46,15 +52,16 @@ fn scale_value() -> f32 {
   return inverseSqrt(f32({{ ATTN_SCALE_DIM }}));
 }
 {% if hasBias %}
 fn load_bias4(base: u32, d4: u32) -> vec4<f32> {
   let offset = base + d4 * 4u;
-  return vec4<f32>(bias[offset], bias[offset + 1u], bias[offset + 2u], bias[offset + 3u]);
 }
 {% endif %}
-@compute @workgroup_size({{ Q_STEP }}, 1, 1)
 fn main(
   @builtin(workgroup_id) wg: vec3<u32>,
   @builtin(local_invocation_id) lid: vec3<u32>{% if USE_SUBGROUPS %},
@@ -99,9 +106,14 @@ fn main(
   // so every lane shares a uniform trip count (subgroup ops stay reconverged);
   // each lane masks its own keys past myMaxKj to NEG_INF.
 {% if CAUSAL %}
   let lastQ = min(wg.x * Q_STEP + Q_STEP - 1u, params.qSeq - 1u);
-  let kvEnd = select(params.kvSeq, min(lastQ + 1u, params.kvSeq), params.isCausal != 0u);
-  let myMaxKj = select(params.kvSeq, min(qi + 1u, params.kvSeq), params.isCausal != 0u);
 {% else %}
   let kvEnd = params.kvSeq;
   let myMaxKj = params.kvSeq;
@@ -193,6 +205,13 @@ fn main(
 {% endfor %}
     previous_max = new_max;
     previous_denom = denom;
     for (var d4: u32 = 0u; d4 < HEAD_DIM_V4; d4 = d4 + 1u) {
 {% if USE_SUBGROUPS %}
@@ -209,21 +228,37 @@ fn main(
       }
 {% endif %}
 {% endif %}
       var acc: vec4<f32> = vec4<f32>(0.0);
 {% for g in range(qkGroups) %}
 {% for lane in range(4) %}
-{% if USE_SUBGROUPS %}
-{% if g < 8 %}
-      acc = acc + vec4<f32>(subgroupShuffle(v_local0, {{ g * 4 + lane }}u)) * qk{{ g }}.{{ components[lane] }};
 {% else %}
-      acc = acc + vec4<f32>(subgroupShuffle(v_local1, {{ (g - 8) * 4 + lane }}u)) * qk{{ g }}.{{ components[lane] }};
 {% endif %}
 {% else %}
-      acc = acc + vec4<f32>(valueTile[d4 * K_STEP + {{ g * 4 + lane }}u]) * qk{{ g }}.{{ components[lane] }};
 {% endif %}
 {% endfor %}
 {% endfor %}
       o_tile[d4] = o_tile[d4] * vec4<f32>(o_ratio) + acc;
     }
 {% if not USE_SUBGROUPS %}
     workgroupBarrier();

+// Register-resident flash prefill uses one 32-lane subgroup per workgroup (fixed
+// by the adapter, or pinned where the adapter can compile exactly 32 lanes). Each
 // lane owns a query row and keeps its Q slice and f32 output in registers. Lanes
 // cooperatively load K/V and broadcast them with subgroupShuffle, so each query
 // computes q·k and p·v without cross-lane reductions or workgroup storage.
 {% set ST = "f16" if usesF16 else "f32" %}
 {% set components = ["x", "y", "z", "w"] %}
 {% set USE_SUBGROUPS = useSubgroups if useSubgroups is defined else true %}
+{% set PIN_SUBGROUP_32 = USE_SUBGROUPS and pinSubgroupSize32 is defined and pinSubgroupSize32 %}
 {% set Q_STEP = qStep if qStep is defined else 32 %}
 {% set CAUSAL = false if (hasCausal is defined and not hasCausal) else true %}
+{% set qHidden = qHidden | default(0) %}
 {% if USE_SUBGROUPS %}
 enable subgroups;
 {% endif %}
+{% if PIN_SUBGROUP_32 %}
+enable subgroup_size_control;
+{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 const HEAD_DIM_V4: u32 = {{ headDimV4 }}u;
   return inverseSqrt(f32({{ ATTN_SCALE_DIM }}));
 }
 {% if hasBias %}
+{% set BW = "" %}
+{% set BC = "" %}
 fn load_bias4(base: u32, d4: u32) -> vec4<f32> {
   let offset = base + d4 * 4u;
+  return vec4<f32>({{ BW }}bias[offset]{{ BC }}, {{ BW }}bias[offset + 1u]{{ BC }}, {{ BW }}bias[offset + 2u]{{ BC }}, {{ BW }}bias[offset + 3u]{{ BC }});
 }
 {% endif %}
+@compute @workgroup_size({{ Q_STEP }}, 1, 1){{ " @subgroup_size(32)" if PIN_SUBGROUP_32 else "" }}
 fn main(
   @builtin(workgroup_id) wg: vec3<u32>,
   @builtin(local_invocation_id) lid: vec3<u32>{% if USE_SUBGROUPS %},
   // so every lane shares a uniform trip count (subgroup ops stay reconverged);
   // each lane masks its own keys past myMaxKj to NEG_INF.
 {% if CAUSAL %}
+{% set FIXED_CAUSAL = false %}
+{% macro key_ceiling(query) %}
+{% if FIXED_CAUSAL %}
+min({{ query }} + 1u, params.kvSeq){% else %}
+select(params.kvSeq, min({{ query }} + 1u, params.kvSeq), params.isCausal != 0u){% endif %}{% endmacro %}
   let lastQ = min(wg.x * Q_STEP + Q_STEP - 1u, params.qSeq - 1u);
+  let kvEnd = {{ key_ceiling("lastQ") }};
+  let myMaxKj = {{ key_ceiling("qi") }};
 {% else %}
   let kvEnd = params.kvSeq;
   let myMaxKj = params.kvSeq;
 {% endfor %}
     previous_max = new_max;
     previous_denom = denom;
+{% set PV_HALF = ST == "f16" and not (precisePv | default(false)) %}
+{% set PV_FLUSH_GROUPS = 4 %}
+{% if PV_HALF %}
+{% for g in range(qkGroups) %}
+    let p{{ g }} = vec4<f16>(qk{{ g }});
+{% endfor %}
+{% endif %}
     for (var d4: u32 = 0u; d4 < HEAD_DIM_V4; d4 = d4 + 1u) {
 {% if USE_SUBGROUPS %}
       }
 {% endif %}
 {% endif %}
+{% if PV_HALF %}
+      var acc_f32: vec4<f32> = vec4<f32>(0.0);
+      var acc: vec4<f16> = vec4<f16>(0.0);
+{% else %}
       var acc: vec4<f32> = vec4<f32>(0.0);
+{% endif %}
 {% for g in range(qkGroups) %}
 {% for lane in range(4) %}
+{% if not USE_SUBGROUPS %}
+{% set pvValue = "valueTile[d4 * K_STEP + " ~ (g * 4 + lane) ~ "u]" %}
+{% elif g < 8 %}
+{% set pvValue = "subgroupShuffle(v_local0, " ~ (g * 4 + lane) ~ "u)" %}
 {% else %}
+{% set pvValue = "subgroupShuffle(v_local1, " ~ ((g - 8) * 4 + lane) ~ "u)" %}
 {% endif %}
+{% if PV_HALF %}
+      acc = fma({{ pvValue }}, vec4<f16>(p{{ g }}.{{ components[lane] }}), acc);
 {% else %}
+      acc = acc + vec4<f32>({{ pvValue }}) * qk{{ g }}.{{ components[lane] }};
 {% endif %}
 {% endfor %}
+{% if PV_HALF and (g + 1) % PV_FLUSH_GROUPS == 0 %}
+      acc_f32 = acc_f32 + vec4<f32>(acc);
+      acc = vec4<f16>(0.0);
+{% endif %}
 {% endfor %}
+{% if PV_HALF %}
+      o_tile[d4] = o_tile[d4] * vec4<f32>(o_ratio) + acc_f32;
+{% else %}
       o_tile[d4] = o_tile[d4] * vec4<f32>(o_ratio) + acc;
+{% endif %}
     }
 {% if not USE_SUBGROUPS %}
     workgroupBarrier();

build/webgpu/attn-materialized-apply-f32.wgsl.jinja CHANGED Viewed

@@ -1,3 +1,4 @@
 {{ env.wgsl.resourceDeclarations }}
 // Register-blocked normalized-scores @ V GEMM with BHSD and BSH specializations.
@@ -39,7 +40,7 @@ var<workgroup> rowMax: array<f32, BM>;
 {% endif %}
 fn v_at(b: u32, h: u32, k: u32, d: u32) -> f32 {
-  return value[(b * params.kvSeq + k) * KV_HIDDEN + h * HEAD_DIM + d];
 }
 fn store_y(b: u32, h: u32, q: u32, d: u32, result: f32) {
@@ -48,7 +49,7 @@ fn store_y(b: u32, h: u32, q: u32, d: u32, result: f32) {
 {% else %}
   let biased = result;
 {% endif %}
-  output[(b * params.qSeq + q) * Q_HIDDEN + h * HEAD_DIM + d] = biased;
 }
 @compute @workgroup_size(WG_DIM, WG_DIM, 1)

+{% set STORAGE_F16 = false %}
 {{ env.wgsl.resourceDeclarations }}
 // Register-blocked normalized-scores @ V GEMM with BHSD and BSH specializations.
 {% endif %}
 fn v_at(b: u32, h: u32, k: u32, d: u32) -> f32 {
+  return {{ "f32(" if STORAGE_F16 else "" }}value[(b * params.kvSeq + k) * KV_HIDDEN + h * HEAD_DIM + d]{{ ")" if STORAGE_F16 else "" }};
 }
 fn store_y(b: u32, h: u32, q: u32, d: u32, result: f32) {
 {% else %}
   let biased = result;
 {% endif %}
+  output[(b * params.qSeq + q) * Q_HIDDEN + h * HEAD_DIM + d] = {{ "f16(" if STORAGE_F16 else "" }}biased{{ ")" if STORAGE_F16 else "" }};
 }
 @compute @workgroup_size(WG_DIM, WG_DIM, 1)

build/webgpu/attn-materialized-rowstats-combine-f32.wgsl.jinja CHANGED Viewed

@@ -13,6 +13,12 @@
 // computing the denominator there would double the exp count. Fold maxima and
 // leave the denominator to the apply pass, which sees every row element anyway.
 {% set MAX_ONLY = maxOnly is defined and maxOnly %}
 {{ env.wgsl.resourceDeclarations }}
 const SLOTS: u32 = {{ statSlots }}u;
@@ -41,21 +47,19 @@ 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);
 }
-{%- endif %}
 {% if stableUsage == "all" %}
 fn exp_shift(value: f32, maxValue: f32) -> f32 {
   return exp(shifted_value(value, maxValue));
 }
-{%- endif %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(
   @builtin(global_invocation_id) gid: vec3<u32>
 ) {
-  let row = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
-  if (row >= params.rows) { return; }
   // `row` already runs over (batch, head, query) together, and the partial
   // layout puts that same product one axis out from the slot, so the stride

 // computing the denominator there would double the exp count. Fold maxima and
 // leave the denominator to the apply pass, which sees every row element anyway.
 {% set MAX_ONLY = maxOnly is defined and maxOnly %}
+{% 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 }}
 const SLOTS: u32 = {{ statSlots }}u;
   let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
   return select(value - maxValue, 0.0, equalFiniteMax);
 }
+{% endif %}
 {% if stableUsage == "all" %}
 fn exp_shift(value: f32, maxValue: f32) -> f32 {
   return exp(shifted_value(value, maxValue));
 }
+{% endif %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(
   @builtin(global_invocation_id) gid: vec3<u32>
 ) {
+{{ flat_index_2d("WG", "row", "params.rows", guardInline=true) }}
   // `row` already runs over (batch, head, query) together, and the partial
   // layout puts that same product one axis out from the slot, so the stride

build/webgpu/attn-materialized-score-f32.wgsl.jinja CHANGED Viewed

@@ -1,3 +1,4 @@
 {{ env.wgsl.resourceDeclarations }}
 // f32 prefill score GEMM with BHSD and packed BSH specializations. It computes
@@ -8,22 +9,14 @@
 // for TM * TN * 4 fused multiply-adds. K keeps a one-word shared stride
 // padding, which changes the transposed workgroup-memory stride and reduces
 // bank-conflict risk on banked implementations.
-{% if scalarAccumulators is not defined %}{% set scalarAccumulators = false %}{% endif %}
-{% if maskIsKeyKeep is not defined %}{% set maskIsKeyKeep = false %}{% endif %}
-{% if scalarAccumulators %}
-{% set TM_JINJA = (materializedQueryTile / materializedWorkgroupDim)|int %}
-{% set TN_JINJA = (materializedKeyTile / materializedWorkgroupDim)|int %}
-{% endif %}
 {% set EMIT_ROW_STATS = emitRowStats is defined and emitRowStats %}
 const Q_HEADS: u32 = {{ qNumHeads }}u;
 const KV_HEADS: u32 = {{ kvNumHeads }}u;
 const HEAD_DIM: u32 = {{ headDim }}u;
 const Q_HIDDEN: u32 = Q_HEADS * HEAD_DIM;
 const KV_HIDDEN: u32 = KV_HEADS * HEAD_DIM;
-{% if maskIsKeyKeep %}
-const NEG_INF: f32 = -3.4028234663852886e38;
-{% endif %}
 const BK: u32 = {{ materializedInnerTile }}u;
 const BM: u32 = {{ materializedQueryTile }}u;
@@ -45,46 +38,20 @@ const STAT_SLOTS: u32 = {{ statSlots }}u;
 var<workgroup> statScratch: array<f32, WG_THREADS>;
 {% endif %}
-{% if EMIT_ROW_STATS %}{% set stableHelperUsage = "constant" %}
-{% set stableUsage = stableHelperUsage if stableHelperUsage is defined else "all" %}
-// 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;
-{% if stableUsage != "constant" %}
-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);
-}
-{%- endif %}
-{% if stableUsage == "all" %}
-fn exp_shift(value: f32, maxValue: f32) -> f32 {
-  return exp(shifted_value(value, maxValue));
-}
-{%- endif %}
 {% endif %}
-{% if ATTN_SCALE_DIM is not defined %}{% set ATTN_SCALE_DIM = "HEAD_DIM" %}{% endif %}
-fn scale_value() -> f32 {
   if (params.scale != 0.0) { return params.scale; }
   return inverseSqrt(f32({{ ATTN_SCALE_DIM }}));
 }
 fn q_at(b: u32, h: u32, q: u32, d: u32) -> f32 {
-  var value = query[(b * params.qSeq + q) * Q_HIDDEN + h * HEAD_DIM + d];
 {% if hasBias %}
   value = value + bias[h * HEAD_DIM + d];
 {% endif %}
@@ -92,7 +59,7 @@ fn q_at(b: u32, h: u32, q: u32, d: u32) -> f32 {
 }
 fn k_at(b: u32, h: u32, k: u32, d: u32) -> f32 {
-  return key[(b * params.kvSeq + k) * KV_HIDDEN + h * HEAD_DIM + d];
 }
 @compute @workgroup_size(WG_DIM, WG_DIM, 1)
@@ -107,17 +74,8 @@ fn main(
   let kBase = wg.x * BN;
   let li = lid.y * WG_DIM + lid.x;
-{% if scalarAccumulators %}
-  // Keep every register-tile cell statically addressable.
-{% for i in range(TM_JINJA) %}
-{% for j in range(TN_JINJA) %}
-  var acc{{ i }}_{{ j }}: f32 = 0.0;
-{% endfor %}
-{% endfor %}
-{% else %}
   var acc: array<f32, TM * TN>;
   for (var i = 0u; i < TM * TN; i = i + 1u) { acc[i] = 0.0; }
-{% endif %}
   for (var dBase = 0u; dBase < HEAD_DIM; dBase = dBase + BK) {
     // Both source tiles are loaded in their physical row-major direction.
@@ -152,19 +110,6 @@ fn main(
     let qr0 = lid.y * TM;
     let kc0 = lid.x * TN;
     for (var dv = 0u; dv < D_VECS; dv = dv + 1u) {
-{% if scalarAccumulators %}
-{% for i in range(TM_JINJA) %}
-      let qv{{ i }} = tileQ[qr0 + {{ i }}u][dv];
-{% endfor %}
-{% for j in range(TN_JINJA) %}
-      let kv{{ j }} = tileK[kc0 + {{ j }}u][dv];
-{% endfor %}
-{% for i in range(TM_JINJA) %}
-{% for j in range(TN_JINJA) %}
-      acc{{ i }}_{{ j }} = acc{{ i }}_{{ j }} + dot(qv{{ i }}, kv{{ j }});
-{% endfor %}
-{% endfor %}
-{% else %}
       var qv: array<vec4<f32>, TM>;
       var kv: array<vec4<f32>, TN>;
       for (var i = 0u; i < TM; i = i + 1u) { qv[i] = tileQ[qr0 + i][dv]; }
@@ -174,7 +119,6 @@ fn main(
           acc[i * TN + j] = acc[i * TN + j] + dot(qv[i], kv[j]);
         }
       }
-{% endif %}
     }
     workgroupBarrier();
   }
@@ -189,28 +133,6 @@ fn main(
   var rowM: array<f32, TM>;
   for (var i = 0u; i < TM; i = i + 1u) { rowM[i] = -FLT_MAX; }
 {% endif %}
-{% if scalarAccumulators %}
-{% for i in range(TM_JINJA) %}
-  {
-    let qi = qi0 + {{ i }}u;
-    if (qi < params.qSeq) {
-{% for j in range(TN_JINJA) %}
-      let ki{{ j }} = ki0 + {{ j }}u;
-      if (ki{{ j }} < params.kvSeq) {
-        var score = acc{{ i }}_{{ j }} * scale;
-{% if maskIsKeyKeep %}
-        score = score + (1.0 - f32(attn_mask[b * params.kvSeq + ki{{ j }}])) * NEG_INF;
-{% endif %}
-        scores[scoreBase + qi * params.kvSeq + ki{{ j }}] = score;
-{% if EMIT_ROW_STATS %}
-        rowM[{{ i }}u] = max(rowM[{{ i }}u], score);
-{% endif %}
-      }
-{% endfor %}
-    }
-  }
-{% endfor %}
-{% else %}
   for (var i = 0u; i < TM; i = i + 1u) {
     let qi = qi0 + i;
     if (qi >= params.qSeq) { continue; }
@@ -225,7 +147,6 @@ fn main(
       }
     }
   }
-{% endif %}
 {% if EMIT_ROW_STATS %}
   // Fold each row's partial (m, d) across the WG_DIM threads that own its columns. Every lane

+{% set STORAGE_F16 = false %}
 {{ env.wgsl.resourceDeclarations }}
 // f32 prefill score GEMM with BHSD and packed BSH specializations. It computes
 // for TM * TN * 4 fused multiply-adds. K keeps a one-word shared stride
 // padding, which changes the transposed workgroup-memory stride and reduces
 // bank-conflict risk on banked implementations.
 {% set EMIT_ROW_STATS = emitRowStats is defined and emitRowStats %}
+{% set statSlots = statSlots | default(0) %}
 const Q_HEADS: u32 = {{ qNumHeads }}u;
 const KV_HEADS: u32 = {{ kvNumHeads }}u;
 const HEAD_DIM: u32 = {{ headDim }}u;
 const Q_HIDDEN: u32 = Q_HEADS * HEAD_DIM;
 const KV_HIDDEN: u32 = KV_HEADS * HEAD_DIM;
 const BK: u32 = {{ materializedInnerTile }}u;
 const BM: u32 = {{ materializedQueryTile }}u;
 var<workgroup> statScratch: array<f32, WG_THREADS>;
 {% endif %}
+{% if EMIT_ROW_STATS %}// 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;
 {% 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 q_at(b: u32, h: u32, q: u32, d: u32) -> f32 {
+  var value = {{ "f32(" if STORAGE_F16 else "" }}query[(b * params.qSeq + q) * Q_HIDDEN + h * HEAD_DIM + d]{{ ")" if STORAGE_F16 else "" }};
 {% if hasBias %}
   value = value + bias[h * HEAD_DIM + d];
 {% endif %}
 }
 fn k_at(b: u32, h: u32, k: u32, d: u32) -> f32 {
+  return {{ "f32(" if STORAGE_F16 else "" }}key[(b * params.kvSeq + k) * KV_HIDDEN + h * HEAD_DIM + d]{{ ")" if STORAGE_F16 else "" }};
 }
 @compute @workgroup_size(WG_DIM, WG_DIM, 1)
   let kBase = wg.x * BN;
   let li = lid.y * WG_DIM + lid.x;
   var acc: array<f32, TM * TN>;
   for (var i = 0u; i < TM * TN; i = i + 1u) { acc[i] = 0.0; }
   for (var dBase = 0u; dBase < HEAD_DIM; dBase = dBase + BK) {
     // Both source tiles are loaded in their physical row-major direction.
     let qr0 = lid.y * TM;
     let kc0 = lid.x * TN;
     for (var dv = 0u; dv < D_VECS; dv = dv + 1u) {
       var qv: array<vec4<f32>, TM>;
       var kv: array<vec4<f32>, TN>;
       for (var i = 0u; i < TM; i = i + 1u) { qv[i] = tileQ[qr0 + i][dv]; }
           acc[i * TN + j] = acc[i * TN + j] + dot(qv[i], kv[j]);
         }
       }
     }
     workgroupBarrier();
   }
   var rowM: array<f32, TM>;
   for (var i = 0u; i < TM; i = i + 1u) { rowM[i] = -FLT_MAX; }
 {% endif %}
   for (var i = 0u; i < TM; i = i + 1u) {
     let qi = qi0 + i;
     if (qi >= params.qSeq) { continue; }
       }
     }
   }
 {% if EMIT_ROW_STATS %}
   // Fold each row's partial (m, d) across the WG_DIM threads that own its columns. Every lane

build/webgpu/attn-materialized-sgmat-f32.wgsl.jinja CHANGED Viewed

@@ -1,7 +1,5 @@
 {% set MT = "f16" if (operandF16 is defined and operandF16) else "f32" %}
-{% if MT == "f16" %}
-enable f16;
-{% endif %}
 enable subgroups;
 {% if pinSubgroupSize32 %}
 enable subgroup_size_control;
@@ -9,7 +7,6 @@ enable subgroup_size_control;
 enable chromium_experimental_subgroup_matrix;
 diagnostic(off, chromium.subgroup_matrix_uniformity);
 {{ env.wgsl.resourceDeclarations }}
 {% set layout = layout | default("bsh") %}
@@ -19,27 +16,21 @@ diagnostic(off, chromium.subgroup_matrix_uniformity);
 {% macro kv_index(seq, d) %}{% if kvHeadMajor %}((b * KV_HEADS + h_kv) * params.kvSeq + {{ seq }}) * HEAD_DIM + {{ d }}{% else %}(b * params.kvSeq + {{ seq }}) * KV_HIDDEN + h_kv * HEAD_DIM + {{ d }}{% endif %}{% endmacro %}
 {% set OUT_ROW_STRIDE = "HEAD_DIM" if headMajor else "HIDDEN" %}
 {% set KV_ROW_STRIDE = "HEAD_DIM" if kvHeadMajor else "KV_HIDDEN" %}
 {% set scorePhase = phase == "score" %}
 {% set FUSED_SOFTMAX = fusedSoftmax is defined and fusedSoftmax %}
 {% set PRIVATE_ROW_STATS = FUSED_SOFTMAX and (materializedSgmatPrivateRowStats is defined and materializedSgmatPrivateRowStats) %}
 {% macro score_value(index, guard) %}
-{% if FUSED_SOFTMAX %}
-{% if PRIVATE_ROW_STATS %}
-select(0.0, exp_shift(scores[{{ index }}], private_softmax_m) / private_softmax_d, {{ guard[1] }})
-{%- else %}
-select(0.0, exp_shift(scores[{{ index }}], softmax_m[{{ guard[0] }}]) / softmax_d[{{ guard[0] }}], {{ guard[1] }})
-{%- endif %}
-{% else %}
-select(0.0, scores[{{ index }}], {{ guard[1] }})
-{%- endif %}
-{% endmacro %}
 {% set TILE_M_VALUE = materializedSgmatQueryTile %}
 {% set TILE_N_VALUE = materializedSgmatKeyTile %}
 {% set TILE_K_VALUE = materializedSgmatInnerTile %}
-{% set SUB_ROWS_VALUE = materializedSgmatSubgroupTileRows if materializedSgmatSubgroupTileRows is defined else 16 %}
-{% set SUB_COLS_VALUE = materializedSgmatSubgroupTileCols if materializedSgmatSubgroupTileCols is defined else 32 %}
-{% set ROW_BLOCKS = (SUB_ROWS_VALUE / 8)|int %}
-{% set COL_BLOCKS = (SUB_COLS_VALUE / 8)|int %}
 {% set SUBGROUP_ROWS = (TILE_M_VALUE / SUB_ROWS_VALUE)|int %}
 {% set SUBGROUP_COLS = (TILE_N_VALUE / SUB_COLS_VALUE)|int %}
 {% set SUBGROUP_COUNT = SUBGROUP_ROWS * SUBGROUP_COLS %}
@@ -47,15 +38,16 @@ select(0.0, scores[{{ index }}], {{ guard[1] }})
 {% if hasBias is not defined %}{% set hasBias = false %}{% endif %}
 {% set EMIT_ROW_STATS = emitRowStats is defined and emitRowStats %}
 {% set DIRECT_SCORE_STORE = materializedSgmatDirectScoreStore and not EMIT_ROW_STATS %}
-{% set DIRECT_APPLY_STORE = materializedSgmatDirectApplyStore and not hasBias and MT == "f32" %}
 {% set DIRECT_OUTPUT_STORE = DIRECT_SCORE_STORE if scorePhase else DIRECT_APPLY_STORE %}
 {% set RUNTIME_DIRECT_STORE = (not DIRECT_OUTPUT_STORE)
    and materializedSgmatRuntimeDirectStore
-   and (scorePhase or (not hasBias and MT == "f32"))
    and not EMIT_ROW_STATS %}
 {% set ANY_DIRECT_STORE = DIRECT_OUTPUT_STORE or RUNTIME_DIRECT_STORE %}
 {% set SCALE_IN_Q = DIRECT_SCORE_STORE or (RUNTIME_DIRECT_STORE and scorePhase) %}
-{% macro q_tile_value(index) %}{% if hasBias %}(query[{{ index }}] + bias[h * HEAD_DIM + k]){% else %}query[{{ index }}]{% endif %}{% if SCALE_IN_Q %} * score_scale{% endif %}{% endmacro %}
 const HEADS: u32 = {{ qNumHeads }}u;
 {% if kvNumHeads is defined %}
@@ -108,14 +100,12 @@ var<workgroup> softmax_m: array<f32, {{ TILE_M_VALUE }}>;
 var<workgroup> softmax_d: array<f32, {{ TILE_M_VALUE }}>;
 {% endif %}
 {% if FUSED_SOFTMAX or EMIT_ROW_STATS %}
-{% set stableUsage = stableHelperUsage if stableHelperUsage is defined else "all" %}
 // 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;
-{% if stableUsage != "constant" %}
 fn is_finite_f32(value: f32) -> bool {
   return select(false, value <= FLT_MAX, value >= -FLT_MAX);
@@ -129,14 +119,10 @@ 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);
 }
-{%- endif %}
-{% if stableUsage == "all" %}
 fn exp_shift(value: f32, maxValue: f32) -> f32 {
   return exp(shifted_value(value, maxValue));
 }
-{%- endif %}
 {% endif %}
 @compute @workgroup_size({{ WORKGROUP_THREADS }}, 1, 1){{ " @subgroup_size(32)" if pinSubgroupSize32 else "" }}
@@ -173,30 +159,18 @@ fn main(
   let inner = params.kvSeq;
 {% endif %}
 {% if FUSED_SOFTMAX %}
-{% if PRIVATE_ROW_STATS %}
-  // In the admitted BM64/BN64/BK32/WG256 loader, four adjacent lanes own the
-  // same query row for every reduction tile. Keep that row's constants private:
-  // this removes both 512 bytes of workgroup storage and the initialization
-  // barrier while preserving the exact exp/divide sequence of the shared-memory
-  // row-stats arm.
-  let private_stat_row =
-    (b * HEADS + h) * params.qSeq + min(m_base + li / 4u, params.qSeq - 1u);
-  let private_softmax_m = rowStats[private_stat_row * 2u];
-  let private_softmax_d = rowStats[private_stat_row * 2u + 1u];
-{% else %}
   // One row-stats pair per query row of the tile. A query tail clamps to the last
   // real row rather than reading past the buffer; those lanes are discarded by the
   // staging guard anyway, and the clamp keeps the denominator non-zero.
   for (var r = li; r < TILE_M; r += WORKGROUP_THREADS) {
-    let stat_row = (b * HEADS + h) * params.qSeq + min(m_base + r, params.qSeq - 1u);
     softmax_m[r] = rowStats[stat_row * 2u];
     softmax_d[r] = rowStats[stat_row * 2u + 1u];
   }
   workgroupBarrier();
-{% endif %}
 {% endif %}
   for (var k_base = 0u; k_base < inner; k_base += TILE_K) {
-{% if phase == "apply" and not FUSED_SOFTMAX and MT == "f32" %}
     // Full interior PV tiles can be loaded directly from storage. Query,
     // reduction, and output-dimension tails use the guarded shared path below.
     if (
@@ -206,7 +180,7 @@ fn main(
     ) {
       for (var step = 0u; step < TILE_K; step += 8u) {
 {% for row_block in range(ROW_BLOCKS) %}
-        let score_offset{{ row_block }} = (b * HEADS + h) * params.qSeq * params.kvSeq
           + (m_base + base_A + {{ row_block * 8 }}u) * params.kvSeq + k_base + step;
         var matA{{ row_block }}: subgroup_matrix_left<f32, 8, 8> =
           subgroupMatrixLoad<subgroup_matrix_left<f32, 8, 8>, row_major>(
@@ -256,7 +230,7 @@ fn main(
       tile_A[a_row * TILE_K + a_col + i] = loaded;
 {% endif %}
 {% else %}
-      let score_base = (b * HEADS + h) * params.qSeq * params.kvSeq;
       tile_A[a_row * TILE_K + a_col + i] =
         {{ "f16(" if MT == "f16" else "" }}{{ score_value("score_base + row * params.kvSeq + k", ["a_row", "row < params.qSeq && k < params.kvSeq"]) }}{{ ")" if MT == "f16" else "" }};
 {% endif %}
@@ -271,20 +245,20 @@ fn main(
 {% if headDim % 32 == 0 %}
       tile_B[b_row * TILE_K + b_col + i] = select(
         {{ "0.0h" if MT == "f16" else "0.0" }},
-        key[{{ kv_index("col", "k") }}],
         col < params.kvSeq
       );
 {% else %}
       var loaded = {{ "0.0h" if MT == "f16" else "0.0" }};
       if (col < params.kvSeq && k < HEAD_DIM) {
-        loaded = key[{{ kv_index("col", "k") }}];
       }
       tile_B[b_row * TILE_K + b_col + i] = loaded;
 {% endif %}
 {% else %}
       tile_B[b_row * TILE_K + b_col + i] = select(
         {{ "0.0h" if MT == "f16" else "0.0" }},
-        value[{{ kv_index("k", "col") }}],
         k < params.kvSeq && col < HEAD_DIM
       );
 {% endif %}
@@ -303,13 +277,13 @@ fn main(
         loaded = {{ q_tile_value(q_index("row", "k")) }};
       }
 {% elif FUSED_SOFTMAX %}
-      let score_base = (b * HEADS + h) * params.qSeq * params.kvSeq;
       let loaded =
         {{ "f16(" if MT == "f16" else "" }}{{ score_value("score_base + row * params.kvSeq + k", ["tile_row", "row < params.qSeq && k < params.kvSeq"]) }}{{ ")" if MT == "f16" else "" }};
 {% else %}
       var loaded = 0.0;
       if (row < params.qSeq && k < params.kvSeq) {
-        let score_base = (b * HEADS + h) * params.qSeq * params.kvSeq;
         loaded = scores[score_base + row * params.kvSeq + k];
       }
 {% endif %}
@@ -324,12 +298,12 @@ fn main(
 {% if scorePhase %}
       var loaded = {{ "0.0h" if MT == "f16" else "0.0" }};
       if (col < params.kvSeq && k < HEAD_DIM) {
-        loaded = key[{{ kv_index("col", "k") }}];
       }
 {% else %}
       var loaded = {{ "0.0h" if MT == "f16" else "0.0" }};
       if (k < params.kvSeq && col < HEAD_DIM) {
-        loaded = value[{{ kv_index("k", "col") }}];
       }
 {% endif %}
       tile_B[idx] = loaded;
@@ -382,7 +356,7 @@ fn main(
 {% for col_block in range(COL_BLOCKS) %}
 {% if scorePhase %}
   let output_offset{{ row_block }}{{ col_block }} =
-    (b * HEADS + h) * params.qSeq * params.kvSeq
       + (m_base + base_A + {{ row_block * 8 }}u) * params.kvSeq
       + n_base + base_B + {{ col_block * 8 }}u;
   subgroupMatrixStore<row_major>(
@@ -448,7 +422,7 @@ fn main(
 {% endif %}
       let scored = result{% if not SCALE_IN_Q %} * scale{% endif %};
       scores[
-        (b * HEADS + h) * params.qSeq * params.kvSeq + row * params.kvSeq + col
       ] = scored;
 {% if EMIT_ROW_STATS %}
       // Softmax sees the STORED value, so the statistics have to be taken on it
@@ -462,7 +436,7 @@ fn main(
 {% if hasBias %}
       // V bias row base: skip the packed Q and K blocks, then index this head.
 {% endif %}
-      output[{{ q_index("row", "col") }}] = {{ "f16(" if MT == "f16" else "" }}result{{ ")" if MT == "f16" else "" }}{% if hasBias %} + bias[2u * HIDDEN + h * HEAD_DIM + col]{% endif %};
 {% endif %}
     }
   }
@@ -490,7 +464,7 @@ fn main(
       // eight consecutive pairs, and the combine pass reads a slot's whole column
       // of rows contiguously.
       let slot = wg.x * {{ SUBGROUP_COLS }}u + subtile_idx;
-      let out_index = (((b * HEADS + h) * STAT_SLOTS + slot) * params.qSeq + stat_row) * 2u;
       scorePartials[out_index] = stat_m{{ row_block }};
       scorePartials[out_index + 1u] = stat_d{{ row_block }};
     }

 {% set MT = "f16" if (operandF16 is defined and operandF16) else "f32" %}
+{% set ST = "f16" if (storageF16 is defined and storageF16) else MT %}
 enable subgroups;
 {% if pinSubgroupSize32 %}
 enable subgroup_size_control;
 enable chromium_experimental_subgroup_matrix;
 diagnostic(off, chromium.subgroup_matrix_uniformity);
 {{ env.wgsl.resourceDeclarations }}
 {% set layout = layout | default("bsh") %}
 {% macro kv_index(seq, d) %}{% if kvHeadMajor %}((b * KV_HEADS + h_kv) * params.kvSeq + {{ seq }}) * HEAD_DIM + {{ d }}{% else %}(b * params.kvSeq + {{ seq }}) * KV_HIDDEN + h_kv * HEAD_DIM + {{ d }}{% endif %}{% endmacro %}
 {% set OUT_ROW_STRIDE = "HEAD_DIM" if headMajor else "HIDDEN" %}
 {% set KV_ROW_STRIDE = "HEAD_DIM" if kvHeadMajor else "KV_HIDDEN" %}
+{% set SCRATCH_HEAD = "(b * HEADS + h)" %}
 {% set scorePhase = phase == "score" %}
 {% set FUSED_SOFTMAX = fusedSoftmax is defined and fusedSoftmax %}
 {% set PRIVATE_ROW_STATS = FUSED_SOFTMAX and (materializedSgmatPrivateRowStats is defined and materializedSgmatPrivateRowStats) %}
 {% macro score_value(index, guard) %}
+{% set rowMax = "private_softmax_m" if PRIVATE_ROW_STATS else "softmax_m[" ~ guard[0] ~ "]" %}
+{% set rowDenom = "private_softmax_d" if PRIVATE_ROW_STATS else "softmax_d[" ~ guard[0] ~ "]" %}
+select(0.0, {{ "exp_shift(scores[" ~ index ~ "], " ~ rowMax ~ ") / " ~ rowDenom if FUSED_SOFTMAX else "scores[" ~ index ~ "]" }}, {{ guard[1] }}){% endmacro %}
 {% set TILE_M_VALUE = materializedSgmatQueryTile %}
 {% set TILE_N_VALUE = materializedSgmatKeyTile %}
 {% set TILE_K_VALUE = materializedSgmatInnerTile %}
+{% set SUB_ROWS_VALUE = 16 %}
+{% set SUB_COLS_VALUE = 32 %}
+{% set ROW_BLOCKS = 2 %}
+{% set COL_BLOCKS = 4 %}
 {% set SUBGROUP_ROWS = (TILE_M_VALUE / SUB_ROWS_VALUE)|int %}
 {% set SUBGROUP_COLS = (TILE_N_VALUE / SUB_COLS_VALUE)|int %}
 {% set SUBGROUP_COUNT = SUBGROUP_ROWS * SUBGROUP_COLS %}
 {% if hasBias is not defined %}{% set hasBias = false %}{% endif %}
 {% set EMIT_ROW_STATS = emitRowStats is defined and emitRowStats %}
 {% set DIRECT_SCORE_STORE = materializedSgmatDirectScoreStore and not EMIT_ROW_STATS %}
+{% set DIRECT_APPLY_STORE = materializedSgmatDirectApplyStore and not hasBias and MT == "f32" and ST == "f32" %}
 {% set DIRECT_OUTPUT_STORE = DIRECT_SCORE_STORE if scorePhase else DIRECT_APPLY_STORE %}
 {% set RUNTIME_DIRECT_STORE = (not DIRECT_OUTPUT_STORE)
    and materializedSgmatRuntimeDirectStore
+   and (scorePhase or (not hasBias and MT == "f32" and ST == "f32"))
    and not EMIT_ROW_STATS %}
 {% set ANY_DIRECT_STORE = DIRECT_OUTPUT_STORE or RUNTIME_DIRECT_STORE %}
 {% set SCALE_IN_Q = DIRECT_SCORE_STORE or (RUNTIME_DIRECT_STORE and scorePhase) %}
+{% macro operand_load(name, index) %}{% if ST != MT %}{{ MT }}({% endif %}{{ name }}[{{ index }}]{% if ST != MT %}){% endif %}{% endmacro %}
+{% macro q_tile_value(index) %}{% if hasBias %}({{ operand_load("query", index) }} + bias[h * HEAD_DIM + k]){% else %}{{ operand_load("query", index) }}{% endif %}{% if SCALE_IN_Q %} * score_scale{% endif %}{% endmacro %}
 const HEADS: u32 = {{ qNumHeads }}u;
 {% if kvNumHeads is defined %}
 var<workgroup> softmax_d: array<f32, {{ TILE_M_VALUE }}>;
 {% endif %}
 {% if FUSED_SOFTMAX or EMIT_ROW_STATS %}
 // 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);
   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));
 }
 {% endif %}
 @compute @workgroup_size({{ WORKGROUP_THREADS }}, 1, 1){{ " @subgroup_size(32)" if pinSubgroupSize32 else "" }}
   let inner = params.kvSeq;
 {% endif %}
 {% if FUSED_SOFTMAX %}
   // One row-stats pair per query row of the tile. A query tail clamps to the last
   // real row rather than reading past the buffer; those lanes are discarded by the
   // staging guard anyway, and the clamp keeps the denominator non-zero.
   for (var r = li; r < TILE_M; r += WORKGROUP_THREADS) {
+    let stat_row = {{ SCRATCH_HEAD }} * params.qSeq + min(m_base + r, params.qSeq - 1u);
     softmax_m[r] = rowStats[stat_row * 2u];
     softmax_d[r] = rowStats[stat_row * 2u + 1u];
   }
   workgroupBarrier();
 {% endif %}
   for (var k_base = 0u; k_base < inner; k_base += TILE_K) {
+{% if phase == "apply" and not FUSED_SOFTMAX and MT == "f32" and ST == "f32" %}
     // Full interior PV tiles can be loaded directly from storage. Query,
     // reduction, and output-dimension tails use the guarded shared path below.
     if (
     ) {
       for (var step = 0u; step < TILE_K; step += 8u) {
 {% for row_block in range(ROW_BLOCKS) %}
+        let score_offset{{ row_block }} = {{ SCRATCH_HEAD }} * params.qSeq * params.kvSeq
           + (m_base + base_A + {{ row_block * 8 }}u) * params.kvSeq + k_base + step;
         var matA{{ row_block }}: subgroup_matrix_left<f32, 8, 8> =
           subgroupMatrixLoad<subgroup_matrix_left<f32, 8, 8>, row_major>(
       tile_A[a_row * TILE_K + a_col + i] = loaded;
 {% endif %}
 {% else %}
+      let score_base = {{ SCRATCH_HEAD }} * params.qSeq * params.kvSeq;
       tile_A[a_row * TILE_K + a_col + i] =
         {{ "f16(" if MT == "f16" else "" }}{{ score_value("score_base + row * params.kvSeq + k", ["a_row", "row < params.qSeq && k < params.kvSeq"]) }}{{ ")" if MT == "f16" else "" }};
 {% endif %}
 {% if headDim % 32 == 0 %}
       tile_B[b_row * TILE_K + b_col + i] = select(
         {{ "0.0h" if MT == "f16" else "0.0" }},
+        {{ operand_load("key", kv_index("col", "k")) }},
         col < params.kvSeq
       );
 {% else %}
       var loaded = {{ "0.0h" if MT == "f16" else "0.0" }};
       if (col < params.kvSeq && k < HEAD_DIM) {
+        loaded = {{ operand_load("key", kv_index("col", "k")) }};
       }
       tile_B[b_row * TILE_K + b_col + i] = loaded;
 {% endif %}
 {% else %}
       tile_B[b_row * TILE_K + b_col + i] = select(
         {{ "0.0h" if MT == "f16" else "0.0" }},
+        {{ operand_load("value", kv_index("k", "col")) }},
         k < params.kvSeq && col < HEAD_DIM
       );
 {% endif %}
         loaded = {{ q_tile_value(q_index("row", "k")) }};
       }
 {% elif FUSED_SOFTMAX %}
+      let score_base = {{ SCRATCH_HEAD }} * params.qSeq * params.kvSeq;
       let loaded =
         {{ "f16(" if MT == "f16" else "" }}{{ score_value("score_base + row * params.kvSeq + k", ["tile_row", "row < params.qSeq && k < params.kvSeq"]) }}{{ ")" if MT == "f16" else "" }};
 {% else %}
       var loaded = 0.0;
       if (row < params.qSeq && k < params.kvSeq) {
+        let score_base = {{ SCRATCH_HEAD }} * params.qSeq * params.kvSeq;
         loaded = scores[score_base + row * params.kvSeq + k];
       }
 {% endif %}
 {% if scorePhase %}
       var loaded = {{ "0.0h" if MT == "f16" else "0.0" }};
       if (col < params.kvSeq && k < HEAD_DIM) {
+        loaded = {{ operand_load("key", kv_index("col", "k")) }};
       }
 {% else %}
       var loaded = {{ "0.0h" if MT == "f16" else "0.0" }};
       if (k < params.kvSeq && col < HEAD_DIM) {
+        loaded = {{ operand_load("value", kv_index("k", "col")) }};
       }
 {% endif %}
       tile_B[idx] = loaded;
 {% for col_block in range(COL_BLOCKS) %}
 {% if scorePhase %}
   let output_offset{{ row_block }}{{ col_block }} =
+    {{ SCRATCH_HEAD }} * params.qSeq * params.kvSeq
       + (m_base + base_A + {{ row_block * 8 }}u) * params.kvSeq
       + n_base + base_B + {{ col_block * 8 }}u;
   subgroupMatrixStore<row_major>(
 {% endif %}
       let scored = result{% if not SCALE_IN_Q %} * scale{% endif %};
       scores[
+        {{ SCRATCH_HEAD }} * params.qSeq * params.kvSeq + row * params.kvSeq + col
       ] = scored;
 {% if EMIT_ROW_STATS %}
       // Softmax sees the STORED value, so the statistics have to be taken on it
 {% if hasBias %}
       // V bias row base: skip the packed Q and K blocks, then index this head.
 {% endif %}
+      output[{{ q_index("row", "col") }}] = {{ "f16(" if ST == "f16" else "" }}result{{ ")" if ST == "f16" else "" }}{% if hasBias %} + bias[2u * HIDDEN + h * HEAD_DIM + col]{% endif %};
 {% endif %}
     }
   }
       // eight consecutive pairs, and the combine pass reads a slot's whole column
       // of rows contiguously.
       let slot = wg.x * {{ SUBGROUP_COLS }}u + subtile_idx;
+      let out_index = (({{ SCRATCH_HEAD }} * STAT_SLOTS + slot) * params.qSeq + stat_row) * 2u;
       scorePartials[out_index] = stat_m{{ row_block }};
       scorePartials[out_index + 1u] = stat_d{{ row_block }};
     }

build/webgpu/attn-materialized-softmax-f32.wgsl.jinja CHANGED Viewed

@@ -1,9 +1,7 @@
 {% set CACHE_VEC4 = cacheVec4 if cacheVec4 is defined else false %}
-{% set CAUSAL_ROWS = causalRows is defined and causalRows %}
-{% set SCALE_IN = scaleInSoftmax is defined and scaleInSoftmax %}
-{% set RUNTIME_COLS = colsFromParams is defined and colsFromParams %}
-{% set COLS_EXPR = "params.keyLen" if RUNTIME_COLS else "COLS" %}
-{% set ROWS_EXPR = rowsExpr if rowsExpr is defined else "params.rows" %}
 {% if useSubgroups %}
 enable subgroups;
 {% endif %}
@@ -12,7 +10,7 @@ enable subgroups;
 // In-place row softmax for materialized f32 attention. The default path reads
 // each scalar twice; the register-cached vec4 path eliminates the second read.
 const WG: u32 = {{ materializedSoftmaxWg }}u;
-{% set SCALED_SCORE = "scores[base + c] * SCALE" if SCALE_IN else "scores[base + c]" %}
 {% if CACHE_VEC4 %}
 const COLS4: u32 = {{ materializedSoftmaxCols4 }}u;
 const CACHE_VECS: u32 = (COLS4 + WG - 1u) / WG;
@@ -38,6 +36,7 @@ fn shifted_value(value: f32, maxValue: f32) -> f32 {
   let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
   return select(value - maxValue, 0.0, equalFiniteMax);
 }
 fn exp_shift(value: f32, maxValue: f32) -> f32 {
   return exp(shifted_value(value, maxValue));
 }
@@ -101,41 +100,9 @@ fn combine_partials(m: f32, d: f32, lidx: u32, sgSize: u32) -> vec2<f32> {
   return combinedMD;
 }
 {% else %}
-{% set mdStreamed = mdStreams is defined %}
-{% set mdStreams = mdStreams if mdStreams is defined else 1 %}
-{% set mdExtent = "WG" if mdStreams == 1 else "WG * " ~ mdStreams ~ "u" %}
 var<workgroup> partialM: array<f32, {{ mdExtent }}>;
 var<workgroup> partialD: array<f32, {{ mdExtent }}>;
-{% if mdStreamed %}
-// In-place fold of {{ mdStreams }} streams. Input partials occupy
-// partialM/partialD; stream s returns its merged pair in slot s * WG.
-fn combine_partials_streams(lidx: u32) {
-  workgroupBarrier();
-  var stride = WG / 2u;
-  loop {
-    if (stride == 0u) {
-      break;
-    }
-    if (lidx < stride) {
-{% for s in range(mdStreams) %}
-      {
-        let slot = {{ s }}u * WG + lidx;
-        let m1 = partialM[slot];
-        let d1 = partialD[slot];
-        let m2 = partialM[slot + stride];
-        let d2 = partialD[slot + stride];
-        let mNew = max(m1, m2);
-        partialD[slot] = d1 * exp_shift(m1, mNew) + d2 * exp_shift(m2, mNew);
-        partialM[slot] = mNew;
-      }
-{% endfor %}
-    }
-    workgroupBarrier();
-    stride = stride / 2u;
-  }
-}
-{% else %}
 fn combine_partials(m: f32, d: f32, lidx: u32) -> vec2<f32> {
   partialM[lidx] = m;
@@ -165,8 +132,6 @@ fn combine_partials(m: f32, d: f32, lidx: u32) -> vec2<f32> {
   return merged;
 }
 {% endif %}
-{% endif %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(

 {% set CACHE_VEC4 = cacheVec4 if cacheVec4 is defined else false %}
+{% set CAUSAL_ROWS = false %}
+{% set COLS_EXPR = "COLS" %}
+{% set ROWS_EXPR = "params.rows" %}
 {% if useSubgroups %}
 enable subgroups;
 {% endif %}
 // In-place row softmax for materialized f32 attention. The default path reads
 // each scalar twice; the register-cached vec4 path eliminates the second read.
 const WG: u32 = {{ materializedSoftmaxWg }}u;
+{% set SCALED_SCORE = "scores[base + c]" %}
 {% if CACHE_VEC4 %}
 const COLS4: u32 = {{ materializedSoftmaxCols4 }}u;
 const CACHE_VECS: u32 = (COLS4 + WG - 1u) / WG;
   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));
 }
   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;
   return merged;
 }
 {% endif %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(

build/webgpu/attn-online-scalar.wgsl.jinja CHANGED Viewed

@@ -1,11 +1,11 @@
 {{ env.wgsl.resourceDeclarations }}
-{% set MASK_BATCH = "batch * params.maskBatchStride + " if layout == "bsh" else "" %}
 // Online-softmax attention fallback with no feature requirements: one
 // workgroup per (batch, head, query token) walks the keys serially; the
-// workgroup cooperates on each q·k dot (tree reduction) and on the running
-// V accumulator, with the online rescale applied per key. This path requires
-// no subgroup or subgroup-matrix features.
 // Layout: rank-3 token-major [batch, seq, heads * headDim].
 // An omitted scale uses 1/sqrt(headDim); explicit zero remains zero through
 // the scaleIsExplicitZero specialization.
@@ -21,18 +21,39 @@ const Q_HIDDEN: u32 = {{ qHidden }}u;
 const KV_HIDDEN: u32 = {{ kvHidden }}u;
 const Q_HEADS: u32 = {{ qNumHeads }}u;
 const KV_HEADS: u32 = {{ kvNumHeads }}u;
-{% set qHeads = "params.qHeads" if headsFromParams else "Q_HEADS" %}
-{% set kvHeads = "params.kvHeads" if headsFromParams else "KV_HEADS" %}
-{% set scale = scale | default("0.0") %}
 const WG: u32 = {{ workgroupSize }}u;
-var<workgroup> partial: array<f32, WG>;
 var<workgroup> running_max: f32;
 var<workgroup> running_denom: f32;
 var<workgroup> running_out: array<f32, HEAD_DIM>;
 var<workgroup> previous_scale: f32;
-{% macro wgsl_tree_reduce_f32(name, mode, buffer="partial", wg="WG", trailingBarrier=true) %}
-fn {{ name }}(value: f32, tid: u32) -> f32 {
   {{ buffer }}[tid] = value;
   workgroupBarrier();
   // Ceil-halving keeps every lane when the workgroup size is not a power of
@@ -42,11 +63,7 @@ fn {{ name }}(value: f32, tid: u32) -> f32 {
   loop {
     let half = (n + 1u) / 2u;
     if (tid < n - half) {
-{% if mode == "max" %}
-      {{ buffer }}[tid] = max({{ buffer }}[tid], {{ buffer }}[tid + half]);
-{% else %}
-      {{ buffer }}[tid] = {{ buffer }}[tid] + {{ buffer }}[tid + half];
-{% endif %}
     }
     workgroupBarrier();
     n = half;
@@ -58,22 +75,16 @@ fn {{ name }}(value: f32, tid: u32) -> f32 {
   // slot 0 here, so the next call's first store must not run until all lanes have read it.
   // `trailingBarrier=false` is safe only when the buffer is never written again before kernel exit.
   let reduced = {{ buffer }}[0];
-{% if trailingBarrier %}
   workgroupBarrier();
-{% endif %}
   return reduced;
 }
 {% endmacro %}
-{{ wgsl_tree_reduce_f32("reduce_sum", "add", "partial", "WG") }}
-{% if ATTN_SCALE_DIM is not defined %}{% set ATTN_SCALE_DIM = "HEAD_DIM" %}{% endif %}
-fn scale_value() -> f32 {
   if (params.scale != 0.0) { return params.scale; }
   return inverseSqrt(f32({{ ATTN_SCALE_DIM }}));
 }
 @compute @workgroup_size(WG, 1, 1)
 fn main(
   @builtin(workgroup_id) wg: vec3<u32>,
@@ -117,7 +128,7 @@ fn main(
   for (var key_token: u32 = minKj; key_token < maxKj; key_token = key_token + 1u) {
     let kRow = kvBase + key_token * kvTokenStride;
-    var partial_dot = 0.0;
     for (var d: u32 = tid; d < HEAD_DIM; d = d + WG) {
       var q_value = f32(query[qBase + d]);
       var k_value = f32(key[kRow + d]);
@@ -126,17 +137,17 @@ fn main(
       q_value = q_value + f32(bias[channel]);
       k_value = k_value + f32(bias[Q_HIDDEN + channel]);
 {% endif %}
-      partial_dot = partial_dot + q_value * k_value;
     }
-    // reduce_sum returns the same partial[0] to every lane, so `score` is already
     // workgroup-uniform; tid==0 advances the running max/denom/scale (broadcast
     // through shared memory by the barrier below).
 {% if hasMask %}
     let maskIndex = {{ MASK_BATCH }}h * params.maskHeadStride + query_token * params.maskSeqStride + key_token;
-    let score = reduce_sum(partial_dot, tid) * scale_value() + f32(attn_mask[maskIndex]);
 {% else %}
-    let score = reduce_sum(partial_dot, tid) * scale_value();
 {% endif %}
     if (tid == 0u) {
       let next_max = max(running_max, score);

 {{ env.wgsl.resourceDeclarations }}
+{% set MASK_BATCH = "batch * params.maskBatchStride + " %}
 // Online-softmax attention fallback with no feature requirements: one
 // workgroup per (batch, head, query token) walks the keys serially; the
+// workgroup cooperates on each q·k dot (compensated tree reduction) and on the
+// running V accumulator, with the online rescale applied per key. This path
+// requires no subgroup or subgroup-matrix features.
 // Layout: rank-3 token-major [batch, seq, heads * headDim].
 // An omitted scale uses 1/sqrt(headDim); explicit zero remains zero through
 // the scaleIsExplicitZero specialization.
 const KV_HIDDEN: u32 = {{ kvHidden }}u;
 const Q_HEADS: u32 = {{ qNumHeads }}u;
 const KV_HEADS: u32 = {{ kvNumHeads }}u;
+{% set qHeads = "Q_HEADS" %}
+{% set kvHeads = "KV_HEADS" %}
 const WG: u32 = {{ workgroupSize }}u;
+{% set dotType = "f32" %}
+// Retain product and addition residuals across a dot product. Explicit fma
+// boundaries preserve the addition error transform under reassociation.
+struct DotAccumulator {
+  hi: {{ dotType }},
+  lo: {{ dotType }},
+}
+fn dot_accumulate(acc: DotAccumulator, a: {{ dotType }}, b: {{ dotType }}) -> DotAccumulator {
+  let product = fma(a, b, {{ dotType }}(0.0));
+  let productError = fma(a, b, -product);
+  let sum = fma(acc.hi, {{ dotType }}(1.0), product);
+  let bv = fma({{ dotType }}(-1.0), acc.hi, sum);
+  let av = fma({{ dotType }}(-1.0), bv, sum);
+  let ae = fma({{ dotType }}(-1.0), av, acc.hi);
+  let be = fma({{ dotType }}(-1.0), bv, product);
+  let error = fma(fma(fma(ae, {{ dotType }}(1.0), be), {{ dotType }}(1.0), acc.lo), {{ dotType }}(1.0), productError);
+  let hi = fma(sum, {{ dotType }}(1.0), error);
+  return DotAccumulator(hi, fma({{ dotType }}(-1.0), fma({{ dotType }}(-1.0), sum, hi), error));
+}
+fn dot_value(acc: DotAccumulator) -> {{ dotType }} {
+  return fma(acc.hi, {{ dotType }}(1.0), acc.lo);
+}
+var<workgroup> partial: array<DotAccumulator, WG>;
 var<workgroup> running_max: f32;
 var<workgroup> running_denom: f32;
 var<workgroup> running_out: array<f32, HEAD_DIM>;
 var<workgroup> previous_scale: f32;
+{% macro wgsl_tree_reduce_f32(name, mode, buffer="partial", wg="WG", trailingBarrier=true, valueType="f32") %}
+fn {{ name }}(value: {{ valueType }}, tid: u32) -> {{ valueType }} {
   {{ buffer }}[tid] = value;
   workgroupBarrier();
   // Ceil-halving keeps every lane when the workgroup size is not a power of
   loop {
     let half = (n + 1u) / 2u;
     if (tid < n - half) {
+      {{ buffer }}[tid] = dot_accumulate(dot_accumulate({{ buffer }}[tid], {{ buffer }}[tid + half].hi, 1.0), {{ buffer }}[tid + half].lo, 1.0);
     }
     workgroupBarrier();
     n = half;
   // slot 0 here, so the next call's first store must not run until all lanes have read it.
   // `trailingBarrier=false` is safe only when the buffer is never written again before kernel exit.
   let reduced = {{ buffer }}[0];
   workgroupBarrier();
   return reduced;
 }
 {% endmacro %}
+{{ wgsl_tree_reduce_f32("reduce_dot", "compensated", "partial", "WG", valueType="DotAccumulator") }}
+{% set ATTN_SCALE_DIM = "HEAD_DIM" %}fn scale_value() -> f32 {
   if (params.scale != 0.0) { return params.scale; }
   return inverseSqrt(f32({{ ATTN_SCALE_DIM }}));
 }
 @compute @workgroup_size(WG, 1, 1)
 fn main(
   @builtin(workgroup_id) wg: vec3<u32>,
   for (var key_token: u32 = minKj; key_token < maxKj; key_token = key_token + 1u) {
     let kRow = kvBase + key_token * kvTokenStride;
+    var partial_dot = DotAccumulator(0.0, 0.0);
     for (var d: u32 = tid; d < HEAD_DIM; d = d + WG) {
       var q_value = f32(query[qBase + d]);
       var k_value = f32(key[kRow + d]);
       q_value = q_value + f32(bias[channel]);
       k_value = k_value + f32(bias[Q_HIDDEN + channel]);
 {% endif %}
+      partial_dot = dot_accumulate(partial_dot, q_value, k_value);
     }
+    // reduce_dot returns the same partial[0] to every lane, so `score` is already
     // workgroup-uniform; tid==0 advances the running max/denom/scale (broadcast
     // through shared memory by the barrier below).
 {% if hasMask %}
     let maskIndex = {{ MASK_BATCH }}h * params.maskHeadStride + query_token * params.maskSeqStride + key_token;
+    let score = dot_value(reduce_dot(partial_dot, tid)) * scale_value() + f32(attn_mask[maskIndex]);
 {% else %}
+    let score = dot_value(reduce_dot(partial_dot, tid)) * scale_value();
 {% endif %}
     if (tid == 0u) {
       let next_max = max(running_max, score);

build/webgpu/attn-small-head-parallel.wgsl.jinja CHANGED Viewed

@@ -9,17 +9,15 @@ const HIDDEN: u32 = {{ qHidden }}u;
 const KV_SEQ: u32 = {{ kvSeq }}u;
 const WG: u32 = 64u;
 const NEG_MAX: f32 = -3.4028234663852886e38;
-{% if ATTN_SCALE_DIM is not defined %}{% set ATTN_SCALE_DIM = "HEAD_DIM" %}{% endif %}
-fn scale_value() -> f32 {
   if (params.scale != 0.0) { return params.scale; }
   return inverseSqrt(f32({{ ATTN_SCALE_DIM }}));
 }
 var<workgroup> scores: array<f32, KV_SEQ>;
 var<workgroup> partial: array<f32, WG>;
-{% macro wgsl_tree_reduce_f32(name, mode, buffer="partial", wg="WG", trailingBarrier=true) %}
-fn {{ name }}(value: f32, tid: u32) -> f32 {
   {{ buffer }}[tid] = value;
   workgroupBarrier();
   // Ceil-halving keeps every lane when the workgroup size is not a power of
@@ -45,16 +43,12 @@ fn {{ name }}(value: f32, tid: u32) -> f32 {
   // slot 0 here, so the next call's first store must not run until all lanes have read it.
   // `trailingBarrier=false` is safe only when the buffer is never written again before kernel exit.
   let reduced = {{ buffer }}[0];
-{% if trailingBarrier %}
   workgroupBarrier();
-{% endif %}
   return reduced;
 }
 {% endmacro %}
 {{ wgsl_tree_reduce_f32("reduce_max", "max", "partial", "WG") }}
 {{ wgsl_tree_reduce_f32("reduce_sum", "add", "partial", "WG") }}
 @compute @workgroup_size(WG, 1, 1)
 fn main(
   @builtin(workgroup_id) wg: vec3<u32>,

 const KV_SEQ: u32 = {{ kvSeq }}u;
 const WG: u32 = 64u;
 const NEG_MAX: f32 = -3.4028234663852886e38;
+{% set ATTN_SCALE_DIM = "HEAD_DIM" %}fn scale_value() -> f32 {
   if (params.scale != 0.0) { return params.scale; }
   return inverseSqrt(f32({{ ATTN_SCALE_DIM }}));
 }
 var<workgroup> scores: array<f32, KV_SEQ>;
 var<workgroup> partial: array<f32, WG>;
+{% macro wgsl_tree_reduce_f32(name, mode, buffer="partial", wg="WG", trailingBarrier=true, valueType="f32") %}
+fn {{ name }}(value: {{ valueType }}, tid: u32) -> {{ valueType }} {
   {{ buffer }}[tid] = value;
   workgroupBarrier();
   // Ceil-halving keeps every lane when the workgroup size is not a power of
   // slot 0 here, so the next call's first store must not run until all lanes have read it.
   // `trailingBarrier=false` is safe only when the buffer is never written again before kernel exit.
   let reduced = {{ buffer }}[0];
   workgroupBarrier();
   return reduced;
 }
 {% endmacro %}
 {{ wgsl_tree_reduce_f32("reduce_max", "max", "partial", "WG") }}
 {{ wgsl_tree_reduce_f32("reduce_sum", "add", "partial", "WG") }}
 @compute @workgroup_size(WG, 1, 1)
 fn main(
   @builtin(workgroup_id) wg: vec3<u32>,

build/webgpu/attn-small-head-value.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,198 @@

+{% if useValueSubgroups %}
+enable subgroups;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+{% set rowType = "f32" if queryBlock == 1 else "vec" ~ queryBlock ~ "<f32>" %}
+{% set dotValueType = rowType %}
+{% set dotType = dotValueType | default("f32") %}
+// Retain product and addition residuals across a dot product. Explicit fma
+// boundaries preserve the addition error transform under reassociation.
+struct DotAccumulator {
+  hi: {{ dotType }},
+  lo: {{ dotType }},
+}
+fn dot_accumulate(acc: DotAccumulator, a: {{ dotType }}, b: {{ dotType }}) -> DotAccumulator {
+  let product = fma(a, b, {{ dotType }}(0.0));
+  let productError = fma(a, b, -product);
+  let sum = fma(acc.hi, {{ dotType }}(1.0), product);
+  let bv = fma({{ dotType }}(-1.0), acc.hi, sum);
+  let av = fma({{ dotType }}(-1.0), bv, sum);
+  let ae = fma({{ dotType }}(-1.0), av, acc.hi);
+  let be = fma({{ dotType }}(-1.0), bv, product);
+  let error = fma(fma(fma(ae, {{ dotType }}(1.0), be), {{ dotType }}(1.0), acc.lo), {{ dotType }}(1.0), productError);
+  let hi = fma(sum, {{ dotType }}(1.0), error);
+  return DotAccumulator(hi, fma({{ dotType }}(-1.0), fma({{ dotType }}(-1.0), sum, hi), error));
+}
+fn dot_value(acc: DotAccumulator) -> {{ dotType }} {
+  return fma(acc.hi, {{ dotType }}(1.0), acc.lo);
+}
+const HEAD_DIM: u32 = {{ headDim }}u;
+const HIDDEN: u32 = {{ qHidden }}u;
+const KV_SEQ: u32 = {{ kvSeq }}u;
+const WG: u32 = {{ valueWorkgroupSize }}u;
+const QUERY_BLOCK: u32 = {{ queryBlock }}u;
+const NEG_MAX: f32 = -3.4028234663852886e38;
+{% set ATTN_SCALE_DIM = "HEAD_DIM" %}fn scale_value() -> f32 {
+  if (params.scale != 0.0) { return params.scale; }
+  return inverseSqrt(f32({{ ATTN_SCALE_DIM }}));
+}
+var<workgroup> scores: array<{{ rowType }}, {{ scoreStorageElements }}u>;
+var<workgroup> partial: array<{{ rowType }}, WG>;
+{% if useValueSubgroups %}
+{% macro reduce(name, collective, mode, identity) %}
+fn {{ name }}(value: {{ rowType }}, sgLane: u32, sgId: u32, numSg: u32) -> {{ rowType }} {
+  let sub = {{ collective }}(value);
+  if (sgLane == 0u) { partial[sgId] = sub; }
+  workgroupBarrier();
+  var total = {{ rowType }}({{ identity }});
+  for (var i = 0u; i < numSg; i = i + 1u) {
+{% if mode == "max" %}
+    total = max(total, partial[i]);
+{% else %}
+    total = total + partial[i];
+{% endif %}
+  }
+  workgroupBarrier();
+  return total;
+}
+{% endmacro %}
+{{ reduce("reduce_max", "subgroupMax", "max", "NEG_MAX") }}
+{{ reduce("reduce_sum", "subgroupAdd", "add", "0.0") }}
+{% else %}
+{% macro wgsl_tree_reduce_f32(name, mode, buffer="partial", wg="WG", trailingBarrier=true, valueType="f32") %}
+fn {{ name }}(value: {{ valueType }}, tid: u32) -> {{ valueType }} {
+  {{ buffer }}[tid] = value;
+  workgroupBarrier();
+  // Ceil-halving keeps every lane when the workgroup size is not a power of
+  // two. For even n this matches the power-of-two tree order; for odd n, lanes
+  // [0, n-half) fold the upper tail while the middle lane carries forward.
+  var n: u32 = {{ wg }};
+  loop {
+    let half = (n + 1u) / 2u;
+    if (tid < n - half) {
+{% if mode == "max" %}
+      {{ buffer }}[tid] = max({{ buffer }}[tid], {{ buffer }}[tid + half]);
+{% else %}
+      {{ buffer }}[tid] = {{ buffer }}[tid] + {{ buffer }}[tid + half];
+{% endif %}
+    }
+    workgroupBarrier();
+    n = half;
+    if (n == 1u) {
+      break;
+    }
+  }
+  // The default trailing barrier makes this helper safe for back-to-back calls: every lane reads
+  // slot 0 here, so the next call's first store must not run until all lanes have read it.
+  // `trailingBarrier=false` is safe only when the buffer is never written again before kernel exit.
+  let reduced = {{ buffer }}[0];
+  workgroupBarrier();
+  return reduced;
+}
+{% endmacro %}
+{{ wgsl_tree_reduce_f32("reduce_max", "max", valueType=rowType) }}
+{{ wgsl_tree_reduce_f32("reduce_sum", "add", valueType=rowType) }}
+{% endif %}
+@compute @workgroup_size(WG, 1, 1)
+fn main(
+  @builtin(workgroup_id) wg: vec3<u32>,
+  @builtin(num_workgroups) nwg: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>
+{% if useValueSubgroups %}
+  , @builtin(subgroup_invocation_id) sgLane: u32
+  , @builtin(subgroup_id) sgId: u32
+  , @builtin(num_subgroups) numSg: u32
+{% endif %}
+) {
+  let h = wg.y;
+  let batch = wg.z;
+  let tid = lid.x;
+  let kvBase = batch * KV_SEQ * HIDDEN + h * HEAD_DIM;
+  for (var query_token = wg.x * QUERY_BLOCK; query_token < params.qSeq; query_token = query_token + nwg.x * QUERY_BLOCK) {
+{% for row in range(queryBlock) %}
+    let qBase_{{ row }} = (batch * params.qSeq + {% if queryTail and row > 0 %}min(query_token + {{ row }}u, params.qSeq - 1u){% else %}query_token{% if row > 0 %} + {{ row }}u{% endif %}{% endif %}) * HIDDEN + h * HEAD_DIM;
+{% endfor %}
+    // Query components stay private throughout the key scan. Their lifetime
+    // ends before the value accumulators need the same registers.
+    {
+{% for d in range(headDim) %}
+      let q_{{ d }} = {% if queryBlock > 1 %}{{ rowType }}({% endif %}{% for row in range(queryBlock) %}{% if not loop.first %}, {% endif %}f32(query[qBase_{{ row }} + {{ d }}u]){% endfor %}{% if queryBlock > 1 %}){% endif %};
+{% endfor %}
+      for (var key_token = tid; key_token < KV_SEQ; key_token = key_token + WG) {
+        let kBase = kvBase + key_token * HIDDEN;
+        var score = DotAccumulator({{ rowType }}(0.0), {{ rowType }}(0.0));
+{% for d in range(headDim) %}
+        score = dot_accumulate(score, q_{{ d }}, {{ rowType }}(f32(key[kBase + {{ d }}u])));
+{% endfor %}
+        scores[key_token] = dot_value(score) * scale_value();
+      }
+    }
+    workgroupBarrier();
+    var laneMax = {{ rowType }}(NEG_MAX);
+    for (var key_token = tid; key_token < KV_SEQ; key_token = key_token + WG) {
+      laneMax = max(laneMax, scores[key_token]);
+    }
+    let rowMax = reduce_max(laneMax, {% if useValueSubgroups %}sgLane, sgId, numSg{% else %}tid{% endif %});
+    var laneSum = {{ rowType }}(0.0);
+    for (var key_token = tid; key_token < KV_SEQ; key_token = key_token + WG) {
+      let probability = exp(scores[key_token] - rowMax);
+      scores[key_token] = probability;
+      laneSum = laneSum + probability;
+    }
+    let invDenom = {{ rowType }}(1.0) / reduce_sum(laneSum, {% if useValueSubgroups %}sgLane, sgId, numSg{% else %}tid{% endif %});
+{% for d in range(headDim) %}
+    var value_acc_{{ d }} = {{ rowType }}(0.0);
+{% endfor %}
+    for (var key_token = tid; key_token < KV_SEQ; key_token = key_token + WG) {
+      let probability = scores[key_token];
+      let vBase = kvBase + key_token * HIDDEN;
+{% for d in range(headDim) %}
+      value_acc_{{ d }} = value_acc_{{ d }} + probability * f32(value[vBase + {{ d }}u]);
+{% endfor %}
+    }
+{% if useValueSubgroups %}
+{% for d in range(headDim) %}
+    let value_sum_{{ d }} = subgroupAdd(value_acc_{{ d }});
+{% endfor %}
+    workgroupBarrier();
+    if (sgLane == 0u) {
+{% for d in range(headDim) %}
+      scores[sgId * HEAD_DIM + {{ d }}u] = value_sum_{{ d }};
+{% endfor %}
+    }
+    workgroupBarrier();
+    if (tid < HEAD_DIM) {
+      var total = {{ rowType }}(0.0);
+      for (var i = 0u; i < numSg; i = i + 1u) { total = total + scores[i * HEAD_DIM + tid]; }
+      let result = total * invDenom;
+{% else %}
+    workgroupBarrier();
+{% for d in range(headDim) %}
+    scores[{{ d }}u * WG + tid] = value_acc_{{ d }};
+{% endfor %}
+    workgroupBarrier();
+    for (var stride = WG / 2u; stride > 0u; stride = stride / 2u) {
+      if (tid < stride) {
+{% for d in range(headDim) %}
+        scores[{{ d }}u * WG + tid] = scores[{{ d }}u * WG + tid] + scores[{{ d }}u * WG + tid + stride];
+{% endfor %}
+      }
+      workgroupBarrier();
+    }
+    if (tid < HEAD_DIM) {
+      let result = scores[tid * WG] * invDenom;
+{% endif %}
+{% for row in range(queryBlock) %}
+{% if queryTail and row > 0 %}
+      if (query_token + {{ row }}u < params.qSeq) {
+{% endif %}
+        output[qBase_{{ row }} + tid] = {{ outputScalar }}(result{% if queryBlock > 1 %}[{{ row }}]{% endif %});
+{% if queryTail and row > 0 %}
+      }
+{% endif %}
+{% endfor %}
+    }
+    workgroupBarrier();
+  }
+}

build/webgpu/bench.json CHANGED Viewed

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

build/webgpu/manifest.json CHANGED Viewed

@@ -57,7 +57,8 @@
     "canPinSubgroupSize32": "device.features.has(\"subgroups\") and device.features.has(\"subgroup-size-control\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize <= 32 and device.adapterInfo.subgroupMaxSize >= 32",
     "pinSubgroupSize32": "canPinSubgroupSize32 and not wave32Adapter",
     "wave32Effective": "wave32Adapter or pinSubgroupSize32",
-    "headDimPlan": "dim(shapes.queryT, 2) / attrs.num_heads if (ranks.queryT == 3 and attrs.num_heads > 0) else 0",
     "qkvDtypesOk": "tensorDtypes.keyT == tensorDtypes.queryT and tensorDtypes.valueT == tensorDtypes.queryT and tensorDtypes.outputT == tensorDtypes.queryT",
     "floatDtypeOk": "(tensorDtypes.queryT == \"float32\" or tensorDtypes.queryT == \"float16\") and f16Ok(tensorDtypes.queryT)",
     "qkvShapeOk": "ranks.queryT == 3 and ranks.keyT == 3 and ranks.valueT == 3 and ranks.outputT == 3 and attrs.num_heads > 0 and dim(shapes.queryT, 2) % attrs.num_heads == 0 and dim(shapes.keyT, 2) == dim(shapes.queryT, 2) and dim(shapes.valueT, 2) == dim(shapes.queryT, 2) and dim(shapes.keyT, 1) == dim(shapes.valueT, 1) and dim(shapes.queryT, 0) == dim(shapes.keyT, 0) and dim(shapes.queryT, 0) == dim(shapes.valueT, 0) and dim(shapes.outputT, 0) == dim(shapes.queryT, 0) and dim(shapes.outputT, 1) == dim(shapes.queryT, 1) and dim(shapes.outputT, 2) == dim(shapes.valueT, 2)",
@@ -66,40 +67,42 @@
     "qkvContractOk": "qkvShapeOk and qkvDtypesOk and floatDtypeOk and noAttnBias",
     "qkvMaskContractOk": "qkvShapeOk and qkvDtypesOk and floatDtypeOk and attnBiasOk",
     "biasOk": "present.biasT and ranks.biasT == 1 and tensorDtypes.biasT == tensorDtypes.queryT and dim(shapes.biasT, 0) == 3 * dim(shapes.queryT, 2)",
-    "q32BroadcastSubgroupLanes": "device.adapterInfo.subgroupMinSize if subgroupsWave32 else 0",
-    "q32BroadcastF32HeadVectors": "headDimPlan / 4 if headDimPlan % 4 == 0 else 0",
-    "q32BroadcastF32RegisterGeometry": "subgroupsWave32 and q32BroadcastF32HeadVectors == q32BroadcastSubgroupLanes",
     "subgroupCluster4": "not narrowSubgroupRange and device.features.has(\"subgroups\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize >= 4 and device.adapterInfo.subgroupMinSize % 4 == 0 and device.adapterInfo.subgroupMaxSize % 4 == 0",
     "subgroupCluster8": "not narrowSubgroupRange and 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",
     "attentionDispatchFits": "dim(shapes.queryT, 0) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and dim(shapes.queryT, 1) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and attrs.num_heads <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
-    "flashHeadOk": "headDimPlan % 4 == 0 and headDimPlan >= 32 and headDimPlan <= 256",
     "flashSizeOk": "flashHeadOk and (dim(shapes.queryT, 1) * attrs.num_heads >= tunables.FLASH_MIN_QUERY_HEADS or (dim(shapes.queryT, 1) == 1 and dim(shapes.keyT, 1) >= 512) or (dim(shapes.queryT, 1) > 1 and dim(shapes.keyT, 1) >= 2048)) and attentionDispatchFits",
     "flashShapeOk": "qkvContractOk and flashSizeOk",
     "flashMaskShapeOk": "qkvMaskContractOk and flashSizeOk",
     "noBiasSplitKCount": "min(tunables.DECODE_MAX_SPLITS if dim(shapes.queryT, 1) == 1 else max(1, ceilDiv(tunables.SPLITK_TARGET_WORKGROUPS, dim(shapes.queryT, 0) * dim(shapes.queryT, 1) * attrs.num_heads)), ceilDiv(dim(shapes.keyT, 1), tunables.DECODE_KEYS_PER_SPLIT))",
     "biasSplitKCount": "min(tunables.DECODE_MAX_SPLITS, ceilDiv(dim(shapes.keyT, 1), tunables.DECODE_KEYS_PER_SPLIT))",
-    "noBiasPartialOutBytes": "dim(shapes.queryT, 0) * dim(shapes.queryT, 1) * attrs.num_heads * noBiasSplitKCount * headDimPlan * 4",
     "noBiasStatsBytes": "2 * dim(shapes.queryT, 0) * dim(shapes.queryT, 1) * attrs.num_heads * noBiasSplitKCount * 4",
     "noBiasSplitScratchFits": "noBiasPartialOutBytes <= device.limits.maxStorageBufferBindingSize and noBiasPartialOutBytes <= device.limits.maxBufferSize and noBiasStatsBytes <= device.limits.maxStorageBufferBindingSize and noBiasStatsBytes <= device.limits.maxBufferSize",
     "noBiasSplitDispatchFits": "dim(shapes.queryT, 1) * noBiasSplitKCount <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
-    "biasPartialOutBytes": "dim(shapes.queryT, 0) * attrs.num_heads * biasSplitKCount * headDimPlan * 4",
     "biasStatsBytes": "2 * dim(shapes.queryT, 0) * attrs.num_heads * biasSplitKCount * 4",
     "biasSplitScratchFits": "biasPartialOutBytes <= device.limits.maxStorageBufferBindingSize and biasPartialOutBytes <= device.limits.maxBufferSize and biasStatsBytes <= device.limits.maxStorageBufferBindingSize and biasStatsBytes <= device.limits.maxBufferSize",
     "biasSplitDispatchFits": "biasSplitKCount <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
     "decodeSplitKNoBiasOk": "qkvContractOk and attrs.unidirectional == 0 and dim(shapes.queryT, 1) == 1 and dim(shapes.keyT, 1) >= 512 and flashHeadOk and attentionDispatchFits and noBiasSplitDispatchFits and noBiasSplitScratchFits",
     "shortQuerySplitKNoBiasOk": "qkvContractOk and attrs.unidirectional == 0 and dim(shapes.queryT, 1) >= 2 and dim(shapes.queryT, 1) <= 16 and dim(shapes.keyT, 1) >= 2048 and flashHeadOk and attentionDispatchFits and noBiasSplitDispatchFits and noBiasSplitScratchFits",
     "decodeSplitKBiasOk": "biasOk and qkvContractOk and attrs.unidirectional == 0 and dim(shapes.queryT, 1) == 1 and dim(shapes.keyT, 1) >= 512 and flashHeadOk and attentionDispatchFits and biasSplitDispatchFits and biasSplitScratchFits",
-    "decodeSplitKPortablePreferred": "tensorDtypes.queryT == \"float32\" and device.features.has(\"subgroups\") and has(device.adapterInfo, \"subgroupMinSize\") and headDimPlan / 4 < device.adapterInfo.subgroupMinSize",
     "materializedScoreBytes": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1) * dim(shapes.keyT, 1) * 4",
     "materializedScoreFits": "materializedScoreBytes <= device.limits.maxStorageBufferBindingSize and materializedScoreBytes <= device.limits.maxBufferSize",
     "materializedWorkgroupSize": "tunables.MATERIALIZED_WORKGROUP_DIM * tunables.MATERIALIZED_WORKGROUP_DIM",
     "materializedScoreStorageBytes": "(tunables.MATERIALIZED_QUERY_TILE * tunables.MATERIALIZED_INNER_TILE + tunables.MATERIALIZED_KEY_TILE * (tunables.MATERIALIZED_INNER_TILE + 4)) * 4",
-    "materializedApplyTileN": "tunables.MATERIALIZED_VALUE_TILE_D128 if headDimPlan == 128 else tunables.MATERIALIZED_VALUE_TILE",
     "materializedApplyStorageBytes": "(tunables.MATERIALIZED_QUERY_TILE * tunables.MATERIALIZED_INNER_TILE + tunables.MATERIALIZED_INNER_TILE * materializedApplyTileN) * 4",
     "materializedTileGeometryOk": "tunables.MATERIALIZED_INNER_TILE % 4 == 0 and (materializedApplyTileN / tunables.MATERIALIZED_WORKGROUP_DIM) % 4 == 0 and tunables.MATERIALIZED_INNER_TILE > 0 and tunables.MATERIALIZED_WORKGROUP_DIM > 0 and tunables.MATERIALIZED_QUERY_TILE % tunables.MATERIALIZED_WORKGROUP_DIM == 0 and tunables.MATERIALIZED_KEY_TILE % tunables.MATERIALIZED_WORKGROUP_DIM == 0 and materializedApplyTileN % tunables.MATERIALIZED_WORKGROUP_DIM == 0 and tunables.MATERIALIZED_QUERY_TILE * tunables.MATERIALIZED_INNER_TILE >= materializedWorkgroupSize and tunables.MATERIALIZED_KEY_TILE * tunables.MATERIALIZED_INNER_TILE >= materializedWorkgroupSize and tunables.MATERIALIZED_INNER_TILE * materializedApplyTileN >= materializedWorkgroupSize and (tunables.MATERIALIZED_QUERY_TILE * tunables.MATERIALIZED_INNER_TILE) % materializedWorkgroupSize == 0 and (tunables.MATERIALIZED_KEY_TILE * tunables.MATERIALIZED_INNER_TILE) % materializedWorkgroupSize == 0 and (tunables.MATERIALIZED_INNER_TILE * materializedApplyTileN) % materializedWorkgroupSize == 0",
     "materializedDeviceOk": "materializedTileGeometryOk and materializedWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.MATERIALIZED_WORKGROUP_DIM <= device.limits.maxComputeWorkgroupSizeX and tunables.MATERIALIZED_WORKGROUP_DIM <= device.limits.maxComputeWorkgroupSizeY and materializedScoreStorageBytes <= device.limits.maxComputeWorkgroupStorageSize and materializedApplyStorageBytes <= device.limits.maxComputeWorkgroupStorageSize",
     "materializedWideSimdOk": "device.features.has(\"subgroups\") or (has(device.adapterInfo, \"subgroupMinSize\") and device.adapterInfo.subgroupMinSize >= 16)",
-    "materializedF32CoreOk": "qkvContractOk and tensorDtypes.queryT == \"float32\" and attrs.unidirectional == 0 and headDimPlan >= 64 and headDimPlan <= 128 and dim(shapes.queryT, 1) >= 512 and dim(shapes.keyT, 1) >= 512 and dim(shapes.queryT, 0) * attrs.num_heads <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(dim(shapes.queryT, 1), tunables.MATERIALIZED_QUERY_TILE) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(dim(shapes.keyT, 1), tunables.MATERIALIZED_KEY_TILE) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and materializedScoreFits and materializedDeviceOk",
     "materializedSoftmaxStorageBytes": "tunables.MATERIALIZED_SOFTMAX_WORKGROUP_SIZE * 8 + 8",
     "materializedSoftmaxResourcesFit": "tunables.MATERIALIZED_SOFTMAX_WORKGROUP_SIZE > 0 and pow2ceil(tunables.MATERIALIZED_SOFTMAX_WORKGROUP_SIZE) == tunables.MATERIALIZED_SOFTMAX_WORKGROUP_SIZE and tunables.MATERIALIZED_SOFTMAX_WORKGROUP_SIZE <= deviceWorkgroupCap and materializedSoftmaxStorageBytes <= device.limits.maxComputeWorkgroupStorageSize",
     "materializedSgmatQueryTile": "tunables.MATERIALIZED_SGMAT_QUERY_TILE",
@@ -123,9 +126,9 @@
     "materializedSgmatResourcesFit": "materializedSgmatGeometryOk and materializedSgmatWorkgroupSize <= deviceWorkgroupCap and materializedSgmatCompactStorageBytes <= device.limits.maxComputeWorkgroupStorageSize",
     "materializedSgmatDispatchFits": "ceilDiv(dim(shapes.keyT, 1), materializedSgmatKeyTile) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(dim(shapes.queryT, 1), materializedSgmatQueryTile) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and dim(shapes.queryT, 0) * attrs.num_heads <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1), min(device.limits.maxComputeWorkgroupsPerDimension, 65535)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
     "materializedSgmatDirectScoreStore": "dim(shapes.queryT, 1) % materializedSgmatQueryTile == 0 and dim(shapes.keyT, 1) % materializedSgmatKeyTile == 0",
-    "materializedSgmatDirectApplyStore": "dim(shapes.queryT, 1) % materializedSgmatQueryTile == 0 and headDimPlan % materializedSgmatKeyTile == 0",
     "materializedSgmatRuntimeDirectStore": "dim(shapes.queryT, 1) >= 2 * materializedSgmatQueryTile and dim(shapes.keyT, 1) >= 2 * materializedSgmatKeyTile",
-    "materializedSgmatCoreOk": "qkvContractOk and tensorDtypes.queryT == \"float32\" and attrs.unidirectional == 0 and headDimPlan >= 32 and headDimPlan <= 256 and dim(shapes.queryT, 1) >= 512 and dim(shapes.keyT, 1) >= 512 and device.features.has(\"subgroups\") and wave32Effective and device.features.has(\"chromium-experimental-subgroup-matrix\") and materializedSgmatBuffersFit and materializedSgmatResourcesFit and materializedSgmatDispatchFits",
     "materializedCachedSoftmaxWg": "min(tunables.MATERIALIZED_CACHED_SOFTMAX_WORKGROUP_SIZE, device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
     "materializedCachedSoftmaxVecsPerLane": "ceilDiv(ceilDiv(dim(shapes.keyT, 1), 4), max(1, materializedCachedSoftmaxWg))",
     "materializedCachedSoftmaxStorageBytes": "materializedCachedSoftmaxWg * 8 + 8",
@@ -135,41 +138,47 @@
     "materializedSgmatOk": "materializedSgmatCoreOk and materializedAdaptiveSoftmaxOk",
     "materializedSgmatFusedOk": "materializedSgmatCoreOk",
     "materializedF32Ok": "materializedF32CoreOk and materializedAdaptiveSoftmaxOk",
-    "clusterTileKWg64": "max(1, min(tunables.FLASH_MAX_TILE_K, floor(device.limits.maxComputeWorkgroupStorageSize / (headDimPlan * (8 if tensorDtypes.queryT != \"float16\" else 4) + tunables.FLASH_CLUSTER_WG_SMALL * 4))))",
-    "clusterTileKWg128": "max(1, min(tunables.FLASH_MAX_TILE_K, floor(device.limits.maxComputeWorkgroupStorageSize / (headDimPlan * (8 if tensorDtypes.queryT != \"float16\" else 4) + tunables.FLASH_CLUSTER_WG_LARGE * 4))))",
-    "smallHeadParallelOk": "qkvContractOk and headDimPlan < 32 and dim(shapes.keyT, 1) >= 64 and dim(shapes.keyT, 1) <= 2048 and attrs.unidirectional == 0 and dim(shapes.queryT, 0) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and attrs.num_heads <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
-    "prefillTiledStorageBytes": "headDimPlan * tunables.PREFILL_QUERY_TILE * 4",
     "prefillTiledDeviceOk": "tunables.PREFILL_QUERY_TILE <= deviceWorkgroupCap and prefillTiledStorageBytes <= device.limits.maxComputeWorkgroupStorageSize",
-    "portableWorkgroupSize": "min(tunables.WORKGROUP_SIZE, pow2ceil(max(1, headDimPlan)))",
-    "portableWorkgroupStorageBytes": "portableWorkgroupSize * 4 + max(1, headDimPlan) * 4 + 16",
     "portableWorkgroupOk": "tunables.WORKGROUP_SIZE > 0 and pow2ceil(tunables.WORKGROUP_SIZE) == tunables.WORKGROUP_SIZE and portableWorkgroupSize <= deviceWorkgroupCap and portableWorkgroupStorageBytes <= device.limits.maxComputeWorkgroupStorageSize",
     "fallbackShapeOk": "qkvContractOk and portableWorkgroupOk",
     "fallbackMaskShapeOk": "qkvMaskContractOk and portableWorkgroupOk",
-    "smallSeqShapeOk": "qkvContractOk and tensorDtypes.queryT == \"float32\" and dim(shapes.queryT, 1) >= 1 and dim(shapes.queryT, 1) <= tunables.SMALL_SEQ_MAX and dim(shapes.keyT, 1) >= 1 and dim(shapes.keyT, 1) <= tunables.SMALL_SEQ_MAX and headDimPlan >= 1 and attrs.unidirectional == 0",
-    "smallSeqPrivateFloats": "dim(shapes.keyT, 1) + headDimPlan",
     "smallSeqWorkgroupSize": "max(32, pow2ceil(dim(shapes.queryT, 1)))",
-    "smallSeqSharedBytes": "dim(shapes.keyT, 1) * headDimPlan * 8",
     "smallSeqResourcesFit": "smallSeqPrivateFloats <= tunables.SMALL_SEQ_MAX_PRIVATE_FLOATS and smallSeqWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup and smallSeqWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX and smallSeqSharedBytes <= device.limits.maxComputeWorkgroupStorageSize",
-    "smallSeqBlockedKvBytes": "dim(shapes.keyT, 1) * headDimPlan * 8",
-    "smallSeqBlockedLaneBytes": "8 + headDimPlan * 4",
     "smallSeqBlockedKeyLanes": "min(pow2ceil(dim(shapes.keyT, 1)), 16 if smallSeqBlockedKvBytes + tunables.SMALL_SEQ_BLOCKED_QUERY_BLOCK * 16 * smallSeqBlockedLaneBytes <= device.limits.maxComputeWorkgroupStorageSize else (8 if smallSeqBlockedKvBytes + tunables.SMALL_SEQ_BLOCKED_QUERY_BLOCK * 8 * smallSeqBlockedLaneBytes <= device.limits.maxComputeWorkgroupStorageSize else 4))",
     "smallSeqBlockedWorkgroupSize": "tunables.SMALL_SEQ_BLOCKED_QUERY_BLOCK * smallSeqBlockedKeyLanes",
     "smallSeqBlockedSharedBytes": "smallSeqBlockedKvBytes + smallSeqBlockedWorkgroupSize * smallSeqBlockedLaneBytes",
-    "smallSeqBlockedShapeOk": "qkvContractOk and tensorDtypes.queryT == \"float32\" and headDimPlan % 4 == 0 and headDimPlan >= 4 and headDimPlan <= tunables.SMALL_SEQ_BLOCKED_MAX_HEAD_DIM and dim(shapes.keyT, 1) >= 1 and dim(shapes.keyT, 1) <= tunables.SMALL_SEQ_BLOCKED_MAX_KV and dim(shapes.queryT, 1) >= 1",
     "smallSeqBlockedFits": "smallSeqBlockedWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup and smallSeqBlockedWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX and smallSeqBlockedSharedBytes <= device.limits.maxComputeWorkgroupStorageSize and ceilDiv(dim(shapes.queryT, 1), tunables.SMALL_SEQ_BLOCKED_QUERY_BLOCK) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and attrs.num_heads <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and dim(shapes.queryT, 0) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
     "smallSeqDispatchFits": "attrs.num_heads <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and dim(shapes.queryT, 0) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
-    "materializedSgmatCoreF16Ok": "qkvContractOk and tensorDtypes.queryT == \"float16\" and tensorDtypes.keyT == \"float16\" and tensorDtypes.valueT == \"float16\" and device.features.has(\"shader-f16\") and attrs.unidirectional == 0 and headDimPlan >= 32 and headDimPlan <= 256 and dim(shapes.queryT, 1) >= 512 and dim(shapes.keyT, 1) >= 512 and device.features.has(\"subgroups\") and wave32Effective and device.features.has(\"chromium-experimental-subgroup-matrix\") and materializedSgmatBuffersFit and materializedSgmatResourcesFit and materializedSgmatDispatchFits",
     "materializedSgmatFusedF16Ok": "materializedSgmatCoreF16Ok",
-    "attentionScaleExpression": "\"select(inverseSqrt(f32(HEAD_DIM)), params.scale, params.scale != 0.0)\""
   },
   "bindings": {
-    "query": { "arg": "queryT", "buffer": "read-only-storage", "elementType": "$inputElement" },
-    "key": { "arg": "keyT", "buffer": "read-only-storage", "elementType": "$inputElement" },
-    "value": { "arg": "valueT", "buffer": "read-only-storage", "elementType": "$inputElement" },
-    "bias": { "arg": "biasT", "buffer": "read-only-storage", "elementType": "$inputScalar" },
-    "output": { "arg": "outputT", "buffer": "storage", "elementType": "$outputElement" },
     "params": {
-      "buffer": "uniform",
       "struct": [
         { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" },
         { "name": "kvSeq", "type": "u32", "value": "dim(shapes.keyT, 1)" },
@@ -177,77 +186,67 @@
         { "name": "isCausal", "type": "u32", "value": "attrs.unidirectional" }
       ]
     },
-    "params_2": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
         { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" },
         { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" }
       ]
     },
-    "q": { "arg": "queryT", "buffer": "read-only-storage", "elementType": "$scalar" },
-    "k": { "arg": "keyT", "buffer": "read-only-storage", "elementType": "$scalar" },
-    "v": { "arg": "valueT", "buffer": "read-only-storage", "elementType": "$scalar" },
-    "y": { "arg": "outputT", "buffer": "storage", "elementType": "$scalar" },
-    "params_4": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
         { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" },
         { "name": "kvSeq", "type": "u32", "value": "dim(shapes.keyT, 1)" },
         { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" }
       ]
     },
-    "query_2": { "arg": "queryT", "name": "query", "buffer": "read-only-storage", "elementType": "f32" },
-    "key_2": { "arg": "keyT", "name": "key", "buffer": "read-only-storage", "elementType": "f32" },
-    "scores": { "scratch": "materializedScores", "buffer": "storage", "elementType": "f32" },
-    "scorePartials": { "scratch": "materializedScorePartials", "buffer": "storage", "elementType": "f32" },
-    "bias_2": { "arg": "biasT", "name": "bias", "buffer": "read-only-storage", "elementType": "f32" },
-    "scorePartials_2": {
       "scratch": "materializedScorePartials",
       "name": "scorePartials",
       "buffer": "read-only-storage",
       "elementType": "f32"
     },
-    "rowStats": { "scratch": "materializedRowStats", "buffer": "storage", "elementType": "f32" },
-    "params_6": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
         { "name": "rows", "type": "u32", "value": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)" }
       ]
     },
-    "scores_2": {
       "scratch": "materializedScores",
       "name": "scores",
       "buffer": "read-only-storage",
       "elementType": "f32"
     },
-    "value_2": { "arg": "valueT", "name": "value", "buffer": "read-only-storage", "elementType": "f32" },
-    "output_2": { "arg": "outputT", "name": "output", "buffer": "storage", "elementType": "f32" },
-    "rowStats_2": {
       "scratch": "materializedRowStats",
       "name": "rowStats",
       "buffer": "read-only-storage",
       "elementType": "f32"
     },
-    "params_7": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
         { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" },
         { "name": "kvSeq", "type": "u32", "value": "dim(shapes.keyT, 1)" }
       ]
     },
-    "attn_mask_2": {
-      "arg": "attentionBiasT",
-      "name": "attn_mask",
-      "buffer": "read-only-storage",
-      "elementType": "$maskElement"
-    },
-    "params_8": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
         { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" },
         { "name": "kvSeq", "type": "u32", "value": "dim(shapes.keyT, 1)" },
@@ -266,53 +265,109 @@
         { "name": "maskSeqStride", "type": "u32", "value": "dim(shapes.attentionBiasT, 3)" }
       ]
     },
-    "query_3": { "arg": "queryT", "name": "query", "buffer": "read-only-storage", "elementType": "$inputVec4" },
-    "key_3": { "arg": "keyT", "name": "key", "buffer": "read-only-storage", "elementType": "$inputVec4" },
-    "value_3": { "arg": "valueT", "name": "value", "buffer": "read-only-storage", "elementType": "$inputVec4" },
-    "partial_out": { "scratch": "partialOut", "buffer": "storage", "elementType": "vec4<f32>" },
-    "partial_stats": { "scratch": "partialStats", "buffer": "storage", "elementType": "vec2<f32>" },
-    "params_9": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
         { "name": "kvSeq", "type": "u32", "value": "dim(shapes.keyT, 1)" },
         { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" }
       ]
     },
-    "partial_out_2": {
       "scratch": "partialOut",
       "name": "partial_out",
       "buffer": "read-only-storage",
       "elementType": "vec4<f32>"
     },
-    "partial_stats_2": {
       "scratch": "partialStats",
       "name": "partial_stats",
       "buffer": "read-only-storage",
       "elementType": "vec2<f32>"
     },
-    "output_3": { "arg": "outputT", "name": "output", "buffer": "storage", "elementType": "$inputVec4" },
-    "scores_3": {
-      "scratch": "materializedScores",
-      "name": "scores",
-      "buffer": "storage",
-      "elementType": "$softmaxElementType"
-    }
   },
   "variants": [
     {
       "id": "qkv_bias_small_seq_blocked",
       "priority": 40,
       "when": ["biasOk", "smallSeqBlockedShapeOk", "not flashShapeOk", "smallSeqBlockedFits"],
       "derive": {
-        "usesF16": false,
-        "scalar": "\"f32\"",
         "inputElement": "\"vec4<f32>\"",
         "outputElement": "\"vec4<f32>\"",
         "inputScalar": "\"f32\"",
-        "hasBias": true,
-        "headDim": "headDimPlan",
-        "headDimV4": "headDimPlan / 4",
         "hidden": "dim(shapes.queryT, 2)",
         "hiddenV4": "dim(shapes.queryT, 2) / 4",
         "kvSeq": "dim(shapes.keyT, 1)",
@@ -329,12 +384,6 @@
             "x": "ceilDiv(dim(shapes.queryT, 1), tunables.SMALL_SEQ_BLOCKED_QUERY_BLOCK)",
             "y": "attrs.num_heads",
             "z": "dim(shapes.queryT, 0)"
-          },
-          "profile": {
-            "op": "\"multi_head_attention\"",
-            "variant": "\"qkv_bias_small_seq_blocked\"",
-            "numHeads": "attrs.num_heads",
-            "headDim": "headDim"
           }
         }
       ]
@@ -344,14 +393,9 @@
       "priority": 40,
       "when": ["not present.biasT", "smallSeqBlockedShapeOk", "not flashShapeOk", "smallSeqBlockedFits"],
       "derive": {
-        "usesF16": false,
-        "scalar": "\"f32\"",
         "inputElement": "\"vec4<f32>\"",
         "outputElement": "\"vec4<f32>\"",
-        "inputScalar": "\"f32\"",
-        "hasBias": false,
-        "headDim": "headDimPlan",
-        "headDimV4": "headDimPlan / 4",
         "hidden": "dim(shapes.queryT, 2)",
         "hiddenV4": "dim(shapes.queryT, 2) / 4",
         "kvSeq": "dim(shapes.keyT, 1)",
@@ -368,12 +412,6 @@
             "x": "ceilDiv(dim(shapes.queryT, 1), tunables.SMALL_SEQ_BLOCKED_QUERY_BLOCK)",
             "y": "attrs.num_heads",
             "z": "dim(shapes.queryT, 0)"
-          },
-          "profile": {
-            "op": "\"multi_head_attention\"",
-            "variant": "\"qkv_no_bias_small_seq_blocked\"",
-            "numHeads": "attrs.num_heads",
-            "headDim": "headDim"
           }
         }
       ]
@@ -383,7 +421,6 @@
       "priority": 60,
       "when": ["not present.biasT", "smallSeqShapeOk", "not flashShapeOk", "smallSeqResourcesFit", "smallSeqDispatchFits"],
       "derive": {
-        "usesF16": false,
         "inputElement": "\"f32\"",
         "outputElement": "\"f32\"",
         "headDim": "dim(shapes.queryT, 2) / attrs.num_heads",
@@ -396,25 +433,18 @@
           "id": "main",
           "name": "MultiHeadAttention",
           "shader": "mha-small-seq.wgsl.jinja",
-          "bindings": ["query", "key", "value", "output", "params_2"],
-          "dispatch": { "x": "attrs.num_heads", "y": "dim(shapes.queryT, 0)" },
-          "profile": {
-            "op": "\"multi_head_attention\"",
-            "variant": "\"qkv_no_bias_small_seq\"",
-            "numHeads": "attrs.num_heads",
-            "headDim": "dim(shapes.queryT, 2) / attrs.num_heads"
-          }
         }
       ]
     },
     {
       "id": "qkv_no_bias_tiled_nosg",
       "priority": 19,
-      "when": ["not present.biasT", "qkvContractOk", "not smallHeadParallelOk", "headDimPlan % 4 == 0", "headDimPlan <= 128", "dim(shapes.queryT, 1) >= 31", "prefillTiledDeviceOk"],
       "supersededBy": ["qkv_no_bias_flash_cluster_nosg", "qkv_no_bias_flash_cluster_lpq4_nosg"],
       "derive": {
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
-        "usesF16": "tensorDtypes.queryT == \"float16\"",
         "blockM": "tunables.PREFILL_QUERY_TILE",
         "vHeadCap": "dim(shapes.queryT, 2) / attrs.num_heads"
       },
@@ -449,8 +479,8 @@
             }
           ],
           "dispatch": {
-            "x": "min(ceilDiv((dim(shapes.outputT, 0) * attrs.num_heads * ceil(dim(shapes.outputT, 1) / 32) * 32), (32)), 65535)",
-            "y": "ceilDiv(ceilDiv((dim(shapes.outputT, 0) * attrs.num_heads * ceil(dim(shapes.outputT, 1) / 32) * 32), (32)), 65535)",
             "z": 1
           }
         }
@@ -459,14 +489,12 @@
     {
       "id": "qkv_bias_flash_q32_broadcast_f32_d128",
       "priority": 30,
-      "when": ["tensorDtypes.queryT == \"float32\"", "biasOk", "attrs.unidirectional == 0", "flashShapeOk", "q32BroadcastF32RegisterGeometry", "dim(shapes.queryT, 1) >= 31", "ceilDiv(dim(shapes.queryT, 1), 32) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "subgroupsWave32"],
       "requires": { "features": ["subgroups"] },
       "derive": {
-        "hasBias": true,
         "hasCausal": false,
         "usesF16": false,
         "scalar": "\"f32\"",
-        "inputVec4": "\"vec4<f32>\"",
         "inputElement": "\"vec4<f32>\"",
         "outputElement": "\"vec4<f32>\"",
         "inputScalar": "\"f32\"",
@@ -484,13 +512,12 @@
           "name": "MultiHeadAttention.FlashQ32BroadcastF32Bias",
           "shader": "attn-flash-q32-broadcast.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
-          "bindings": ["query", "key", "value", "bias", "output", "params_4"],
           "dispatch": {
             "x": "ceilDiv(dim(shapes.queryT, 1), 32)",
             "y": "attrs.num_heads",
             "z": "dim(shapes.queryT, 0)"
-          },
-          "subgroupCollectivesWidth": 32
         }
       ]
     },
@@ -499,8 +526,6 @@
       "priority": 10,
       "when": ["not present.biasT", "smallHeadParallelOk"],
       "derive": {
-        "usesF16": "tensorDtypes.queryT == \"float16\"",
-        "inputScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "outputScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "inputElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "outputElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
@@ -511,9 +536,9 @@
       "passes": [
         {
           "id": "main",
-          "name": "MultiHeadAttentionSmallHeadParallel",
           "shader": "attn-small-head-parallel.wgsl.jinja",
-          "bindings": ["query", "key", "value", "output", "params_2"],
           "dispatch": {
             "x": "min(dim(shapes.queryT, 1), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
             "y": "attrs.num_heads",
@@ -525,12 +550,9 @@
     {
       "id": "qkv_no_bias_tiled_attn_bias_nosg",
       "priority": 17,
-      "when": ["not present.biasT", "qkvMaskContractOk", "not smallHeadParallelOk", "headDimPlan % 4 == 0", "headDimPlan <= 128", "dim(shapes.queryT, 1) >= 31", "prefillTiledDeviceOk"],
       "derive": {
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
-        "usesF16": "tensorDtypes.queryT == \"float16\"",
-        "hasMask": true,
-        "maskIsBool": false,
         "blockM": "tunables.PREFILL_QUERY_TILE",
         "vHeadCap": "dim(shapes.queryT, 2) / attrs.num_heads"
       },
@@ -577,8 +599,8 @@
             }
           ],
           "dispatch": {
-            "x": "min(ceilDiv((dim(shapes.outputT, 0) * attrs.num_heads * ceil(dim(shapes.outputT, 1) / 32) * 32), (32)), 65535)",
-            "y": "ceilDiv(ceilDiv((dim(shapes.outputT, 0) * attrs.num_heads * ceil(dim(shapes.outputT, 1) / 32) * 32), (32)), 65535)",
             "z": 1
           }
         }
@@ -594,10 +616,7 @@
       },
       "derive": {
         "qNumHeads": "attrs.num_heads",
-        "headDim": "headDimPlan",
         "qHidden": "dim(shapes.queryT, 2)",
-        "hasBias": false,
-        "useSubgroups": true,
         "statSlots": "materializedSgmatStatSlots",
         "statQuerySeq": "dim(shapes.queryT, 1)"
       },
@@ -616,19 +635,18 @@
           "name": "MultiHeadAttention.MaterializedScoresSgmat",
           "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
           "derive": { "phase": "\"score\"", "emitRowStats": true },
-          "bindings": ["query_2", "key_2", "scores", "scorePartials", "params_4"],
           "dispatch": {
             "x": "ceilDiv(dim(shapes.keyT, 1), materializedSgmatKeyTile)",
             "y": "ceilDiv(dim(shapes.queryT, 1), materializedSgmatQueryTile)",
             "z": "dim(shapes.queryT, 0) * attrs.num_heads"
-          },
-          "subgroupCollectivesWidth": 32
         },
         {
           "id": "rowstats",
           "name": "MultiHeadAttention.MaterializedRowStatsCombine",
           "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja",
-          "bindings": ["scorePartials_2", "rowStats", "params_6"],
           "dispatch": {
             "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)), (materializedRowStatsWg)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
             "y": 1,
@@ -640,7 +658,7 @@
           "name": "MultiHeadAttention.MaterializedApplySgmat",
           "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
           "derive": { "phase": "\"apply\"", "fusedSoftmax": true },
-          "bindings": ["scores_2", "value_2", "output_2", "rowStats_2", "params_7"],
           "dispatch": {
             "x": "ceilDiv(headDim, materializedSgmatKeyTile)",
             "y": "ceilDiv(dim(shapes.queryT, 1), materializedSgmatQueryTile)",
@@ -659,10 +677,7 @@
       },
       "derive": {
         "qNumHeads": "attrs.num_heads",
-        "headDim": "headDimPlan",
         "qHidden": "dim(shapes.queryT, 2)",
-        "hasBias": true,
-        "useSubgroups": true,
         "statSlots": "materializedSgmatStatSlots",
         "statQuerySeq": "dim(shapes.queryT, 1)"
       },
@@ -681,19 +696,18 @@
           "name": "MultiHeadAttention.MaterializedScoresSgmatBias",
           "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
           "derive": { "phase": "\"score\"", "emitRowStats": true },
-          "bindings": ["query_2", "key_2", "bias_2", "scores", "scorePartials", "params_4"],
           "dispatch": {
             "x": "ceilDiv(dim(shapes.keyT, 1), materializedSgmatKeyTile)",
             "y": "ceilDiv(dim(shapes.queryT, 1), materializedSgmatQueryTile)",
             "z": "dim(shapes.queryT, 0) * attrs.num_heads"
-          },
-          "subgroupCollectivesWidth": 32
         },
         {
           "id": "rowstats",
           "name": "MultiHeadAttention.MaterializedRowStatsCombineBias",
           "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja",
-          "bindings": ["scorePartials_2", "rowStats", "params_6"],
           "dispatch": {
             "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)), (materializedRowStatsWg)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
             "y": 1,
@@ -705,7 +719,7 @@
           "name": "MultiHeadAttention.MaterializedApplySgmatBias",
           "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
           "derive": { "phase": "\"apply\"", "fusedSoftmax": true },
-          "bindings": ["scores_2", "value_2", "bias_2", "output_2", "rowStats_2", "params_7"],
           "dispatch": {
             "x": "ceilDiv(headDim, materializedSgmatKeyTile)",
             "y": "ceilDiv(dim(shapes.queryT, 1), materializedSgmatQueryTile)",
@@ -724,10 +738,7 @@
       },
       "derive": {
         "qNumHeads": "attrs.num_heads",
-        "headDim": "headDimPlan",
         "qHidden": "dim(shapes.queryT, 2)",
-        "hasBias": false,
-        "useSubgroups": true,
         "statSlots": "materializedSgmatStatSlots",
         "statQuerySeq": "dim(shapes.queryT, 1)",
         "operandF16": true,
@@ -749,19 +760,18 @@
           "name": "MultiHeadAttention.MaterializedScoresSgmatF16",
           "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
           "derive": { "phase": "\"score\"", "emitRowStats": true },
-          "bindings": ["query", "key", "scores", "scorePartials", "params_4"],
           "dispatch": {
             "x": "ceilDiv(dim(shapes.keyT, 1), materializedSgmatKeyTile)",
             "y": "ceilDiv(dim(shapes.queryT, 1), materializedSgmatQueryTile)",
             "z": "dim(shapes.queryT, 0) * attrs.num_heads"
-          },
-          "subgroupCollectivesWidth": 32
         },
         {
           "id": "rowstats",
           "name": "MultiHeadAttention.MaterializedRowStatsCombine",
           "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja",
-          "bindings": ["scorePartials_2", "rowStats", "params_6"],
           "dispatch": {
             "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)), (materializedRowStatsWg)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
             "y": 1,
@@ -773,7 +783,7 @@
           "name": "MultiHeadAttention.MaterializedApplySgmatF16",
           "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
           "derive": { "phase": "\"apply\"", "fusedSoftmax": true },
-          "bindings": ["scores_2", "value", "output", "rowStats_2", "params_7"],
           "dispatch": {
             "x": "ceilDiv(headDim, materializedSgmatKeyTile)",
             "y": "ceilDiv(dim(shapes.queryT, 1), materializedSgmatQueryTile)",
@@ -785,15 +795,12 @@
     {
       "id": "qkv_no_bias_flash_cluster_lpq4_nosg",
       "priority": 20,
-      "when": ["not present.biasT", "flashShapeOk", "headDimPlan % 16 == 0", "headDimPlan % 32 != 0", "dim(shapes.queryT, 1) >= 31"],
       "requires": {},
       "derive": {
-        "hasBias": false,
         "hasCausal": true,
-        "hasWindow": false,
         "usesF16": "tensorDtypes.queryT == \"float16\"",
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
-        "inputVec4": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "inputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "outputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "qNumHeads": "attrs.num_heads",
@@ -826,15 +833,12 @@
     {
       "id": "qkv_no_bias_flash_cluster_nosg",
       "priority": 20,
-      "when": ["not present.biasT", "flashShapeOk", "headDimPlan % 32 == 0", "dim(shapes.queryT, 1) >= 31"],
       "requires": {},
       "derive": {
-        "hasBias": false,
         "hasCausal": true,
-        "hasWindow": false,
         "usesF16": "tensorDtypes.queryT == \"float16\"",
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
-        "inputVec4": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "inputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "outputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "qNumHeads": "attrs.num_heads",
@@ -867,15 +871,12 @@
     {
       "id": "qkv_bias_flash_cluster_nosg",
       "priority": 19,
-      "when": ["biasOk", "flashShapeOk", "headDimPlan % 32 == 0", "dim(shapes.queryT, 1) >= 31"],
       "requires": {},
       "derive": {
-        "hasBias": true,
         "hasCausal": true,
-        "hasWindow": false,
         "usesF16": "tensorDtypes.queryT == \"float16\"",
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
-        "inputVec4": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "inputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "outputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "qNumHeads": "attrs.num_heads",
@@ -910,15 +911,12 @@
     {
       "id": "qkv_no_bias_flash_cluster_lpq4",
       "priority": 22,
-      "when": ["not present.biasT", "flashShapeOk", "headDimPlan % 16 == 0", "headDimPlan % 32 != 0", "dim(shapes.queryT, 1) >= 31", "subgroupCluster4"],
       "requires": { "features": ["subgroups"] },
       "derive": {
-        "hasBias": false,
         "hasCausal": true,
-        "hasWindow": false,
         "usesF16": "tensorDtypes.queryT == \"float16\"",
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
-        "inputVec4": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "inputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "outputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "qNumHeads": "attrs.num_heads",
@@ -942,23 +940,19 @@
             "x": "ceilDiv(dim(shapes.queryT, 1), TILE_Q)",
             "y": "attrs.num_heads",
             "z": "dim(shapes.queryT, 0)"
-          },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "qkv_no_bias_flash_cluster",
       "priority": 22,
-      "when": ["not present.biasT", "flashShapeOk", "headDimPlan % 32 == 0", "dim(shapes.queryT, 1) >= 31", "subgroupCluster8"],
       "requires": { "features": ["subgroups"] },
       "derive": {
-        "hasBias": false,
         "hasCausal": true,
-        "hasWindow": false,
         "usesF16": "tensorDtypes.queryT == \"float16\"",
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
-        "inputVec4": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "inputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "outputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "qNumHeads": "attrs.num_heads",
@@ -982,23 +976,19 @@
             "x": "ceilDiv(dim(shapes.queryT, 1), TILE_Q)",
             "y": "attrs.num_heads",
             "z": "dim(shapes.queryT, 0)"
-          },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "qkv_bias_flash_cluster",
       "priority": 21,
-      "when": ["biasOk", "flashShapeOk", "headDimPlan % 32 == 0", "dim(shapes.queryT, 1) >= 31", "subgroupCluster8"],
       "requires": { "features": ["subgroups"] },
       "derive": {
-        "hasBias": true,
         "hasCausal": true,
-        "hasWindow": false,
         "usesF16": "tensorDtypes.queryT == \"float16\"",
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
-        "inputVec4": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "inputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "outputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "qNumHeads": "attrs.num_heads",
@@ -1024,25 +1014,19 @@
             "x": "ceilDiv(dim(shapes.queryT, 1), TILE_Q)",
             "y": "attrs.num_heads",
             "z": "dim(shapes.queryT, 0)"
-          },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "qkv_no_bias_flash_cluster_attn_bias",
       "priority": 22,
-      "when": ["not present.biasT", "flashMaskShapeOk", "headDimPlan % 32 == 0", "dim(shapes.queryT, 1) >= 31", "subgroupCluster8"],
       "requires": { "features": ["subgroups"] },
       "derive": {
-        "hasBias": false,
         "hasCausal": true,
-        "hasWindow": false,
-        "hasMask": true,
-        "maskIsBool": false,
         "usesF16": "tensorDtypes.queryT == \"float16\"",
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
-        "inputVec4": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "inputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "outputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "maskElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
@@ -1062,30 +1046,24 @@
           "name": "MultiHeadAttention.Flash",
           "shader": "attn-flash-prefill-cluster.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
-          "bindings": ["query", "key", "value", "attn_mask_2", "output", "params_8"],
           "dispatch": {
             "x": "ceilDiv(dim(shapes.queryT, 1), TILE_Q)",
             "y": "attrs.num_heads",
             "z": "dim(shapes.queryT, 0)"
-          },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "qkv_bias_flash_cluster_attn_bias",
       "priority": 21,
-      "when": ["biasOk", "flashMaskShapeOk", "headDimPlan % 32 == 0", "dim(shapes.queryT, 1) >= 31", "subgroupCluster8"],
       "requires": { "features": ["subgroups"] },
       "derive": {
-        "hasBias": true,
         "hasCausal": true,
-        "hasWindow": false,
-        "hasMask": true,
-        "maskIsBool": false,
         "usesF16": "tensorDtypes.queryT == \"float16\"",
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
-        "inputVec4": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "inputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "outputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "maskElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
@@ -1107,13 +1085,12 @@
           "name": "MultiHeadAttention.Flash",
           "shader": "attn-flash-prefill-cluster.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
-          "bindings": ["query", "key", "value", "attn_mask_2", "bias", "output", "params_8"],
           "dispatch": {
             "x": "ceilDiv(dim(shapes.queryT, 1), TILE_Q)",
             "y": "attrs.num_heads",
             "z": "dim(shapes.queryT, 0)"
-          },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
@@ -1123,9 +1100,7 @@
       "when": ["not present.biasT", "decodeSplitKNoBiasOk or shortQuerySplitKNoBiasOk"],
       "requires": {},
       "derive": {
-        "combineSubgroups": false,
         "useSubgroups": false,
-        "hasWindow": false,
         "splitQueries": true,
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "inputVec4": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
@@ -1157,7 +1132,7 @@
           "name": "MultiHeadAttention.DecodeSplitKNoSg",
           "shader": "attn-flash-decode-splitk.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
-          "bindings": ["query_3", "key_3", "value_3", "partial_out", "partial_stats", "params_9"],
           "dispatch": {
             "x": "dim(shapes.queryT, 1) * noBiasSplitKCount",
             "y": "attrs.num_heads",
@@ -1169,7 +1144,7 @@
           "name": "MultiHeadAttention.DecodeSplitKMergeNoSg",
           "shader": "attn-flash-decode-splitk-merge.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
-          "bindings": ["partial_out_2", "partial_stats_2", "output_3"],
           "dispatch": { "x": "dim(shapes.queryT, 1)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
         }
       ]
@@ -1180,10 +1155,7 @@
       "when": ["biasOk", "decodeSplitKBiasOk"],
       "requires": {},
       "derive": {
-        "combineSubgroups": false,
         "useSubgroups": false,
-        "hasBias": true,
-        "hasWindow": false,
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "inputVec4": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "inputScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
@@ -1215,7 +1187,7 @@
           "name": "MultiHeadAttention.DecodeSplitKBiasNoSg",
           "shader": "attn-flash-decode-splitk.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
-          "bindings": ["query_3", "key_3", "value_3", "bias", "partial_out", "partial_stats", "params_9"],
           "dispatch": { "x": "biasSplitKCount", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
         },
         {
@@ -1223,7 +1195,7 @@
           "name": "MultiHeadAttention.DecodeSplitKMergeBiasNoSg",
           "shader": "attn-flash-decode-splitk-merge.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
-          "bindings": ["partial_out_2", "partial_stats_2", "bias", "output_3"],
           "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
         }
       ]
@@ -1235,8 +1207,6 @@
       "demoteWhen": ["decodeSplitKPortablePreferred"],
       "requires": { "features": ["subgroups"] },
       "derive": {
-        "combineSubgroups": true,
-        "hasWindow": false,
         "splitQueries": true,
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "inputVec4": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
@@ -1268,20 +1238,19 @@
           "name": "MultiHeadAttention.DecodeSplitK",
           "shader": "attn-flash-decode-splitk.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
-          "bindings": ["query_3", "key_3", "value_3", "partial_out", "partial_stats", "params_9"],
           "dispatch": {
             "x": "dim(shapes.queryT, 1) * noBiasSplitKCount",
             "y": "attrs.num_heads",
             "z": "dim(shapes.queryT, 0)"
-          },
-          "subgroupCollectivesWidth": "portable"
         },
         {
           "id": "merge",
           "name": "MultiHeadAttention.DecodeSplitKMerge",
           "shader": "attn-flash-decode-splitk-merge.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
-          "bindings": ["partial_out_2", "partial_stats_2", "output_3"],
           "dispatch": { "x": "dim(shapes.queryT, 1)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
         }
       ]
@@ -1292,9 +1261,6 @@
       "when": ["biasOk", "decodeSplitKBiasOk"],
       "requires": { "features": ["subgroups"] },
       "derive": {
-        "combineSubgroups": true,
-        "hasBias": true,
-        "hasWindow": false,
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "inputVec4": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "inputScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
@@ -1326,16 +1292,15 @@
           "name": "MultiHeadAttention.DecodeSplitKBias",
           "shader": "attn-flash-decode-splitk.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
-          "bindings": ["query_3", "key_3", "value_3", "bias", "partial_out", "partial_stats", "params_9"],
-          "dispatch": { "x": "biasSplitKCount", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" },
-          "subgroupCollectivesWidth": "portable"
         },
         {
           "id": "merge",
           "name": "MultiHeadAttention.DecodeSplitKMergeBias",
           "shader": "attn-flash-decode-splitk-merge.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
-          "bindings": ["partial_out_2", "partial_stats_2", "bias", "output_3"],
           "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
         }
       ]
@@ -1348,7 +1313,6 @@
       "derive": {
         "qNumHeads": "attrs.num_heads",
         "kvNumHeads": "attrs.num_heads",
-        "headDim": "headDimPlan",
         "materializedInnerTile": "tunables.MATERIALIZED_INNER_TILE",
         "materializedQueryTile": "tunables.MATERIALIZED_QUERY_TILE",
         "materializedKeyTile": "tunables.MATERIALIZED_KEY_TILE",
@@ -1358,7 +1322,6 @@
         "materializedSoftmaxCols4": "dim(shapes.keyT, 1) / 4",
         "softmaxElementType": "\"vec4<f32>\" if materializedCachedSoftmaxOk else \"f32\"",
         "applyTileN": "materializedApplyTileN",
-        "hasBias": false,
         "useSubgroups": "device.features.has(\"subgroups\")"
       },
       "intermediates": [
@@ -1374,7 +1337,7 @@
           "name": "MultiHeadAttention.MaterializedScores",
           "shader": "attn-materialized-score-f32.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
-          "bindings": ["query_2", "key_2", "scores", "params_4"],
           "dispatch": {
             "x": "ceilDiv(dim(shapes.keyT, 1), tunables.MATERIALIZED_KEY_TILE)",
             "y": "ceilDiv(dim(shapes.queryT, 1), tunables.MATERIALIZED_QUERY_TILE)",
@@ -1386,20 +1349,19 @@
           "name": "MultiHeadAttention.MaterializedSoftmax",
           "shader": "attn-materialized-softmax-f32.wgsl.jinja",
           "derive": { "cacheVec4": "materializedCachedSoftmaxOk" },
-          "bindings": ["scores_3", "params_6"],
           "dispatch": {
             "x": "min(dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1), 65535)",
             "y": "ceilDiv(dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1), 65535)",
             "z": 1
-          },
-          "subgroupCollectivesWidth": "portable"
         },
         {
           "id": "apply",
           "name": "MultiHeadAttention.MaterializedApply",
           "shader": "attn-materialized-apply-f32.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
-          "bindings": ["scores_2", "value_2", "output_2", "params_7"],
           "dispatch": {
             "x": "ceilDiv(headDim, applyTileN)",
             "y": "ceilDiv(dim(shapes.queryT, 1), tunables.MATERIALIZED_QUERY_TILE)",
@@ -1415,7 +1377,6 @@
       "derive": {
         "qNumHeads": "attrs.num_heads",
         "kvNumHeads": "attrs.num_heads",
-        "headDim": "headDimPlan",
         "materializedInnerTile": "tunables.MATERIALIZED_INNER_TILE",
         "materializedQueryTile": "tunables.MATERIALIZED_QUERY_TILE",
         "materializedKeyTile": "tunables.MATERIALIZED_KEY_TILE",
@@ -1425,7 +1386,6 @@
         "materializedSoftmaxCols4": "dim(shapes.keyT, 1) / 4",
         "softmaxElementType": "\"vec4<f32>\" if materializedCachedSoftmaxOk else \"f32\"",
         "applyTileN": "materializedApplyTileN",
-        "hasBias": true,
         "useSubgroups": "device.features.has(\"subgroups\")"
       },
       "intermediates": [
@@ -1441,7 +1401,7 @@
           "name": "MultiHeadAttention.MaterializedScoresBias",
           "shader": "attn-materialized-score-f32.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
-          "bindings": ["query_2", "key_2", "bias_2", "scores", "params_4"],
           "dispatch": {
             "x": "ceilDiv(dim(shapes.keyT, 1), tunables.MATERIALIZED_KEY_TILE)",
             "y": "ceilDiv(dim(shapes.queryT, 1), tunables.MATERIALIZED_QUERY_TILE)",
@@ -1453,20 +1413,19 @@
           "name": "MultiHeadAttention.MaterializedSoftmaxBias",
           "shader": "attn-materialized-softmax-f32.wgsl.jinja",
           "derive": { "cacheVec4": "materializedCachedSoftmaxOk" },
-          "bindings": ["scores_3", "params_6"],
           "dispatch": {
             "x": "min(dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1), 65535)",
             "y": "ceilDiv(dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1), 65535)",
             "z": 1
-          },
-          "subgroupCollectivesWidth": "portable"
         },
         {
           "id": "apply",
           "name": "MultiHeadAttention.MaterializedApplyBias",
           "shader": "attn-materialized-apply-f32.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
-          "bindings": ["scores_2", "value_2", "bias_2", "output_2", "params_7"],
           "dispatch": {
             "x": "ceilDiv(headDim, applyTileN)",
             "y": "ceilDiv(dim(shapes.queryT, 1), tunables.MATERIALIZED_QUERY_TILE)",
@@ -1483,14 +1442,11 @@
       "derive": {
         "qNumHeads": "attrs.num_heads",
         "kvNumHeads": "attrs.num_heads",
-        "headDim": "headDimPlan",
         "materializedInnerTile": "tunables.MATERIALIZED_INNER_TILE",
         "materializedQueryTile": "tunables.MATERIALIZED_QUERY_TILE",
         "materializedKeyTile": "tunables.MATERIALIZED_KEY_TILE",
         "materializedWorkgroupDim": "tunables.MATERIALIZED_WORKGROUP_DIM",
         "applyTileN": "materializedApplyTileN",
-        "hasBias": false,
-        "useSubgroups": "device.features.has(\"subgroups\")",
         "statSlots": "materializedGemmStatSlots",
         "statQuerySeq": "dim(shapes.queryT, 1)"
       },
@@ -1509,7 +1465,7 @@
           "name": "MultiHeadAttention.MaterializedScoresFused",
           "shader": "attn-materialized-score-f32.wgsl.jinja",
           "derive": { "layout": "\"bsh\"", "emitRowStats": true },
-          "bindings": ["query_2", "key_2", "scores", "scorePartials", "params_4"],
           "dispatch": {
             "x": "ceilDiv(dim(shapes.keyT, 1), tunables.MATERIALIZED_KEY_TILE)",
             "y": "ceilDiv(dim(shapes.queryT, 1), tunables.MATERIALIZED_QUERY_TILE)",
@@ -1521,7 +1477,7 @@
           "name": "MultiHeadAttention.MaterializedGemmRowStatsCombine",
           "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja",
           "derive": { "maxOnly": true },
-          "bindings": ["scorePartials_2", "rowStats", "params_6"],
           "dispatch": {
             "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)), (materializedRowStatsWg)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
             "y": 1,
@@ -1533,7 +1489,7 @@
           "name": "MultiHeadAttention.MaterializedApplyFused",
           "shader": "attn-materialized-apply-f32.wgsl.jinja",
           "derive": { "layout": "\"bsh\"", "fusedSoftmax": true },
-          "bindings": ["scores_2", "value_2", "output_2", "rowStats_2", "params_7"],
           "dispatch": {
             "x": "ceilDiv(headDim, applyTileN)",
             "y": "ceilDiv(dim(shapes.queryT, 1), tunables.MATERIALIZED_QUERY_TILE)",
@@ -1549,14 +1505,11 @@
       "derive": {
         "qNumHeads": "attrs.num_heads",
         "kvNumHeads": "attrs.num_heads",
-        "headDim": "headDimPlan",
         "materializedInnerTile": "tunables.MATERIALIZED_INNER_TILE",
         "materializedQueryTile": "tunables.MATERIALIZED_QUERY_TILE",
         "materializedKeyTile": "tunables.MATERIALIZED_KEY_TILE",
         "materializedWorkgroupDim": "tunables.MATERIALIZED_WORKGROUP_DIM",
         "applyTileN": "materializedApplyTileN",
-        "hasBias": true,
-        "useSubgroups": "device.features.has(\"subgroups\")",
         "statSlots": "materializedGemmStatSlots",
         "statQuerySeq": "dim(shapes.queryT, 1)"
       },
@@ -1575,7 +1528,7 @@
           "name": "MultiHeadAttention.MaterializedScoresBiasFused",
           "shader": "attn-materialized-score-f32.wgsl.jinja",
           "derive": { "layout": "\"bsh\"", "emitRowStats": true },
-          "bindings": ["query_2", "key_2", "bias_2", "scores", "scorePartials", "params_4"],
           "dispatch": {
             "x": "ceilDiv(dim(shapes.keyT, 1), tunables.MATERIALIZED_KEY_TILE)",
             "y": "ceilDiv(dim(shapes.queryT, 1), tunables.MATERIALIZED_QUERY_TILE)",
@@ -1587,7 +1540,7 @@
           "name": "MultiHeadAttention.MaterializedGemmRowStatsCombineBias",
           "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja",
           "derive": { "maxOnly": true },
-          "bindings": ["scorePartials_2", "rowStats", "params_6"],
           "dispatch": {
             "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)), (materializedRowStatsWg)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
             "y": 1,
@@ -1599,7 +1552,7 @@
           "name": "MultiHeadAttention.MaterializedApplyBiasFused",
           "shader": "attn-materialized-apply-f32.wgsl.jinja",
           "derive": { "layout": "\"bsh\"", "fusedSoftmax": true },
-          "bindings": ["scores_2", "value_2", "bias_2", "output_2", "rowStats_2", "params_7"],
           "dispatch": {
             "x": "ceilDiv(headDim, applyTileN)",
             "y": "ceilDiv(dim(shapes.queryT, 1), tunables.MATERIALIZED_QUERY_TILE)",
@@ -1611,12 +1564,11 @@
     {
       "id": "qkv_no_bias_flash_q32_broadcast",
       "priority": 30,
-      "when": ["tensorDtypes.queryT == \"float16\"", "not present.biasT", "flashShapeOk", "(dim(shapes.queryT, 2) / attrs.num_heads) % 4 == 0", "(dim(shapes.queryT, 2) / attrs.num_heads) >= 64", "(dim(shapes.queryT, 2) / attrs.num_heads) <= 256", "dim(shapes.queryT, 1) >= 31", "subgroupsWave32"],
       "requires": { "features": ["subgroups", "shader-f16"] },
       "derive": {
         "usesF16": true,
         "scalar": "\"f16\"",
-        "inputVec4": "\"vec4<f16>\"",
         "inputElement": "\"vec4<f16>\"",
         "outputElement": "\"vec4<f16>\"",
         "qNumHeads": "attrs.num_heads",
@@ -1637,8 +1589,7 @@
             "x": "ceilDiv(dim(shapes.queryT, 1), 32)",
             "y": "attrs.num_heads",
             "z": "dim(shapes.queryT, 0)"
-          },
-          "subgroupCollectivesWidth": 32
         }
       ]
     },
@@ -1650,7 +1601,6 @@
       "derive": {
         "usesF16": true,
         "scalar": "\"f16\"",
-        "inputVec4": "\"vec4<f16>\"",
         "inputElement": "\"vec4<f16>\"",
         "outputElement": "\"vec4<f16>\"",
         "qNumHeads": "attrs.num_heads",
@@ -1681,14 +1631,8 @@
       "priority": 0,
       "when": ["not present.biasT and fallbackMaskShapeOk"],
       "derive": {
-        "headsFromParams": false,
-        "hasBias": false,
         "hasCausal": true,
-        "hasWindow": false,
         "hasKeyLimit": false,
-        "hasMask": true,
-        "scaleFallbackRsqrt": true,
-        "usesF16": "tensorDtypes.queryT == \"float16\"",
         "inputScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "outputScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "inputElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
@@ -1707,7 +1651,7 @@
           "name": "MultiHeadAttention",
           "shader": "attn-online-scalar.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
-          "bindings": ["query", "key", "value", "attn_mask_2", "output", "params_8"],
           "dispatch": {
             "x": "min(dim(shapes.queryT, 1), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
             "y": "attrs.num_heads",
@@ -1721,14 +1665,8 @@
       "priority": 0,
       "when": ["biasOk and fallbackMaskShapeOk"],
       "derive": {
-        "headsFromParams": false,
-        "hasBias": true,
         "hasCausal": true,
-        "hasWindow": false,
         "hasKeyLimit": false,
-        "hasMask": true,
-        "scaleFallbackRsqrt": true,
-        "usesF16": "tensorDtypes.queryT == \"float16\"",
         "inputScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "outputScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "inputElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
@@ -1747,7 +1685,7 @@
           "name": "MultiHeadAttention",
           "shader": "attn-online-scalar.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
-          "bindings": ["query", "key", "value", "attn_mask_2", "bias", "output", "params_8"],
           "dispatch": {
             "x": "min(dim(shapes.queryT, 1), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
             "y": "attrs.num_heads",
@@ -1761,13 +1699,8 @@
       "priority": 0,
       "when": ["not present.biasT and fallbackShapeOk and not flashShapeOk"],
       "derive": {
-        "headsFromParams": false,
-        "hasBias": false,
         "hasCausal": true,
-        "hasWindow": false,
         "hasKeyLimit": false,
-        "scaleFallbackRsqrt": true,
-        "usesF16": "tensorDtypes.queryT == \"float16\"",
         "inputScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "outputScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "inputElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
@@ -1799,13 +1732,8 @@
       "priority": 0,
       "when": ["biasOk and fallbackShapeOk and not flashShapeOk"],
       "derive": {
-        "headsFromParams": false,
-        "hasBias": true,
         "hasCausal": true,
-        "hasWindow": false,
         "hasKeyLimit": false,
-        "scaleFallbackRsqrt": true,
-        "usesF16": "tensorDtypes.queryT == \"float16\"",
         "inputScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "outputScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "inputElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
@@ -1837,17 +1765,10 @@
       "priority": 20,
       "when": ["device.features.has(\"subgroups\")", "not present.biasT", "flashShapeOk"],
       "derive": {
-        "headsFromParams": false,
-        "hasBias": false,
         "hasCausal": true,
-        "hasWindow": false,
         "combineSubgroups": true,
-        "hasMask": false,
-        "maskIsBool": false,
         "maskElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
-        "usesF16": "tensorDtypes.queryT == \"float16\"",
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
-        "inputVec4": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "inputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "outputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "qNumHeads": "attrs.num_heads",
@@ -1864,8 +1785,7 @@
           "shader": "attn-flash-online.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
           "bindings": ["query", "key", "value", "output", "params"],
-          "dispatch": { "x": "dim(shapes.queryT, 1)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
@@ -1874,17 +1794,10 @@
       "priority": 20,
       "when": ["device.features.has(\"subgroups\")", "biasOk", "flashShapeOk"],
       "derive": {
-        "headsFromParams": false,
-        "hasBias": true,
         "hasCausal": true,
-        "hasWindow": false,
         "combineSubgroups": true,
-        "hasMask": false,
-        "maskIsBool": false,
         "maskElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
-        "usesF16": "tensorDtypes.queryT == \"float16\"",
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
-        "inputVec4": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "inputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "outputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "inputScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
@@ -1903,8 +1816,7 @@
           "shader": "attn-flash-online.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
           "bindings": ["query", "key", "value", "bias", "output", "params"],
-          "dispatch": { "x": "dim(shapes.queryT, 1)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
@@ -1913,17 +1825,10 @@
       "priority": 18,
       "when": ["true", "not present.biasT", "flashShapeOk"],
       "derive": {
-        "headsFromParams": false,
-        "hasBias": false,
         "hasCausal": true,
-        "hasWindow": false,
         "combineSubgroups": false,
-        "hasMask": false,
-        "maskIsBool": false,
         "maskElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
-        "usesF16": "tensorDtypes.queryT == \"float16\"",
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
-        "inputVec4": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "inputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "outputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "qNumHeads": "attrs.num_heads",
@@ -1949,17 +1854,10 @@
       "priority": 17,
       "when": ["true", "biasOk", "flashShapeOk"],
       "derive": {
-        "headsFromParams": false,
-        "hasBias": true,
         "hasCausal": true,
-        "hasWindow": false,
         "combineSubgroups": false,
-        "hasMask": false,
-        "maskIsBool": false,
         "maskElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
-        "usesF16": "tensorDtypes.queryT == \"float16\"",
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
-        "inputVec4": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "inputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "outputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "inputScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
@@ -1987,17 +1885,10 @@
       "priority": 20,
       "when": ["device.features.has(\"subgroups\")", "not present.biasT", "flashMaskShapeOk"],
       "derive": {
-        "headsFromParams": false,
-        "hasBias": false,
         "hasCausal": true,
-        "hasWindow": false,
         "combineSubgroups": true,
-        "hasMask": true,
-        "maskIsBool": false,
         "maskElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
-        "usesF16": "tensorDtypes.queryT == \"float16\"",
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
-        "inputVec4": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "inputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "outputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "qNumHeads": "attrs.num_heads",
@@ -2013,9 +1904,8 @@
           "name": "MultiHeadAttention.Flash",
           "shader": "attn-flash-online.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
-          "bindings": ["query", "key", "value", "attn_mask_2", "output", "params_8"],
-          "dispatch": { "x": "dim(shapes.queryT, 1)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
@@ -2024,17 +1914,10 @@
       "priority": 20,
       "when": ["device.features.has(\"subgroups\")", "biasOk", "flashMaskShapeOk"],
       "derive": {
-        "headsFromParams": false,
-        "hasBias": true,
         "hasCausal": true,
-        "hasWindow": false,
         "combineSubgroups": true,
-        "hasMask": true,
-        "maskIsBool": false,
         "maskElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
-        "usesF16": "tensorDtypes.queryT == \"float16\"",
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
-        "inputVec4": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "inputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "outputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "inputScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
@@ -2052,9 +1935,8 @@
           "name": "MultiHeadAttention.Flash",
           "shader": "attn-flash-online.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
-          "bindings": ["query", "key", "value", "attn_mask_2", "bias", "output", "params_8"],
-          "dispatch": { "x": "dim(shapes.queryT, 1)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
@@ -2063,17 +1945,10 @@
       "priority": 18,
       "when": ["true", "not present.biasT", "flashMaskShapeOk"],
       "derive": {
-        "headsFromParams": false,
-        "hasBias": false,
         "hasCausal": true,
-        "hasWindow": false,
         "combineSubgroups": false,
-        "hasMask": true,
-        "maskIsBool": false,
         "maskElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
-        "usesF16": "tensorDtypes.queryT == \"float16\"",
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
-        "inputVec4": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "inputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "outputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "qNumHeads": "attrs.num_heads",
@@ -2089,7 +1964,7 @@
           "name": "MultiHeadAttention.NoBiasOnlineFlashNoSg",
           "shader": "attn-flash-online.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
-          "bindings": ["query", "key", "value", "attn_mask_2", "output", "params_8"],
           "dispatch": { "x": "dim(shapes.queryT, 1)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
         }
       ]
@@ -2099,17 +1974,10 @@
       "priority": 17,
       "when": ["true", "biasOk", "flashMaskShapeOk"],
       "derive": {
-        "headsFromParams": false,
-        "hasBias": true,
         "hasCausal": true,
-        "hasWindow": false,
         "combineSubgroups": false,
-        "hasMask": true,
-        "maskIsBool": false,
         "maskElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
-        "usesF16": "tensorDtypes.queryT == \"float16\"",
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
-        "inputVec4": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "inputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "outputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "inputScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
@@ -2127,7 +1995,7 @@
           "name": "MultiHeadAttention.BiasOnlineFlashNoSg",
           "shader": "attn-flash-online.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
-          "bindings": ["query", "key", "value", "attn_mask_2", "bias", "output", "params_8"],
           "dispatch": { "x": "dim(shapes.queryT, 1)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
         }
       ]
@@ -2142,13 +2010,11 @@
       },
       "derive": {
         "qNumHeads": "attrs.num_heads",
-        "headDim": "headDimPlan",
         "qHidden": "dim(shapes.queryT, 2)",
         "materializedSoftmaxWg": "materializedCachedSoftmaxWg if materializedCachedSoftmaxOk else tunables.MATERIALIZED_SOFTMAX_WORKGROUP_SIZE",
         "materializedSoftmaxCols": "dim(shapes.keyT, 1)",
         "materializedSoftmaxCols4": "dim(shapes.keyT, 1) / 4",
         "softmaxElementType": "\"vec4<f32>\" if materializedCachedSoftmaxOk else \"f32\"",
-        "hasBias": false,
         "useSubgroups": true
       },
       "intermediates": [
@@ -2164,7 +2030,7 @@
           "name": "MultiHeadAttention.MaterializedScoresSgmat",
           "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
           "derive": { "phase": "\"score\"" },
-          "bindings": ["query_2", "key_2", "scores", "params_4"],
           "dispatch": {
             "x": "ceilDiv(dim(shapes.keyT, 1), materializedSgmatKeyTile)",
             "y": "ceilDiv(dim(shapes.queryT, 1), materializedSgmatQueryTile)",
@@ -2176,20 +2042,19 @@
           "name": "MultiHeadAttention.MaterializedSoftmaxSgmat",
           "shader": "attn-materialized-softmax-f32.wgsl.jinja",
           "derive": { "cacheVec4": "materializedCachedSoftmaxOk" },
-          "bindings": ["scores_3", "params_6"],
           "dispatch": {
             "x": "min(dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1), 65535)",
             "y": "ceilDiv(dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1), 65535)",
             "z": 1
-          },
-          "subgroupCollectivesWidth": "portable"
         },
         {
           "id": "apply",
           "name": "MultiHeadAttention.MaterializedApplySgmat",
           "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
           "derive": { "phase": "\"apply\"" },
-          "bindings": ["scores_2", "value_2", "output_2", "params_7"],
           "dispatch": {
             "x": "ceilDiv(headDim, materializedSgmatKeyTile)",
             "y": "ceilDiv(dim(shapes.queryT, 1), materializedSgmatQueryTile)",
@@ -2208,13 +2073,11 @@
       },
       "derive": {
         "qNumHeads": "attrs.num_heads",
-        "headDim": "headDimPlan",
         "qHidden": "dim(shapes.queryT, 2)",
         "materializedSoftmaxWg": "materializedCachedSoftmaxWg if materializedCachedSoftmaxOk else tunables.MATERIALIZED_SOFTMAX_WORKGROUP_SIZE",
         "materializedSoftmaxCols": "dim(shapes.keyT, 1)",
         "materializedSoftmaxCols4": "dim(shapes.keyT, 1) / 4",
         "softmaxElementType": "\"vec4<f32>\" if materializedCachedSoftmaxOk else \"f32\"",
-        "hasBias": true,
         "useSubgroups": true
       },
       "intermediates": [
@@ -2230,7 +2093,7 @@
           "name": "MultiHeadAttention.MaterializedScoresSgmatBias",
           "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
           "derive": { "phase": "\"score\"" },
-          "bindings": ["query_2", "key_2", "bias_2", "scores", "params_4"],
           "dispatch": {
             "x": "ceilDiv(dim(shapes.keyT, 1), materializedSgmatKeyTile)",
             "y": "ceilDiv(dim(shapes.queryT, 1), materializedSgmatQueryTile)",
@@ -2242,20 +2105,19 @@
           "name": "MultiHeadAttention.MaterializedSoftmaxSgmatBias",
           "shader": "attn-materialized-softmax-f32.wgsl.jinja",
           "derive": { "cacheVec4": "materializedCachedSoftmaxOk" },
-          "bindings": ["scores_3", "params_6"],
           "dispatch": {
             "x": "min(dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1), 65535)",
             "y": "ceilDiv(dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1), 65535)",
             "z": 1
-          },
-          "subgroupCollectivesWidth": "portable"
         },
         {
           "id": "apply",
           "name": "MultiHeadAttention.MaterializedApplySgmatBias",
           "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
           "derive": { "phase": "\"apply\"" },
-          "bindings": ["scores_2", "value_2", "bias_2", "output_2", "params_7"],
           "dispatch": {
             "x": "ceilDiv(headDim, materializedSgmatKeyTile)",
             "y": "ceilDiv(dim(shapes.queryT, 1), materializedSgmatQueryTile)",

     "canPinSubgroupSize32": "device.features.has(\"subgroups\") and device.features.has(\"subgroup-size-control\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize <= 32 and device.adapterInfo.subgroupMaxSize >= 32",
     "pinSubgroupSize32": "canPinSubgroupSize32 and not wave32Adapter",
     "wave32Effective": "wave32Adapter or pinSubgroupSize32",
+    "wave32SubgroupsUsable": "subgroupsWave32 or pinSubgroupSize32",
+    "headDim": "dim(shapes.queryT, 2) / attrs.num_heads if (ranks.queryT == 3 and attrs.num_heads > 0) else 0",
     "qkvDtypesOk": "tensorDtypes.keyT == tensorDtypes.queryT and tensorDtypes.valueT == tensorDtypes.queryT and tensorDtypes.outputT == tensorDtypes.queryT",
     "floatDtypeOk": "(tensorDtypes.queryT == \"float32\" or tensorDtypes.queryT == \"float16\") and f16Ok(tensorDtypes.queryT)",
     "qkvShapeOk": "ranks.queryT == 3 and ranks.keyT == 3 and ranks.valueT == 3 and ranks.outputT == 3 and attrs.num_heads > 0 and dim(shapes.queryT, 2) % attrs.num_heads == 0 and dim(shapes.keyT, 2) == dim(shapes.queryT, 2) and dim(shapes.valueT, 2) == dim(shapes.queryT, 2) and dim(shapes.keyT, 1) == dim(shapes.valueT, 1) and dim(shapes.queryT, 0) == dim(shapes.keyT, 0) and dim(shapes.queryT, 0) == dim(shapes.valueT, 0) and dim(shapes.outputT, 0) == dim(shapes.queryT, 0) and dim(shapes.outputT, 1) == dim(shapes.queryT, 1) and dim(shapes.outputT, 2) == dim(shapes.valueT, 2)",
     "qkvContractOk": "qkvShapeOk and qkvDtypesOk and floatDtypeOk and noAttnBias",
     "qkvMaskContractOk": "qkvShapeOk and qkvDtypesOk and floatDtypeOk and attnBiasOk",
     "biasOk": "present.biasT and ranks.biasT == 1 and tensorDtypes.biasT == tensorDtypes.queryT and dim(shapes.biasT, 0) == 3 * dim(shapes.queryT, 2)",
+    "hasBias": "present.biasT",
+    "hasMask": "present.attentionBiasT",
+    "q32BroadcastSubgroupLanes": "32 if wave32SubgroupsUsable else 0",
+    "q32BroadcastF32HeadVectors": "headDim / 4 if headDim % 4 == 0 else 0",
+    "q32BroadcastF32RegisterGeometry": "wave32SubgroupsUsable and q32BroadcastF32HeadVectors == q32BroadcastSubgroupLanes",
     "subgroupCluster4": "not narrowSubgroupRange and device.features.has(\"subgroups\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize >= 4 and device.adapterInfo.subgroupMinSize % 4 == 0 and device.adapterInfo.subgroupMaxSize % 4 == 0",
     "subgroupCluster8": "not narrowSubgroupRange and 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",
     "attentionDispatchFits": "dim(shapes.queryT, 0) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and dim(shapes.queryT, 1) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and attrs.num_heads <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
+    "flashHeadOk": "headDim % 4 == 0 and headDim >= 32 and headDim <= 256",
     "flashSizeOk": "flashHeadOk and (dim(shapes.queryT, 1) * attrs.num_heads >= tunables.FLASH_MIN_QUERY_HEADS or (dim(shapes.queryT, 1) == 1 and dim(shapes.keyT, 1) >= 512) or (dim(shapes.queryT, 1) > 1 and dim(shapes.keyT, 1) >= 2048)) and attentionDispatchFits",
     "flashShapeOk": "qkvContractOk and flashSizeOk",
     "flashMaskShapeOk": "qkvMaskContractOk and flashSizeOk",
     "noBiasSplitKCount": "min(tunables.DECODE_MAX_SPLITS if dim(shapes.queryT, 1) == 1 else max(1, ceilDiv(tunables.SPLITK_TARGET_WORKGROUPS, dim(shapes.queryT, 0) * dim(shapes.queryT, 1) * attrs.num_heads)), ceilDiv(dim(shapes.keyT, 1), tunables.DECODE_KEYS_PER_SPLIT))",
     "biasSplitKCount": "min(tunables.DECODE_MAX_SPLITS, ceilDiv(dim(shapes.keyT, 1), tunables.DECODE_KEYS_PER_SPLIT))",
+    "noBiasPartialOutBytes": "dim(shapes.queryT, 0) * dim(shapes.queryT, 1) * attrs.num_heads * noBiasSplitKCount * headDim * 4",
     "noBiasStatsBytes": "2 * dim(shapes.queryT, 0) * dim(shapes.queryT, 1) * attrs.num_heads * noBiasSplitKCount * 4",
     "noBiasSplitScratchFits": "noBiasPartialOutBytes <= device.limits.maxStorageBufferBindingSize and noBiasPartialOutBytes <= device.limits.maxBufferSize and noBiasStatsBytes <= device.limits.maxStorageBufferBindingSize and noBiasStatsBytes <= device.limits.maxBufferSize",
     "noBiasSplitDispatchFits": "dim(shapes.queryT, 1) * noBiasSplitKCount <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
+    "biasPartialOutBytes": "dim(shapes.queryT, 0) * attrs.num_heads * biasSplitKCount * headDim * 4",
     "biasStatsBytes": "2 * dim(shapes.queryT, 0) * attrs.num_heads * biasSplitKCount * 4",
     "biasSplitScratchFits": "biasPartialOutBytes <= device.limits.maxStorageBufferBindingSize and biasPartialOutBytes <= device.limits.maxBufferSize and biasStatsBytes <= device.limits.maxStorageBufferBindingSize and biasStatsBytes <= device.limits.maxBufferSize",
     "biasSplitDispatchFits": "biasSplitKCount <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
     "decodeSplitKNoBiasOk": "qkvContractOk and attrs.unidirectional == 0 and dim(shapes.queryT, 1) == 1 and dim(shapes.keyT, 1) >= 512 and flashHeadOk and attentionDispatchFits and noBiasSplitDispatchFits and noBiasSplitScratchFits",
     "shortQuerySplitKNoBiasOk": "qkvContractOk and attrs.unidirectional == 0 and dim(shapes.queryT, 1) >= 2 and dim(shapes.queryT, 1) <= 16 and dim(shapes.keyT, 1) >= 2048 and flashHeadOk and attentionDispatchFits and noBiasSplitDispatchFits and noBiasSplitScratchFits",
     "decodeSplitKBiasOk": "biasOk and qkvContractOk and attrs.unidirectional == 0 and dim(shapes.queryT, 1) == 1 and dim(shapes.keyT, 1) >= 512 and flashHeadOk and attentionDispatchFits and biasSplitDispatchFits and biasSplitScratchFits",
+    "decodeSplitKPortablePreferred": "tensorDtypes.queryT == \"float32\" and device.features.has(\"subgroups\") and has(device.adapterInfo, \"subgroupMinSize\") and headDim / 4 < device.adapterInfo.subgroupMinSize",
     "materializedScoreBytes": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1) * dim(shapes.keyT, 1) * 4",
     "materializedScoreFits": "materializedScoreBytes <= device.limits.maxStorageBufferBindingSize and materializedScoreBytes <= device.limits.maxBufferSize",
     "materializedWorkgroupSize": "tunables.MATERIALIZED_WORKGROUP_DIM * tunables.MATERIALIZED_WORKGROUP_DIM",
     "materializedScoreStorageBytes": "(tunables.MATERIALIZED_QUERY_TILE * tunables.MATERIALIZED_INNER_TILE + tunables.MATERIALIZED_KEY_TILE * (tunables.MATERIALIZED_INNER_TILE + 4)) * 4",
+    "materializedApplyTileN": "tunables.MATERIALIZED_VALUE_TILE_D128 if headDim == 128 else tunables.MATERIALIZED_VALUE_TILE",
     "materializedApplyStorageBytes": "(tunables.MATERIALIZED_QUERY_TILE * tunables.MATERIALIZED_INNER_TILE + tunables.MATERIALIZED_INNER_TILE * materializedApplyTileN) * 4",
     "materializedTileGeometryOk": "tunables.MATERIALIZED_INNER_TILE % 4 == 0 and (materializedApplyTileN / tunables.MATERIALIZED_WORKGROUP_DIM) % 4 == 0 and tunables.MATERIALIZED_INNER_TILE > 0 and tunables.MATERIALIZED_WORKGROUP_DIM > 0 and tunables.MATERIALIZED_QUERY_TILE % tunables.MATERIALIZED_WORKGROUP_DIM == 0 and tunables.MATERIALIZED_KEY_TILE % tunables.MATERIALIZED_WORKGROUP_DIM == 0 and materializedApplyTileN % tunables.MATERIALIZED_WORKGROUP_DIM == 0 and tunables.MATERIALIZED_QUERY_TILE * tunables.MATERIALIZED_INNER_TILE >= materializedWorkgroupSize and tunables.MATERIALIZED_KEY_TILE * tunables.MATERIALIZED_INNER_TILE >= materializedWorkgroupSize and tunables.MATERIALIZED_INNER_TILE * materializedApplyTileN >= materializedWorkgroupSize and (tunables.MATERIALIZED_QUERY_TILE * tunables.MATERIALIZED_INNER_TILE) % materializedWorkgroupSize == 0 and (tunables.MATERIALIZED_KEY_TILE * tunables.MATERIALIZED_INNER_TILE) % materializedWorkgroupSize == 0 and (tunables.MATERIALIZED_INNER_TILE * materializedApplyTileN) % materializedWorkgroupSize == 0",
     "materializedDeviceOk": "materializedTileGeometryOk and materializedWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.MATERIALIZED_WORKGROUP_DIM <= device.limits.maxComputeWorkgroupSizeX and tunables.MATERIALIZED_WORKGROUP_DIM <= device.limits.maxComputeWorkgroupSizeY and materializedScoreStorageBytes <= device.limits.maxComputeWorkgroupStorageSize and materializedApplyStorageBytes <= device.limits.maxComputeWorkgroupStorageSize",
     "materializedWideSimdOk": "device.features.has(\"subgroups\") or (has(device.adapterInfo, \"subgroupMinSize\") and device.adapterInfo.subgroupMinSize >= 16)",
+    "materializedF32CoreOk": "qkvContractOk and tensorDtypes.queryT == \"float32\" and attrs.unidirectional == 0 and headDim >= 64 and headDim <= 128 and dim(shapes.queryT, 1) >= 512 and dim(shapes.keyT, 1) >= 512 and dim(shapes.queryT, 0) * attrs.num_heads <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(dim(shapes.queryT, 1), tunables.MATERIALIZED_QUERY_TILE) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(dim(shapes.keyT, 1), tunables.MATERIALIZED_KEY_TILE) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and materializedScoreFits and materializedDeviceOk",
     "materializedSoftmaxStorageBytes": "tunables.MATERIALIZED_SOFTMAX_WORKGROUP_SIZE * 8 + 8",
     "materializedSoftmaxResourcesFit": "tunables.MATERIALIZED_SOFTMAX_WORKGROUP_SIZE > 0 and pow2ceil(tunables.MATERIALIZED_SOFTMAX_WORKGROUP_SIZE) == tunables.MATERIALIZED_SOFTMAX_WORKGROUP_SIZE and tunables.MATERIALIZED_SOFTMAX_WORKGROUP_SIZE <= deviceWorkgroupCap and materializedSoftmaxStorageBytes <= device.limits.maxComputeWorkgroupStorageSize",
     "materializedSgmatQueryTile": "tunables.MATERIALIZED_SGMAT_QUERY_TILE",
     "materializedSgmatResourcesFit": "materializedSgmatGeometryOk and materializedSgmatWorkgroupSize <= deviceWorkgroupCap and materializedSgmatCompactStorageBytes <= device.limits.maxComputeWorkgroupStorageSize",
     "materializedSgmatDispatchFits": "ceilDiv(dim(shapes.keyT, 1), materializedSgmatKeyTile) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(dim(shapes.queryT, 1), materializedSgmatQueryTile) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and dim(shapes.queryT, 0) * attrs.num_heads <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1), min(device.limits.maxComputeWorkgroupsPerDimension, 65535)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
     "materializedSgmatDirectScoreStore": "dim(shapes.queryT, 1) % materializedSgmatQueryTile == 0 and dim(shapes.keyT, 1) % materializedSgmatKeyTile == 0",
+    "materializedSgmatDirectApplyStore": "dim(shapes.queryT, 1) % materializedSgmatQueryTile == 0 and headDim % materializedSgmatKeyTile == 0",
     "materializedSgmatRuntimeDirectStore": "dim(shapes.queryT, 1) >= 2 * materializedSgmatQueryTile and dim(shapes.keyT, 1) >= 2 * materializedSgmatKeyTile",
+    "materializedSgmatCoreOk": "qkvContractOk and tensorDtypes.queryT == \"float32\" and attrs.unidirectional == 0 and headDim >= 32 and headDim <= 256 and dim(shapes.queryT, 1) >= 512 and dim(shapes.keyT, 1) >= 512 and device.features.has(\"subgroups\") and wave32Effective and device.features.has(\"chromium-experimental-subgroup-matrix\") and materializedSgmatBuffersFit and materializedSgmatResourcesFit and materializedSgmatDispatchFits",
     "materializedCachedSoftmaxWg": "min(tunables.MATERIALIZED_CACHED_SOFTMAX_WORKGROUP_SIZE, device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
     "materializedCachedSoftmaxVecsPerLane": "ceilDiv(ceilDiv(dim(shapes.keyT, 1), 4), max(1, materializedCachedSoftmaxWg))",
     "materializedCachedSoftmaxStorageBytes": "materializedCachedSoftmaxWg * 8 + 8",
     "materializedSgmatOk": "materializedSgmatCoreOk and materializedAdaptiveSoftmaxOk",
     "materializedSgmatFusedOk": "materializedSgmatCoreOk",
     "materializedF32Ok": "materializedF32CoreOk and materializedAdaptiveSoftmaxOk",
+    "clusterTileKWg64": "max(1, min(tunables.FLASH_MAX_TILE_K, floor(device.limits.maxComputeWorkgroupStorageSize / (headDim * (8 if tensorDtypes.queryT != \"float16\" else 4) + tunables.FLASH_CLUSTER_WG_SMALL * 4))))",
+    "clusterTileKWg128": "max(1, min(tunables.FLASH_MAX_TILE_K, floor(device.limits.maxComputeWorkgroupStorageSize / (headDim * (8 if tensorDtypes.queryT != \"float16\" else 4) + tunables.FLASH_CLUSTER_WG_LARGE * 4))))",
+    "smallHeadShapeOk": "qkvContractOk and headDim < 32 and dim(shapes.keyT, 1) >= 64 and attrs.unidirectional == 0 and dim(shapes.queryT, 0) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and attrs.num_heads <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
+    "smallHeadParallelOk": "smallHeadShapeOk and dim(shapes.keyT, 1) <= 2048",
+    "prefillTiledStorageBytes": "headDim * tunables.PREFILL_QUERY_TILE * 4",
     "prefillTiledDeviceOk": "tunables.PREFILL_QUERY_TILE <= deviceWorkgroupCap and prefillTiledStorageBytes <= device.limits.maxComputeWorkgroupStorageSize",
+    "portableWorkgroupSize": "min(tunables.WORKGROUP_SIZE, pow2ceil(max(1, headDim)))",
+    "portableWorkgroupStorageBytes": "portableWorkgroupSize * 8 + max(1, headDim) * 4 + 16",
     "portableWorkgroupOk": "tunables.WORKGROUP_SIZE > 0 and pow2ceil(tunables.WORKGROUP_SIZE) == tunables.WORKGROUP_SIZE and portableWorkgroupSize <= deviceWorkgroupCap and portableWorkgroupStorageBytes <= device.limits.maxComputeWorkgroupStorageSize",
     "fallbackShapeOk": "qkvContractOk and portableWorkgroupOk",
     "fallbackMaskShapeOk": "qkvMaskContractOk and portableWorkgroupOk",
+    "smallSeqShapeOk": "qkvContractOk and tensorDtypes.queryT == \"float32\" and dim(shapes.queryT, 1) >= 1 and dim(shapes.queryT, 1) <= tunables.SMALL_SEQ_MAX and dim(shapes.keyT, 1) >= 1 and dim(shapes.keyT, 1) <= tunables.SMALL_SEQ_MAX and headDim >= 1 and attrs.unidirectional == 0",
+    "smallSeqPrivateFloats": "dim(shapes.keyT, 1) + headDim",
     "smallSeqWorkgroupSize": "max(32, pow2ceil(dim(shapes.queryT, 1)))",
+    "smallSeqSharedBytes": "dim(shapes.keyT, 1) * headDim * 8",
     "smallSeqResourcesFit": "smallSeqPrivateFloats <= tunables.SMALL_SEQ_MAX_PRIVATE_FLOATS and smallSeqWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup and smallSeqWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX and smallSeqSharedBytes <= device.limits.maxComputeWorkgroupStorageSize",
+    "smallSeqBlockedKvBytes": "dim(shapes.keyT, 1) * headDim * 8",
+    "smallSeqBlockedLaneBytes": "8 + headDim * 4",
     "smallSeqBlockedKeyLanes": "min(pow2ceil(dim(shapes.keyT, 1)), 16 if smallSeqBlockedKvBytes + tunables.SMALL_SEQ_BLOCKED_QUERY_BLOCK * 16 * smallSeqBlockedLaneBytes <= device.limits.maxComputeWorkgroupStorageSize else (8 if smallSeqBlockedKvBytes + tunables.SMALL_SEQ_BLOCKED_QUERY_BLOCK * 8 * smallSeqBlockedLaneBytes <= device.limits.maxComputeWorkgroupStorageSize else 4))",
     "smallSeqBlockedWorkgroupSize": "tunables.SMALL_SEQ_BLOCKED_QUERY_BLOCK * smallSeqBlockedKeyLanes",
     "smallSeqBlockedSharedBytes": "smallSeqBlockedKvBytes + smallSeqBlockedWorkgroupSize * smallSeqBlockedLaneBytes",
+    "smallSeqBlockedShapeOk": "qkvContractOk and tensorDtypes.queryT == \"float32\" and headDim % 4 == 0 and headDim >= 4 and headDim <= tunables.SMALL_SEQ_BLOCKED_MAX_HEAD_DIM and dim(shapes.keyT, 1) >= 1 and dim(shapes.keyT, 1) <= tunables.SMALL_SEQ_BLOCKED_MAX_KV and dim(shapes.queryT, 1) >= 1",
     "smallSeqBlockedFits": "smallSeqBlockedWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup and smallSeqBlockedWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX and smallSeqBlockedSharedBytes <= device.limits.maxComputeWorkgroupStorageSize and ceilDiv(dim(shapes.queryT, 1), tunables.SMALL_SEQ_BLOCKED_QUERY_BLOCK) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and attrs.num_heads <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and dim(shapes.queryT, 0) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
     "smallSeqDispatchFits": "attrs.num_heads <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and dim(shapes.queryT, 0) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
+    "materializedSgmatCoreF16Ok": "qkvContractOk and tensorDtypes.queryT == \"float16\" and tensorDtypes.keyT == \"float16\" and tensorDtypes.valueT == \"float16\" and device.features.has(\"shader-f16\") and attrs.unidirectional == 0 and headDim >= 32 and headDim <= 256 and dim(shapes.queryT, 1) >= 512 and dim(shapes.keyT, 1) >= 512 and device.features.has(\"subgroups\") and wave32Effective and device.features.has(\"chromium-experimental-subgroup-matrix\") and materializedSgmatBuffersFit and materializedSgmatResourcesFit and materializedSgmatDispatchFits",
     "materializedSgmatFusedF16Ok": "materializedSgmatCoreF16Ok",
+    "attentionScaleExpression": "\"select(inverseSqrt(f32(HEAD_DIM)), params.scale, params.scale != 0.0)\"",
+    "smallHeadValueWg": "pow(2, log2ceil(min(64, device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX) + 1) - 1)",
+    "smallHeadValueStorage": "max(dim(shapes.keyT, 1), smallHeadValueWg * headDim)",
+    "smallHeadValueQueryBlock": "2 if dim(shapes.queryT, 1) >= 2 and (smallHeadValueStorage + smallHeadValueWg) * 8 <= device.limits.maxComputeWorkgroupStorageSize else 1",
+    "smallHeadValueEstimatedSteps": "headDim * (ceilDiv(dim(shapes.keyT, 1), smallHeadValueWg) + 2 * log2ceil(smallHeadValueWg) * (2 if tensorDtypes.queryT == \"float16\" else 1) + 2)",
+    "smallHeadValueSgReductionSteps": "ceilDiv(smallHeadValueWg, max(1, device.adapterInfo.subgroupMinSize)) if has(device.adapterInfo, \"subgroupMinSize\") else log2ceil(smallHeadValueWg)",
+    "smallHeadValueSgEstimatedSteps": "headDim * (ceilDiv(dim(shapes.keyT, 1), smallHeadValueWg) + 2 * smallHeadValueSgReductionSteps * (2 if tensorDtypes.queryT == \"float16\" else 1) + 2)"
   },
   "bindings": {
+    "query": { "arg": "queryT", "elementType": "$inputElement" },
+    "key": { "arg": "keyT", "elementType": "$inputElement" },
+    "value": { "arg": "valueT", "elementType": "$inputElement" },
+    "bias": { "arg": "biasT", "elementType": "$inputScalar" },
+    "output": { "arg": "outputT", "elementType": "$outputElement" },
     "params": {
       "struct": [
         { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" },
         { "name": "kvSeq", "type": "u32", "value": "dim(shapes.keyT, 1)" },
         { "name": "isCausal", "type": "u32", "value": "attrs.unidirectional" }
       ]
     },
+    "params_main": {
       "name": "params",
       "struct": [
         { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" },
         { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" }
       ]
     },
+    "q": { "arg": "queryT", "elementType": "$scalar" },
+    "k": { "arg": "keyT", "elementType": "$scalar" },
+    "v": { "arg": "valueT", "elementType": "$scalar" },
+    "y": { "arg": "outputT", "elementType": "$scalar" },
+    "params_scores": {
       "name": "params",
       "struct": [
         { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" },
         { "name": "kvSeq", "type": "u32", "value": "dim(shapes.keyT, 1)" },
         { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" }
       ]
     },
+    "query_f32": { "arg": "queryT", "name": "query", "elementType": "f32" },
+    "key_f32": { "arg": "keyT", "name": "key", "elementType": "f32" },
+    "scores": { "scratch": "materializedScores", "elementType": "f32" },
+    "scorePartials": { "scratch": "materializedScorePartials", "elementType": "f32" },
+    "bias_f32": { "arg": "biasT", "name": "bias", "elementType": "f32" },
+    "scorePartials_f32": {
       "scratch": "materializedScorePartials",
       "name": "scorePartials",
       "buffer": "read-only-storage",
       "elementType": "f32"
     },
+    "rowStats": { "scratch": "materializedRowStats", "elementType": "f32" },
+    "params_rows": {
       "name": "params",
       "struct": [
         { "name": "rows", "type": "u32", "value": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)" }
       ]
     },
+    "scores_f32": {
       "scratch": "materializedScores",
       "name": "scores",
       "buffer": "read-only-storage",
       "elementType": "f32"
     },
+    "value_f32": { "arg": "valueT", "name": "value", "elementType": "f32" },
+    "output_f32": { "arg": "outputT", "name": "output", "elementType": "f32" },
+    "rowStats_f32": {
       "scratch": "materializedRowStats",
       "name": "rowStats",
       "buffer": "read-only-storage",
       "elementType": "f32"
     },
+    "params_apply": {
       "name": "params",
       "struct": [
         { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" },
         { "name": "kvSeq", "type": "u32", "value": "dim(shapes.keyT, 1)" }
       ]
     },
+    "attn_mask_main": { "arg": "attentionBiasT", "name": "attn_mask", "elementType": "$maskElement" },
+    "params__uniform": {
       "name": "params",
       "struct": [
         { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" },
         { "name": "kvSeq", "type": "u32", "value": "dim(shapes.keyT, 1)" },
         { "name": "maskSeqStride", "type": "u32", "value": "dim(shapes.attentionBiasT, 3)" }
       ]
     },
+    "query_query_t": { "arg": "queryT", "name": "query", "elementType": "$inputVec4" },
+    "key_key_t": { "arg": "keyT", "name": "key", "elementType": "$inputVec4" },
+    "value_value_t": { "arg": "valueT", "name": "value", "elementType": "$inputVec4" },
+    "partial_out": { "scratch": "partialOut", "elementType": "vec4<f32>" },
+    "partial_stats": { "scratch": "partialStats", "elementType": "vec2<f32>" },
+    "params_kv_seq_scale": {
       "name": "params",
       "struct": [
         { "name": "kvSeq", "type": "u32", "value": "dim(shapes.keyT, 1)" },
         { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" }
       ]
     },
+    "partial_out_merge": {
       "scratch": "partialOut",
       "name": "partial_out",
       "buffer": "read-only-storage",
       "elementType": "vec4<f32>"
     },
+    "partial_stats_merge": {
       "scratch": "partialStats",
       "name": "partial_stats",
       "buffer": "read-only-storage",
       "elementType": "vec2<f32>"
     },
+    "output_merge": { "arg": "outputT", "name": "output", "elementType": "$inputVec4" },
+    "scores_softmax": { "scratch": "materializedScores", "name": "scores", "elementType": "$softmaxElementType" }
   },
   "variants": [
+    {
+      "id": "qkv_no_bias_small_head_value_subgroups",
+      "priority": 12,
+      "when": ["not present.biasT", "smallHeadShapeOk", "headDim > 0", "headDim <= smallHeadValueWg", "(smallHeadValueStorage + smallHeadValueWg) * 4 * smallHeadValueQueryBlock <= device.limits.maxComputeWorkgroupStorageSize", "device.wgslLanguageFeatures.has(\"subgroup_id\")"],
+      "derive": {
+        "outputScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
+        "inputElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
+        "outputElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
+        "qHidden": "dim(shapes.queryT, 2)",
+        "headDim": "dim(shapes.queryT, 2) / attrs.num_heads",
+        "kvSeq": "dim(shapes.keyT, 1)",
+        "valueWorkgroupSize": "smallHeadValueWg",
+        "scoreStorageElements": "smallHeadValueStorage",
+        "queryBlock": "smallHeadValueQueryBlock",
+        "queryTail": "dim(shapes.queryT, 1) % smallHeadValueQueryBlock != 0",
+        "useValueSubgroups": true
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "MultiHeadAttention.SmallHeadValueSubgroups",
+          "shader": "attn-small-head-value.wgsl.jinja",
+          "bindings": ["query", "key", "value", "output", "params_main"],
+          "dispatch": {
+            "x": "min(ceilDiv(dim(shapes.queryT, 1), smallHeadValueQueryBlock), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
+            "y": "attrs.num_heads",
+            "z": "dim(shapes.queryT, 0)"
+          }
+        }
+      ],
+      "requires": { "features": ["subgroups"] },
+      "demoteWhen": ["dim(shapes.keyT, 1) <= smallHeadValueSgEstimatedSteps"]
+    },
+    {
+      "id": "qkv_no_bias_small_head_value_tree",
+      "priority": 11,
+      "when": ["not present.biasT", "smallHeadShapeOk", "headDim > 0", "headDim <= smallHeadValueWg", "(smallHeadValueStorage + smallHeadValueWg) * 4 * smallHeadValueQueryBlock <= device.limits.maxComputeWorkgroupStorageSize", "true"],
+      "derive": {
+        "outputScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
+        "inputElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
+        "outputElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
+        "qHidden": "dim(shapes.queryT, 2)",
+        "headDim": "dim(shapes.queryT, 2) / attrs.num_heads",
+        "kvSeq": "dim(shapes.keyT, 1)",
+        "valueWorkgroupSize": "smallHeadValueWg",
+        "scoreStorageElements": "smallHeadValueStorage",
+        "queryBlock": "smallHeadValueQueryBlock",
+        "queryTail": "dim(shapes.queryT, 1) % smallHeadValueQueryBlock != 0",
+        "useValueSubgroups": false
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "MultiHeadAttention.SmallHeadValueTree",
+          "shader": "attn-small-head-value.wgsl.jinja",
+          "bindings": ["query", "key", "value", "output", "params_main"],
+          "dispatch": {
+            "x": "min(ceilDiv(dim(shapes.queryT, 1), smallHeadValueQueryBlock), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
+            "y": "attrs.num_heads",
+            "z": "dim(shapes.queryT, 0)"
+          }
+        }
+      ],
+      "requires": { "features": [] },
+      "demoteWhen": ["dim(shapes.keyT, 1) <= smallHeadValueEstimatedSteps"]
+    },
     {
       "id": "qkv_bias_small_seq_blocked",
       "priority": 40,
       "when": ["biasOk", "smallSeqBlockedShapeOk", "not flashShapeOk", "smallSeqBlockedFits"],
       "derive": {
         "inputElement": "\"vec4<f32>\"",
         "outputElement": "\"vec4<f32>\"",
         "inputScalar": "\"f32\"",
+        "headDimV4": "headDim / 4",
         "hidden": "dim(shapes.queryT, 2)",
         "hiddenV4": "dim(shapes.queryT, 2) / 4",
         "kvSeq": "dim(shapes.keyT, 1)",
             "x": "ceilDiv(dim(shapes.queryT, 1), tunables.SMALL_SEQ_BLOCKED_QUERY_BLOCK)",
             "y": "attrs.num_heads",
             "z": "dim(shapes.queryT, 0)"
           }
         }
       ]
       "priority": 40,
       "when": ["not present.biasT", "smallSeqBlockedShapeOk", "not flashShapeOk", "smallSeqBlockedFits"],
       "derive": {
         "inputElement": "\"vec4<f32>\"",
         "outputElement": "\"vec4<f32>\"",
+        "headDimV4": "headDim / 4",
         "hidden": "dim(shapes.queryT, 2)",
         "hiddenV4": "dim(shapes.queryT, 2) / 4",
         "kvSeq": "dim(shapes.keyT, 1)",
             "x": "ceilDiv(dim(shapes.queryT, 1), tunables.SMALL_SEQ_BLOCKED_QUERY_BLOCK)",
             "y": "attrs.num_heads",
             "z": "dim(shapes.queryT, 0)"
           }
         }
       ]
       "priority": 60,
       "when": ["not present.biasT", "smallSeqShapeOk", "not flashShapeOk", "smallSeqResourcesFit", "smallSeqDispatchFits"],
       "derive": {
         "inputElement": "\"f32\"",
         "outputElement": "\"f32\"",
         "headDim": "dim(shapes.queryT, 2) / attrs.num_heads",
           "id": "main",
           "name": "MultiHeadAttention",
           "shader": "mha-small-seq.wgsl.jinja",
+          "bindings": ["query", "key", "value", "output", "params_main"],
+          "dispatch": { "x": "attrs.num_heads", "y": "dim(shapes.queryT, 0)" }
         }
       ]
     },
     {
       "id": "qkv_no_bias_tiled_nosg",
       "priority": 19,
+      "when": ["not present.biasT", "qkvContractOk", "not smallHeadParallelOk", "headDim % 4 == 0", "headDim <= 128", "dim(shapes.queryT, 1) >= 31", "prefillTiledDeviceOk"],
       "supersededBy": ["qkv_no_bias_flash_cluster_nosg", "qkv_no_bias_flash_cluster_lpq4_nosg"],
       "derive": {
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "blockM": "tunables.PREFILL_QUERY_TILE",
         "vHeadCap": "dim(shapes.queryT, 2) / attrs.num_heads"
       },
             }
           ],
           "dispatch": {
+            "x": "min(ceilDiv((dim(shapes.outputT, 0) * attrs.num_heads * ceil(dim(shapes.outputT, 1) / blockM) * blockM), (blockM)), 65535)",
+            "y": "ceilDiv(ceilDiv((dim(shapes.outputT, 0) * attrs.num_heads * ceil(dim(shapes.outputT, 1) / blockM) * blockM), (blockM)), 65535)",
             "z": 1
           }
         }
     {
       "id": "qkv_bias_flash_q32_broadcast_f32_d128",
       "priority": 30,
+      "when": ["tensorDtypes.queryT == \"float32\"", "biasOk", "attrs.unidirectional == 0", "flashShapeOk", "q32BroadcastF32RegisterGeometry", "dim(shapes.queryT, 1) >= 31", "ceilDiv(dim(shapes.queryT, 1), 32) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "wave32SubgroupsUsable"],
       "requires": { "features": ["subgroups"] },
       "derive": {
         "hasCausal": false,
         "usesF16": false,
         "scalar": "\"f32\"",
         "inputElement": "\"vec4<f32>\"",
         "outputElement": "\"vec4<f32>\"",
         "inputScalar": "\"f32\"",
           "name": "MultiHeadAttention.FlashQ32BroadcastF32Bias",
           "shader": "attn-flash-q32-broadcast.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
+          "bindings": ["query", "key", "value", "bias", "output", "params_scores"],
           "dispatch": {
             "x": "ceilDiv(dim(shapes.queryT, 1), 32)",
             "y": "attrs.num_heads",
             "z": "dim(shapes.queryT, 0)"
+          }
         }
       ]
     },
       "priority": 10,
       "when": ["not present.biasT", "smallHeadParallelOk"],
       "derive": {
         "outputScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "inputElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "outputElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
       "passes": [
         {
           "id": "main",
+          "name": "MultiHeadAttention.SmallHeadParallel",
           "shader": "attn-small-head-parallel.wgsl.jinja",
+          "bindings": ["query", "key", "value", "output", "params_main"],
           "dispatch": {
             "x": "min(dim(shapes.queryT, 1), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
             "y": "attrs.num_heads",
     {
       "id": "qkv_no_bias_tiled_attn_bias_nosg",
       "priority": 17,
+      "when": ["not present.biasT", "qkvMaskContractOk", "not smallHeadParallelOk", "headDim % 4 == 0", "headDim <= 128", "dim(shapes.queryT, 1) >= 31", "prefillTiledDeviceOk"],
       "derive": {
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "blockM": "tunables.PREFILL_QUERY_TILE",
         "vHeadCap": "dim(shapes.queryT, 2) / attrs.num_heads"
       },
             }
           ],
           "dispatch": {
+            "x": "min(ceilDiv((dim(shapes.outputT, 0) * attrs.num_heads * ceil(dim(shapes.outputT, 1) / blockM) * blockM), (blockM)), 65535)",
+            "y": "ceilDiv(ceilDiv((dim(shapes.outputT, 0) * attrs.num_heads * ceil(dim(shapes.outputT, 1) / blockM) * blockM), (blockM)), 65535)",
             "z": 1
           }
         }
       },
       "derive": {
         "qNumHeads": "attrs.num_heads",
         "qHidden": "dim(shapes.queryT, 2)",
         "statSlots": "materializedSgmatStatSlots",
         "statQuerySeq": "dim(shapes.queryT, 1)"
       },
           "name": "MultiHeadAttention.MaterializedScoresSgmat",
           "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
           "derive": { "phase": "\"score\"", "emitRowStats": true },
+          "bindings": ["query_f32", "key_f32", "scores", "scorePartials", "params_scores"],
           "dispatch": {
             "x": "ceilDiv(dim(shapes.keyT, 1), materializedSgmatKeyTile)",
             "y": "ceilDiv(dim(shapes.queryT, 1), materializedSgmatQueryTile)",
             "z": "dim(shapes.queryT, 0) * attrs.num_heads"
+          }
         },
         {
           "id": "rowstats",
           "name": "MultiHeadAttention.MaterializedRowStatsCombine",
           "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja",
+          "bindings": ["scorePartials_f32", "rowStats", "params_rows"],
           "dispatch": {
             "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)), (materializedRowStatsWg)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
             "y": 1,
           "name": "MultiHeadAttention.MaterializedApplySgmat",
           "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
           "derive": { "phase": "\"apply\"", "fusedSoftmax": true },
+          "bindings": ["scores_f32", "value_f32", "output_f32", "rowStats_f32", "params_apply"],
           "dispatch": {
             "x": "ceilDiv(headDim, materializedSgmatKeyTile)",
             "y": "ceilDiv(dim(shapes.queryT, 1), materializedSgmatQueryTile)",
       },
       "derive": {
         "qNumHeads": "attrs.num_heads",
         "qHidden": "dim(shapes.queryT, 2)",
         "statSlots": "materializedSgmatStatSlots",
         "statQuerySeq": "dim(shapes.queryT, 1)"
       },
           "name": "MultiHeadAttention.MaterializedScoresSgmatBias",
           "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
           "derive": { "phase": "\"score\"", "emitRowStats": true },
+          "bindings": ["query_f32", "key_f32", "bias_f32", "scores", "scorePartials", "params_scores"],
           "dispatch": {
             "x": "ceilDiv(dim(shapes.keyT, 1), materializedSgmatKeyTile)",
             "y": "ceilDiv(dim(shapes.queryT, 1), materializedSgmatQueryTile)",
             "z": "dim(shapes.queryT, 0) * attrs.num_heads"
+          }
         },
         {
           "id": "rowstats",
           "name": "MultiHeadAttention.MaterializedRowStatsCombineBias",
           "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja",
+          "bindings": ["scorePartials_f32", "rowStats", "params_rows"],
           "dispatch": {
             "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)), (materializedRowStatsWg)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
             "y": 1,
           "name": "MultiHeadAttention.MaterializedApplySgmatBias",
           "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
           "derive": { "phase": "\"apply\"", "fusedSoftmax": true },
+          "bindings": ["scores_f32", "value_f32", "bias_f32", "output_f32", "rowStats_f32", "params_apply"],
           "dispatch": {
             "x": "ceilDiv(headDim, materializedSgmatKeyTile)",
             "y": "ceilDiv(dim(shapes.queryT, 1), materializedSgmatQueryTile)",
       },
       "derive": {
         "qNumHeads": "attrs.num_heads",
         "qHidden": "dim(shapes.queryT, 2)",
         "statSlots": "materializedSgmatStatSlots",
         "statQuerySeq": "dim(shapes.queryT, 1)",
         "operandF16": true,
           "name": "MultiHeadAttention.MaterializedScoresSgmatF16",
           "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
           "derive": { "phase": "\"score\"", "emitRowStats": true },
+          "bindings": ["query", "key", "scores", "scorePartials", "params_scores"],
           "dispatch": {
             "x": "ceilDiv(dim(shapes.keyT, 1), materializedSgmatKeyTile)",
             "y": "ceilDiv(dim(shapes.queryT, 1), materializedSgmatQueryTile)",
             "z": "dim(shapes.queryT, 0) * attrs.num_heads"
+          }
         },
         {
           "id": "rowstats",
           "name": "MultiHeadAttention.MaterializedRowStatsCombine",
           "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja",
+          "bindings": ["scorePartials_f32", "rowStats", "params_rows"],
           "dispatch": {
             "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)), (materializedRowStatsWg)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
             "y": 1,
           "name": "MultiHeadAttention.MaterializedApplySgmatF16",
           "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
           "derive": { "phase": "\"apply\"", "fusedSoftmax": true },
+          "bindings": ["scores_f32", "value", "output", "rowStats_f32", "params_apply"],
           "dispatch": {
             "x": "ceilDiv(headDim, materializedSgmatKeyTile)",
             "y": "ceilDiv(dim(shapes.queryT, 1), materializedSgmatQueryTile)",
     {
       "id": "qkv_no_bias_flash_cluster_lpq4_nosg",
       "priority": 20,
+      "when": ["not present.biasT", "flashShapeOk", "headDim % 16 == 0", "headDim % 32 != 0", "dim(shapes.queryT, 1) >= 31"],
       "requires": {},
       "derive": {
         "hasCausal": true,
         "usesF16": "tensorDtypes.queryT == \"float16\"",
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "inputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "outputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "qNumHeads": "attrs.num_heads",
     {
       "id": "qkv_no_bias_flash_cluster_nosg",
       "priority": 20,
+      "when": ["not present.biasT", "flashShapeOk", "headDim % 32 == 0", "dim(shapes.queryT, 1) >= 31"],
       "requires": {},
       "derive": {
         "hasCausal": true,
         "usesF16": "tensorDtypes.queryT == \"float16\"",
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "inputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "outputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "qNumHeads": "attrs.num_heads",
     {
       "id": "qkv_bias_flash_cluster_nosg",
       "priority": 19,
+      "when": ["biasOk", "flashShapeOk", "headDim % 32 == 0", "dim(shapes.queryT, 1) >= 31"],
       "requires": {},
       "derive": {
         "hasCausal": true,
         "usesF16": "tensorDtypes.queryT == \"float16\"",
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "inputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "outputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "qNumHeads": "attrs.num_heads",
     {
       "id": "qkv_no_bias_flash_cluster_lpq4",
       "priority": 22,
+      "when": ["not present.biasT", "flashShapeOk", "headDim % 16 == 0", "headDim % 32 != 0", "dim(shapes.queryT, 1) >= 31", "subgroupCluster4"],
       "requires": { "features": ["subgroups"] },
       "derive": {
         "hasCausal": true,
         "usesF16": "tensorDtypes.queryT == \"float16\"",
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "inputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "outputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "qNumHeads": "attrs.num_heads",
             "x": "ceilDiv(dim(shapes.queryT, 1), TILE_Q)",
             "y": "attrs.num_heads",
             "z": "dim(shapes.queryT, 0)"
+          }
         }
       ]
     },
     {
       "id": "qkv_no_bias_flash_cluster",
       "priority": 22,
+      "when": ["not present.biasT", "flashShapeOk", "headDim % 32 == 0", "dim(shapes.queryT, 1) >= 31", "subgroupCluster8"],
       "requires": { "features": ["subgroups"] },
       "derive": {
         "hasCausal": true,
         "usesF16": "tensorDtypes.queryT == \"float16\"",
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "inputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "outputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "qNumHeads": "attrs.num_heads",
             "x": "ceilDiv(dim(shapes.queryT, 1), TILE_Q)",
             "y": "attrs.num_heads",
             "z": "dim(shapes.queryT, 0)"
+          }
         }
       ]
     },
     {
       "id": "qkv_bias_flash_cluster",
       "priority": 21,
+      "when": ["biasOk", "flashShapeOk", "headDim % 32 == 0", "dim(shapes.queryT, 1) >= 31", "subgroupCluster8"],
       "requires": { "features": ["subgroups"] },
       "derive": {
         "hasCausal": true,
         "usesF16": "tensorDtypes.queryT == \"float16\"",
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "inputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "outputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "qNumHeads": "attrs.num_heads",
             "x": "ceilDiv(dim(shapes.queryT, 1), TILE_Q)",
             "y": "attrs.num_heads",
             "z": "dim(shapes.queryT, 0)"
+          }
         }
       ]
     },
     {
       "id": "qkv_no_bias_flash_cluster_attn_bias",
       "priority": 22,
+      "when": ["not present.biasT", "flashMaskShapeOk", "headDim % 32 == 0", "dim(shapes.queryT, 1) >= 31", "subgroupCluster8"],
       "requires": { "features": ["subgroups"] },
       "derive": {
         "hasCausal": true,
         "usesF16": "tensorDtypes.queryT == \"float16\"",
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "inputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "outputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "maskElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
           "name": "MultiHeadAttention.Flash",
           "shader": "attn-flash-prefill-cluster.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
+          "bindings": ["query", "key", "value", "attn_mask_main", "output", "params__uniform"],
           "dispatch": {
             "x": "ceilDiv(dim(shapes.queryT, 1), TILE_Q)",
             "y": "attrs.num_heads",
             "z": "dim(shapes.queryT, 0)"
+          }
         }
       ]
     },
     {
       "id": "qkv_bias_flash_cluster_attn_bias",
       "priority": 21,
+      "when": ["biasOk", "flashMaskShapeOk", "headDim % 32 == 0", "dim(shapes.queryT, 1) >= 31", "subgroupCluster8"],
       "requires": { "features": ["subgroups"] },
       "derive": {
         "hasCausal": true,
         "usesF16": "tensorDtypes.queryT == \"float16\"",
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "inputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "outputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "maskElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
           "name": "MultiHeadAttention.Flash",
           "shader": "attn-flash-prefill-cluster.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
+          "bindings": ["query", "key", "value", "attn_mask_main", "bias", "output", "params__uniform"],
           "dispatch": {
             "x": "ceilDiv(dim(shapes.queryT, 1), TILE_Q)",
             "y": "attrs.num_heads",
             "z": "dim(shapes.queryT, 0)"
+          }
         }
       ]
     },
       "when": ["not present.biasT", "decodeSplitKNoBiasOk or shortQuerySplitKNoBiasOk"],
       "requires": {},
       "derive": {
         "useSubgroups": false,
         "splitQueries": true,
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "inputVec4": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
           "name": "MultiHeadAttention.DecodeSplitKNoSg",
           "shader": "attn-flash-decode-splitk.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
+          "bindings": ["query_query_t", "key_key_t", "value_value_t", "partial_out", "partial_stats", "params_kv_seq_scale"],
           "dispatch": {
             "x": "dim(shapes.queryT, 1) * noBiasSplitKCount",
             "y": "attrs.num_heads",
           "name": "MultiHeadAttention.DecodeSplitKMergeNoSg",
           "shader": "attn-flash-decode-splitk-merge.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
+          "bindings": ["partial_out_merge", "partial_stats_merge", "output_merge"],
           "dispatch": { "x": "dim(shapes.queryT, 1)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
         }
       ]
       "when": ["biasOk", "decodeSplitKBiasOk"],
       "requires": {},
       "derive": {
         "useSubgroups": false,
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "inputVec4": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "inputScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
           "name": "MultiHeadAttention.DecodeSplitKBiasNoSg",
           "shader": "attn-flash-decode-splitk.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
+          "bindings": ["query_query_t", "key_key_t", "value_value_t", "bias", "partial_out", "partial_stats", "params_kv_seq_scale"],
           "dispatch": { "x": "biasSplitKCount", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
         },
         {
           "name": "MultiHeadAttention.DecodeSplitKMergeBiasNoSg",
           "shader": "attn-flash-decode-splitk-merge.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
+          "bindings": ["partial_out_merge", "partial_stats_merge", "bias", "output_merge"],
           "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
         }
       ]
       "demoteWhen": ["decodeSplitKPortablePreferred"],
       "requires": { "features": ["subgroups"] },
       "derive": {
         "splitQueries": true,
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "inputVec4": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
           "name": "MultiHeadAttention.DecodeSplitK",
           "shader": "attn-flash-decode-splitk.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
+          "bindings": ["query_query_t", "key_key_t", "value_value_t", "partial_out", "partial_stats", "params_kv_seq_scale"],
           "dispatch": {
             "x": "dim(shapes.queryT, 1) * noBiasSplitKCount",
             "y": "attrs.num_heads",
             "z": "dim(shapes.queryT, 0)"
+          }
         },
         {
           "id": "merge",
           "name": "MultiHeadAttention.DecodeSplitKMerge",
           "shader": "attn-flash-decode-splitk-merge.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
+          "bindings": ["partial_out_merge", "partial_stats_merge", "output_merge"],
           "dispatch": { "x": "dim(shapes.queryT, 1)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
         }
       ]
       "when": ["biasOk", "decodeSplitKBiasOk"],
       "requires": { "features": ["subgroups"] },
       "derive": {
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "inputVec4": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "inputScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
           "name": "MultiHeadAttention.DecodeSplitKBias",
           "shader": "attn-flash-decode-splitk.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
+          "bindings": ["query_query_t", "key_key_t", "value_value_t", "bias", "partial_out", "partial_stats", "params_kv_seq_scale"],
+          "dispatch": { "x": "biasSplitKCount", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
         },
         {
           "id": "merge",
           "name": "MultiHeadAttention.DecodeSplitKMergeBias",
           "shader": "attn-flash-decode-splitk-merge.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
+          "bindings": ["partial_out_merge", "partial_stats_merge", "bias", "output_merge"],
           "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
         }
       ]
       "derive": {
         "qNumHeads": "attrs.num_heads",
         "kvNumHeads": "attrs.num_heads",
         "materializedInnerTile": "tunables.MATERIALIZED_INNER_TILE",
         "materializedQueryTile": "tunables.MATERIALIZED_QUERY_TILE",
         "materializedKeyTile": "tunables.MATERIALIZED_KEY_TILE",
         "materializedSoftmaxCols4": "dim(shapes.keyT, 1) / 4",
         "softmaxElementType": "\"vec4<f32>\" if materializedCachedSoftmaxOk else \"f32\"",
         "applyTileN": "materializedApplyTileN",
         "useSubgroups": "device.features.has(\"subgroups\")"
       },
       "intermediates": [
           "name": "MultiHeadAttention.MaterializedScores",
           "shader": "attn-materialized-score-f32.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
+          "bindings": ["query_f32", "key_f32", "scores", "params_scores"],
           "dispatch": {
             "x": "ceilDiv(dim(shapes.keyT, 1), tunables.MATERIALIZED_KEY_TILE)",
             "y": "ceilDiv(dim(shapes.queryT, 1), tunables.MATERIALIZED_QUERY_TILE)",
           "name": "MultiHeadAttention.MaterializedSoftmax",
           "shader": "attn-materialized-softmax-f32.wgsl.jinja",
           "derive": { "cacheVec4": "materializedCachedSoftmaxOk" },
+          "bindings": ["scores_softmax", "params_rows"],
           "dispatch": {
             "x": "min(dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1), 65535)",
             "y": "ceilDiv(dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1), 65535)",
             "z": 1
+          }
         },
         {
           "id": "apply",
           "name": "MultiHeadAttention.MaterializedApply",
           "shader": "attn-materialized-apply-f32.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
+          "bindings": ["scores_f32", "value_f32", "output_f32", "params_apply"],
           "dispatch": {
             "x": "ceilDiv(headDim, applyTileN)",
             "y": "ceilDiv(dim(shapes.queryT, 1), tunables.MATERIALIZED_QUERY_TILE)",
       "derive": {
         "qNumHeads": "attrs.num_heads",
         "kvNumHeads": "attrs.num_heads",
         "materializedInnerTile": "tunables.MATERIALIZED_INNER_TILE",
         "materializedQueryTile": "tunables.MATERIALIZED_QUERY_TILE",
         "materializedKeyTile": "tunables.MATERIALIZED_KEY_TILE",
         "materializedSoftmaxCols4": "dim(shapes.keyT, 1) / 4",
         "softmaxElementType": "\"vec4<f32>\" if materializedCachedSoftmaxOk else \"f32\"",
         "applyTileN": "materializedApplyTileN",
         "useSubgroups": "device.features.has(\"subgroups\")"
       },
       "intermediates": [
           "name": "MultiHeadAttention.MaterializedScoresBias",
           "shader": "attn-materialized-score-f32.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
+          "bindings": ["query_f32", "key_f32", "bias_f32", "scores", "params_scores"],
           "dispatch": {
             "x": "ceilDiv(dim(shapes.keyT, 1), tunables.MATERIALIZED_KEY_TILE)",
             "y": "ceilDiv(dim(shapes.queryT, 1), tunables.MATERIALIZED_QUERY_TILE)",
           "name": "MultiHeadAttention.MaterializedSoftmaxBias",
           "shader": "attn-materialized-softmax-f32.wgsl.jinja",
           "derive": { "cacheVec4": "materializedCachedSoftmaxOk" },
+          "bindings": ["scores_softmax", "params_rows"],
           "dispatch": {
             "x": "min(dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1), 65535)",
             "y": "ceilDiv(dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1), 65535)",
             "z": 1
+          }
         },
         {
           "id": "apply",
           "name": "MultiHeadAttention.MaterializedApplyBias",
           "shader": "attn-materialized-apply-f32.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
+          "bindings": ["scores_f32", "value_f32", "bias_f32", "output_f32", "params_apply"],
           "dispatch": {
             "x": "ceilDiv(headDim, applyTileN)",
             "y": "ceilDiv(dim(shapes.queryT, 1), tunables.MATERIALIZED_QUERY_TILE)",
       "derive": {
         "qNumHeads": "attrs.num_heads",
         "kvNumHeads": "attrs.num_heads",
         "materializedInnerTile": "tunables.MATERIALIZED_INNER_TILE",
         "materializedQueryTile": "tunables.MATERIALIZED_QUERY_TILE",
         "materializedKeyTile": "tunables.MATERIALIZED_KEY_TILE",
         "materializedWorkgroupDim": "tunables.MATERIALIZED_WORKGROUP_DIM",
         "applyTileN": "materializedApplyTileN",
         "statSlots": "materializedGemmStatSlots",
         "statQuerySeq": "dim(shapes.queryT, 1)"
       },
           "name": "MultiHeadAttention.MaterializedScoresFused",
           "shader": "attn-materialized-score-f32.wgsl.jinja",
           "derive": { "layout": "\"bsh\"", "emitRowStats": true },
+          "bindings": ["query_f32", "key_f32", "scores", "scorePartials", "params_scores"],
           "dispatch": {
             "x": "ceilDiv(dim(shapes.keyT, 1), tunables.MATERIALIZED_KEY_TILE)",
             "y": "ceilDiv(dim(shapes.queryT, 1), tunables.MATERIALIZED_QUERY_TILE)",
           "name": "MultiHeadAttention.MaterializedGemmRowStatsCombine",
           "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja",
           "derive": { "maxOnly": true },
+          "bindings": ["scorePartials_f32", "rowStats", "params_rows"],
           "dispatch": {
             "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)), (materializedRowStatsWg)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
             "y": 1,
           "name": "MultiHeadAttention.MaterializedApplyFused",
           "shader": "attn-materialized-apply-f32.wgsl.jinja",
           "derive": { "layout": "\"bsh\"", "fusedSoftmax": true },
+          "bindings": ["scores_f32", "value_f32", "output_f32", "rowStats_f32", "params_apply"],
           "dispatch": {
             "x": "ceilDiv(headDim, applyTileN)",
             "y": "ceilDiv(dim(shapes.queryT, 1), tunables.MATERIALIZED_QUERY_TILE)",
       "derive": {
         "qNumHeads": "attrs.num_heads",
         "kvNumHeads": "attrs.num_heads",
         "materializedInnerTile": "tunables.MATERIALIZED_INNER_TILE",
         "materializedQueryTile": "tunables.MATERIALIZED_QUERY_TILE",
         "materializedKeyTile": "tunables.MATERIALIZED_KEY_TILE",
         "materializedWorkgroupDim": "tunables.MATERIALIZED_WORKGROUP_DIM",
         "applyTileN": "materializedApplyTileN",
         "statSlots": "materializedGemmStatSlots",
         "statQuerySeq": "dim(shapes.queryT, 1)"
       },
           "name": "MultiHeadAttention.MaterializedScoresBiasFused",
           "shader": "attn-materialized-score-f32.wgsl.jinja",
           "derive": { "layout": "\"bsh\"", "emitRowStats": true },
+          "bindings": ["query_f32", "key_f32", "bias_f32", "scores", "scorePartials", "params_scores"],
           "dispatch": {
             "x": "ceilDiv(dim(shapes.keyT, 1), tunables.MATERIALIZED_KEY_TILE)",
             "y": "ceilDiv(dim(shapes.queryT, 1), tunables.MATERIALIZED_QUERY_TILE)",
           "name": "MultiHeadAttention.MaterializedGemmRowStatsCombineBias",
           "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja",
           "derive": { "maxOnly": true },
+          "bindings": ["scorePartials_f32", "rowStats", "params_rows"],
           "dispatch": {
             "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)), (materializedRowStatsWg)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
             "y": 1,
           "name": "MultiHeadAttention.MaterializedApplyBiasFused",
           "shader": "attn-materialized-apply-f32.wgsl.jinja",
           "derive": { "layout": "\"bsh\"", "fusedSoftmax": true },
+          "bindings": ["scores_f32", "value_f32", "bias_f32", "output_f32", "rowStats_f32", "params_apply"],
           "dispatch": {
             "x": "ceilDiv(headDim, applyTileN)",
             "y": "ceilDiv(dim(shapes.queryT, 1), tunables.MATERIALIZED_QUERY_TILE)",
     {
       "id": "qkv_no_bias_flash_q32_broadcast",
       "priority": 30,
+      "when": ["tensorDtypes.queryT == \"float16\"", "not present.biasT", "flashShapeOk", "(dim(shapes.queryT, 2) / attrs.num_heads) % 4 == 0", "(dim(shapes.queryT, 2) / attrs.num_heads) >= 64", "(dim(shapes.queryT, 2) / attrs.num_heads) <= 256", "dim(shapes.queryT, 1) >= 31", "wave32SubgroupsUsable"],
       "requires": { "features": ["subgroups", "shader-f16"] },
       "derive": {
         "usesF16": true,
         "scalar": "\"f16\"",
         "inputElement": "\"vec4<f16>\"",
         "outputElement": "\"vec4<f16>\"",
         "qNumHeads": "attrs.num_heads",
             "x": "ceilDiv(dim(shapes.queryT, 1), 32)",
             "y": "attrs.num_heads",
             "z": "dim(shapes.queryT, 0)"
+          }
         }
       ]
     },
       "derive": {
         "usesF16": true,
         "scalar": "\"f16\"",
         "inputElement": "\"vec4<f16>\"",
         "outputElement": "\"vec4<f16>\"",
         "qNumHeads": "attrs.num_heads",
       "priority": 0,
       "when": ["not present.biasT and fallbackMaskShapeOk"],
       "derive": {
         "hasCausal": true,
         "hasKeyLimit": false,
         "inputScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "outputScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "inputElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
           "name": "MultiHeadAttention",
           "shader": "attn-online-scalar.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
+          "bindings": ["query", "key", "value", "attn_mask_main", "output", "params__uniform"],
           "dispatch": {
             "x": "min(dim(shapes.queryT, 1), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
             "y": "attrs.num_heads",
       "priority": 0,
       "when": ["biasOk and fallbackMaskShapeOk"],
       "derive": {
         "hasCausal": true,
         "hasKeyLimit": false,
         "inputScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "outputScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "inputElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
           "name": "MultiHeadAttention",
           "shader": "attn-online-scalar.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
+          "bindings": ["query", "key", "value", "attn_mask_main", "bias", "output", "params__uniform"],
           "dispatch": {
             "x": "min(dim(shapes.queryT, 1), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
             "y": "attrs.num_heads",
       "priority": 0,
       "when": ["not present.biasT and fallbackShapeOk and not flashShapeOk"],
       "derive": {
         "hasCausal": true,
         "hasKeyLimit": false,
         "inputScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "outputScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "inputElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
       "priority": 0,
       "when": ["biasOk and fallbackShapeOk and not flashShapeOk"],
       "derive": {
         "hasCausal": true,
         "hasKeyLimit": false,
         "inputScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "outputScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "inputElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
       "priority": 20,
       "when": ["device.features.has(\"subgroups\")", "not present.biasT", "flashShapeOk"],
       "derive": {
         "hasCausal": true,
         "combineSubgroups": true,
         "maskElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "inputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "outputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "qNumHeads": "attrs.num_heads",
           "shader": "attn-flash-online.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
           "bindings": ["query", "key", "value", "output", "params"],
+          "dispatch": { "x": "dim(shapes.queryT, 1)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
         }
       ]
     },
       "priority": 20,
       "when": ["device.features.has(\"subgroups\")", "biasOk", "flashShapeOk"],
       "derive": {
         "hasCausal": true,
         "combineSubgroups": true,
         "maskElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "inputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "outputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "inputScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
           "shader": "attn-flash-online.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
           "bindings": ["query", "key", "value", "bias", "output", "params"],
+          "dispatch": { "x": "dim(shapes.queryT, 1)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
         }
       ]
     },
       "priority": 18,
       "when": ["true", "not present.biasT", "flashShapeOk"],
       "derive": {
         "hasCausal": true,
         "combineSubgroups": false,
         "maskElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "inputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "outputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "qNumHeads": "attrs.num_heads",
       "priority": 17,
       "when": ["true", "biasOk", "flashShapeOk"],
       "derive": {
         "hasCausal": true,
         "combineSubgroups": false,
         "maskElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "inputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "outputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "inputScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
       "priority": 20,
       "when": ["device.features.has(\"subgroups\")", "not present.biasT", "flashMaskShapeOk"],
       "derive": {
         "hasCausal": true,
         "combineSubgroups": true,
         "maskElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "inputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "outputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "qNumHeads": "attrs.num_heads",
           "name": "MultiHeadAttention.Flash",
           "shader": "attn-flash-online.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
+          "bindings": ["query", "key", "value", "attn_mask_main", "output", "params__uniform"],
+          "dispatch": { "x": "dim(shapes.queryT, 1)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
         }
       ]
     },
       "priority": 20,
       "when": ["device.features.has(\"subgroups\")", "biasOk", "flashMaskShapeOk"],
       "derive": {
         "hasCausal": true,
         "combineSubgroups": true,
         "maskElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "inputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "outputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "inputScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
           "name": "MultiHeadAttention.Flash",
           "shader": "attn-flash-online.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
+          "bindings": ["query", "key", "value", "attn_mask_main", "bias", "output", "params__uniform"],
+          "dispatch": { "x": "dim(shapes.queryT, 1)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
         }
       ]
     },
       "priority": 18,
       "when": ["true", "not present.biasT", "flashMaskShapeOk"],
       "derive": {
         "hasCausal": true,
         "combineSubgroups": false,
         "maskElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "inputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "outputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "qNumHeads": "attrs.num_heads",
           "name": "MultiHeadAttention.NoBiasOnlineFlashNoSg",
           "shader": "attn-flash-online.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
+          "bindings": ["query", "key", "value", "attn_mask_main", "output", "params__uniform"],
           "dispatch": { "x": "dim(shapes.queryT, 1)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
         }
       ]
       "priority": 17,
       "when": ["true", "biasOk", "flashMaskShapeOk"],
       "derive": {
         "hasCausal": true,
         "combineSubgroups": false,
         "maskElement": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "scalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
         "inputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "outputElement": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"",
         "inputScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"",
           "name": "MultiHeadAttention.BiasOnlineFlashNoSg",
           "shader": "attn-flash-online.wgsl.jinja",
           "derive": { "layout": "\"bsh\"" },
+          "bindings": ["query", "key", "value", "attn_mask_main", "bias", "output", "params__uniform"],
           "dispatch": { "x": "dim(shapes.queryT, 1)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }
         }
       ]
       },
       "derive": {
         "qNumHeads": "attrs.num_heads",
         "qHidden": "dim(shapes.queryT, 2)",
         "materializedSoftmaxWg": "materializedCachedSoftmaxWg if materializedCachedSoftmaxOk else tunables.MATERIALIZED_SOFTMAX_WORKGROUP_SIZE",
         "materializedSoftmaxCols": "dim(shapes.keyT, 1)",
         "materializedSoftmaxCols4": "dim(shapes.keyT, 1) / 4",
         "softmaxElementType": "\"vec4<f32>\" if materializedCachedSoftmaxOk else \"f32\"",
         "useSubgroups": true
       },
       "intermediates": [
           "name": "MultiHeadAttention.MaterializedScoresSgmat",
           "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
           "derive": { "phase": "\"score\"" },
+          "bindings": ["query_f32", "key_f32", "scores", "params_scores"],
           "dispatch": {
             "x": "ceilDiv(dim(shapes.keyT, 1), materializedSgmatKeyTile)",
             "y": "ceilDiv(dim(shapes.queryT, 1), materializedSgmatQueryTile)",
           "name": "MultiHeadAttention.MaterializedSoftmaxSgmat",
           "shader": "attn-materialized-softmax-f32.wgsl.jinja",
           "derive": { "cacheVec4": "materializedCachedSoftmaxOk" },
+          "bindings": ["scores_softmax", "params_rows"],
           "dispatch": {
             "x": "min(dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1), 65535)",
             "y": "ceilDiv(dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1), 65535)",
             "z": 1
+          }
         },
         {
           "id": "apply",
           "name": "MultiHeadAttention.MaterializedApplySgmat",
           "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
           "derive": { "phase": "\"apply\"" },
+          "bindings": ["scores_f32", "value_f32", "output_f32", "params_apply"],
           "dispatch": {
             "x": "ceilDiv(headDim, materializedSgmatKeyTile)",
             "y": "ceilDiv(dim(shapes.queryT, 1), materializedSgmatQueryTile)",
       },
       "derive": {
         "qNumHeads": "attrs.num_heads",
         "qHidden": "dim(shapes.queryT, 2)",
         "materializedSoftmaxWg": "materializedCachedSoftmaxWg if materializedCachedSoftmaxOk else tunables.MATERIALIZED_SOFTMAX_WORKGROUP_SIZE",
         "materializedSoftmaxCols": "dim(shapes.keyT, 1)",
         "materializedSoftmaxCols4": "dim(shapes.keyT, 1) / 4",
         "softmaxElementType": "\"vec4<f32>\" if materializedCachedSoftmaxOk else \"f32\"",
         "useSubgroups": true
       },
       "intermediates": [
           "name": "MultiHeadAttention.MaterializedScoresSgmatBias",
           "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
           "derive": { "phase": "\"score\"" },
+          "bindings": ["query_f32", "key_f32", "bias_f32", "scores", "params_scores"],
           "dispatch": {
             "x": "ceilDiv(dim(shapes.keyT, 1), materializedSgmatKeyTile)",
             "y": "ceilDiv(dim(shapes.queryT, 1), materializedSgmatQueryTile)",
           "name": "MultiHeadAttention.MaterializedSoftmaxSgmatBias",
           "shader": "attn-materialized-softmax-f32.wgsl.jinja",
           "derive": { "cacheVec4": "materializedCachedSoftmaxOk" },
+          "bindings": ["scores_softmax", "params_rows"],
           "dispatch": {
             "x": "min(dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1), 65535)",
             "y": "ceilDiv(dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1), 65535)",
             "z": 1
+          }
         },
         {
           "id": "apply",
           "name": "MultiHeadAttention.MaterializedApplySgmatBias",
           "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
           "derive": { "phase": "\"apply\"" },
+          "bindings": ["scores_f32", "value_f32", "bias_f32", "output_f32", "params_apply"],
           "dispatch": {
             "x": "ceilDiv(headDim, materializedSgmatKeyTile)",
             "y": "ceilDiv(dim(shapes.queryT, 1), materializedSgmatQueryTile)",

build/webgpu/metadata.json CHANGED Viewed

@@ -1,36 +1,39 @@
 {
   "name": "com.microsoft.MultiHeadAttention",
-  "id": "_com_microsoft_multiheadattention_webgpu_5144476",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "attention-rank4-tiled.wgsl.jinja": "gYl58ecIVTxJJkvnK/6lNcbPKaRT24cfz5Yx0MxxxxE=",
-      "attn-flash-decode-splitk-merge.wgsl.jinja": "XSAMF/uzQKJg9pZZDeqmViaVXuj47Vd7Pk404Z1J48I=",
-      "attn-flash-decode-splitk.wgsl.jinja": "dqxX8taw1tpo/QQkFr2zJ7TEEtQtOspTYEYHADXtESs=",
-      "attn-flash-online.wgsl.jinja": "hCxO/6VGkyixwyoT/kuYeHgcEkFoRwk/ZmiLlnwa0jM=",
-      "attn-flash-prefill-cluster.wgsl.jinja": "igbsg5xRuFz7alnwmL74hhh1JUqfDB5PfUslsEz5rjA=",
-      "attn-flash-q32-broadcast.wgsl.jinja": "8qgKXydwiH2shac/fkdHi+FR+c8pOMDFbNp9Zs8KAnc=",
-      "attn-materialized-apply-f32.wgsl.jinja": "aF7JSJj8rQN8wGTV7IWmxkqMDbXOnGBSI3tsI1CxZNk=",
-      "attn-materialized-rowstats-combine-f32.wgsl.jinja": "+/5m03sN44zTn9q7cquDea2UNhEO1ujzm1J7yPai4OY=",
-      "attn-materialized-score-f32.wgsl.jinja": "gCvBC1T1CLEmNxcpo04CLBQdFBkEVMJ9Xhv1wbQxCsI=",
-      "attn-materialized-sgmat-f32.wgsl.jinja": "Ez9UJggPXMZ6OyzUQu8xLtg8bZd+lTCoEGDwLkNHNL0=",
-      "attn-materialized-softmax-f32.wgsl.jinja": "DCuz4gFKXDCdR9UMnfD+mYFjBghnv12GxtEsPKH52mE=",
-      "attn-online-scalar.wgsl.jinja": "6+hU7MzwHz69fjciv94xsPxN+o9rwm42u9Cin2C+AL8=",
-      "attn-small-head-parallel.wgsl.jinja": "oxFE6hQZKTRugz9hwuS+cob00PskT8iYqP4e+YWtHKA=",
-      "bench.json": "b/WdaAhEP5NpRXNB5w1mivRNTfALrBDJisDONxTf4Yo=",
-      "manifest.json": "+/bNCAjkG5OoutFq7O4vQJqb2yMxzICsKvFXsRE50bA=",
-      "mha-small-seq-blocked.wgsl.jinja": "zznGERrPco5uH9tKiQ2QE+ArO3vho/nEv69wT8IgY1U=",
-      "mha-small-seq.wgsl.jinja": "UsbUIkvgWL/RgOBSm7lhPGIatA36yTQR4Z2gpHVLqx4=",
-      "test.json": "NlMA/wlmGSpogAuhSZ/kcQ01dd3M+CvkPF0d/VBmIpM="
     }
   },
-  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
   "webgpu": {
-    "manifestSpec": "2.0",
     "variants": {
       "qkv_bias_small_seq_blocked": ["mha-small-seq-blocked.wgsl.jinja"],
       "qkv_no_bias_small_seq_blocked": ["mha-small-seq-blocked.wgsl.jinja"],
       "qkv_no_bias_small_seq": ["mha-small-seq.wgsl.jinja"],

 {
   "name": "com.microsoft.MultiHeadAttention",
+  "id": "_com_microsoft_multiheadattention_webgpu_e942a19",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "attention-rank4-tiled.wgsl.jinja": "084x5bbpxeFR95/Q4nGzLrLnk9JJrd9XA1KJHxwU9CY=",
+      "attn-flash-decode-splitk-merge.wgsl.jinja": "JIhoYhGkuat/Wrvix6ddD19+bmzHnc0yJdwWbW9o+nw=",
+      "attn-flash-decode-splitk.wgsl.jinja": "tDj5VlltkM0N99vjpj6z5wtlIv00benlaAboa4e5vWI=",
+      "attn-flash-online.wgsl.jinja": "1DGpp4kzd/v4gzAviJKnYGCT7BwDomzXK3He6CIT+1E=",
+      "attn-flash-prefill-cluster.wgsl.jinja": "pjDvKLZ3CSluSAPNfb9j3U7yKEy+wVOfOAyZ1jdrjv8=",
+      "attn-flash-q32-broadcast.wgsl.jinja": "puOnD1vGGeVg6HmaUi8k4JAZJviJM1dXWJyMZ2xuURM=",
+      "attn-materialized-apply-f32.wgsl.jinja": "E2Vas8y/zcjHQb4Ks5TPnC4/lAhVRawXR5S1OwJB1Tc=",
+      "attn-materialized-rowstats-combine-f32.wgsl.jinja": "CBltqk0CN9vvscb6D29o1T43vHWZYlU9fFpO9av2Sa8=",
+      "attn-materialized-score-f32.wgsl.jinja": "bxW7HdJcF47NTbyXVHcaWrtCO0xCJ8pW5GW9/FTuGGg=",
+      "attn-materialized-sgmat-f32.wgsl.jinja": "3e4ZOgVD5XSnYjVldRVq5QL72KiloCkr91OZj/VU0ZU=",
+      "attn-materialized-softmax-f32.wgsl.jinja": "bpHIAAuUa+nrniTKxXEScXT5rJr5nObXY/sJ81hby5A=",
+      "attn-online-scalar.wgsl.jinja": "AbpGvmzka15nPEj3QZPyMaW48Lo/9TqtndP0SdyqV/U=",
+      "attn-small-head-parallel.wgsl.jinja": "Y0otPUGeAa4DzQUi+LnWLRTATc1i/awpIipRqbGDoQ0=",
+      "attn-small-head-value.wgsl.jinja": "+6NLq7AJT3lRkAwsW2AvjGE0nhkD7dkBLqh6SpOFjy4=",
+      "bench.json": "ateXBEoOVVMHNfiuXPmDW6sKzU1NRVRA9ndzslgsD2A=",
+      "manifest.json": "VPWToEdcoP2Mr6k2oQuf/BhwUZdbcPl9J5YLZsMBMaE=",
+      "mha-small-seq-blocked.wgsl.jinja": "WudxiAnGnY7dSEnlhHtwvQKvm4AN3ngQ9cbBfI6OUBM=",
+      "mha-small-seq.wgsl.jinja": "1MR+JuEbffrS+QkuAZzFzWCLPX4VX9RJ5lmNIOm0/o0=",
+      "test.json": "IEMAo5OqrMxRC7xc5Y/8xmjWObXuWw1SJhJeGGzDhwc="
     }
   },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
   "webgpu": {
+    "manifestSpec": "2.1",
     "variants": {
+      "qkv_no_bias_small_head_value_subgroups": ["attn-small-head-value.wgsl.jinja"],
+      "qkv_no_bias_small_head_value_tree": ["attn-small-head-value.wgsl.jinja"],
       "qkv_bias_small_seq_blocked": ["mha-small-seq-blocked.wgsl.jinja"],
       "qkv_no_bias_small_seq_blocked": ["mha-small-seq-blocked.wgsl.jinja"],
       "qkv_no_bias_small_seq": ["mha-small-seq.wgsl.jinja"],

build/webgpu/mha-small-seq-blocked.wgsl.jinja CHANGED Viewed

@@ -37,19 +37,20 @@ fn shifted_value(value: f32, maxValue: f32) -> f32 {
   let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
   return select(value - maxValue, 0.0, equalFiniteMax);
 }
 fn exp_shift(value: f32, maxValue: f32) -> f32 {
   return exp(shifted_value(value, maxValue));
 }
-{% if ATTN_SCALE_DIM is not defined %}{% set ATTN_SCALE_DIM = "HEAD_DIM" %}{% endif %}
-fn scale_value() -> f32 {
   if (params.scale != 0.0) { return params.scale; }
   return inverseSqrt(f32({{ ATTN_SCALE_DIM }}));
 }
 {% if hasBias %}
 fn load_bias4(base: u32, d4: u32) -> vec4<f32> {
   let offset = base + d4 * 4u;
-  return vec4<f32>(bias[offset], bias[offset + 1u], bias[offset + 2u], bias[offset + 3u]);
 }
 {% endif %}

   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));
 }
+{% set ATTN_SCALE_DIM = "HEAD_DIM" %}fn scale_value() -> f32 {
   if (params.scale != 0.0) { return params.scale; }
   return inverseSqrt(f32({{ ATTN_SCALE_DIM }}));
 }
 {% if hasBias %}
+{% set BW = "" %}
+{% set BC = "" %}
 fn load_bias4(base: u32, d4: u32) -> vec4<f32> {
   let offset = base + d4 * 4u;
+  return vec4<f32>({{ BW }}bias[offset]{{ BC }}, {{ BW }}bias[offset + 1u]{{ BC }}, {{ BW }}bias[offset + 2u]{{ BC }}, {{ BW }}bias[offset + 3u]{{ BC }});
 }
 {% endif %}

build/webgpu/mha-small-seq.wgsl.jinja CHANGED Viewed

@@ -1,4 +1,26 @@
 {{ env.wgsl.resourceDeclarations }}
 // Whole-head attention for tiny bidirectional sequences without bias or mask.
 // One workgroup owns one (head, batch) and stages the complete K and V planes in
@@ -9,13 +31,11 @@ const HEAD_DIM: u32 = {{ headDim }}u;
 const KV_SEQ: u32 = {{ kvSeq }}u;
 const HIDDEN: u32 = {{ hidden }}u;
 const WG: u32 = {{ workgroupSize }}u;
-{% if ATTN_SCALE_DIM is not defined %}{% set ATTN_SCALE_DIM = "HEAD_DIM" %}{% endif %}
-fn scale_value() -> f32 {
   if (params.scale != 0.0) { return params.scale; }
   return inverseSqrt(f32({{ ATTN_SCALE_DIM }}));
 }
 var<workgroup> kShared: array<f32, KV_SEQ * HEAD_DIM>;
 var<workgroup> vShared: array<f32, KV_SEQ * HEAD_DIM>;
@@ -56,11 +76,11 @@ fn main(
   var scores: array<f32, KV_SEQ>;
   var maxScore = -3.4028234663852886e38;
   for (var k = 0u; k < KV_SEQ; k++) {
-    var dot = 0.0;
     for (var d = 0u; d < HEAD_DIM; d++) {
-      dot += qRow[d] * kShared[k * HEAD_DIM + d];
     }
-    let s = dot * scale;
     scores[k] = s;
     maxScore = max(maxScore, s);
   }

 {{ env.wgsl.resourceDeclarations }}
+{% set dotType = "f32" %}
+// Retain product and addition residuals across a dot product. Explicit fma
+// boundaries preserve the addition error transform under reassociation.
+struct DotAccumulator {
+  hi: {{ dotType }},
+  lo: {{ dotType }},
+}
+fn dot_accumulate(acc: DotAccumulator, a: {{ dotType }}, b: {{ dotType }}) -> DotAccumulator {
+  let product = fma(a, b, {{ dotType }}(0.0));
+  let productError = fma(a, b, -product);
+  let sum = fma(acc.hi, {{ dotType }}(1.0), product);
+  let bv = fma({{ dotType }}(-1.0), acc.hi, sum);
+  let av = fma({{ dotType }}(-1.0), bv, sum);
+  let ae = fma({{ dotType }}(-1.0), av, acc.hi);
+  let be = fma({{ dotType }}(-1.0), bv, product);
+  let error = fma(fma(fma(ae, {{ dotType }}(1.0), be), {{ dotType }}(1.0), acc.lo), {{ dotType }}(1.0), productError);
+  let hi = fma(sum, {{ dotType }}(1.0), error);
+  return DotAccumulator(hi, fma({{ dotType }}(-1.0), fma({{ dotType }}(-1.0), sum, hi), error));
+}
+fn dot_value(acc: DotAccumulator) -> {{ dotType }} {
+  return fma(acc.hi, {{ dotType }}(1.0), acc.lo);
+}
 // Whole-head attention for tiny bidirectional sequences without bias or mask.
 // One workgroup owns one (head, batch) and stages the complete K and V planes in
 const KV_SEQ: u32 = {{ kvSeq }}u;
 const HIDDEN: u32 = {{ hidden }}u;
 const WG: u32 = {{ workgroupSize }}u;
+{% set ATTN_SCALE_DIM = "HEAD_DIM" %}fn scale_value() -> f32 {
   if (params.scale != 0.0) { return params.scale; }
   return inverseSqrt(f32({{ ATTN_SCALE_DIM }}));
 }
 var<workgroup> kShared: array<f32, KV_SEQ * HEAD_DIM>;
 var<workgroup> vShared: array<f32, KV_SEQ * HEAD_DIM>;
   var scores: array<f32, KV_SEQ>;
   var maxScore = -3.4028234663852886e38;
   for (var k = 0u; k < KV_SEQ; k++) {
+    var dot = DotAccumulator(0.0, 0.0);
     for (var d = 0u; d < HEAD_DIM; d++) {
+      dot = dot_accumulate(dot, qRow[d], kShared[k * HEAD_DIM + d]);
     }
+    let s = dot_value(dot) * scale;
     scores[k] = s;
     maxScore = max(maxScore, s);
   }

build/webgpu/test.json CHANGED Viewed

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