Xenova HF Staff commited on 4 days ago

Commit

ba9df6e

verified ·

1 Parent(s): 78d4fcb

sync 6fdf6301e2bb

Browse files

Files changed (17) hide show

README.md +16 -28
build/webgpu/attention-rank4-tiled.wgsl.jinja +1 -25
build/webgpu/attn-flash-decode-splitk-merge.wgsl.jinja +5 -34
build/webgpu/attn-flash-decode-splitk.wgsl.jinja +21 -109
build/webgpu/attn-flash-online.wgsl.jinja +15 -95
build/webgpu/attn-flash-prefill-cluster.wgsl.jinja +112 -100
build/webgpu/attn-flash-q32-broadcast.wgsl.jinja +47 -8
build/webgpu/attn-materialized-rowstats-combine-f32.wgsl.jinja +8 -2
build/webgpu/attn-materialized-sgmat-f32.wgsl.jinja +37 -64
build/webgpu/attn-online-scalar.wgsl.jinja +53 -28
build/webgpu/bench.json +603 -6
build/webgpu/gqa-attention.wgsl.jinja +55 -38
build/webgpu/gqa-present.wgsl.jinja +70 -42
build/webgpu/gqa-qprep.wgsl.jinja +7 -5
build/webgpu/manifest.json +0 -0
build/webgpu/metadata.json +27 -37
build/webgpu/test.json +0 -0

README.md CHANGED Viewed

@@ -87,19 +87,16 @@ One implementation is selected per call from the device capabilities, the reques
 - `window_shift_materialized_sgmat_f32` — Materializes causal chunked-prefill attention with a windowed cache. The shift pass compacts surviving rows and appends the chunk; score and apply passes enforce both the sliding-window floor and causal bound.
 - `share_append_materialized_sgmat_f32` — Materializes causal chunked-prefill attention while updating a shared-capacity cache in place. Score and apply passes bound every causal tile by the live length from `seqlens_k`, not the buffer capacity, so right-padded batches remain left-aligned.
 - `new_kv_share_append_split` — Retains the existing cache and appends new key/value rows in separate passes before portable attention. It avoids rebuilding unchanged cache positions when the past and present allocations share capacity.
 - `qkv_present_tiled_nosg` — Portable tiled prefill route that computes attention online and writes the present cache separately. It is used when the flash shape is valid but no suitable subgroup route is admissible.
 - `quant_int8_decode_splitk` — Appends int8-quantized key/value rows, partitions cached decode across the key axis, and merges partial online-softmax results. Used when one workgroup per head exposes too little independent work.
 - `qkv_present_flash_splitk` — Partitions direct Q/K/V attention across the key axis, merges partial online-softmax results, and writes the present cache separately. It serves short-query shapes needing more key-axis parallelism.
 - `qkv_present_flash_cluster` — Computes clustered online-softmax prefill directly from Q/K/V and writes the present cache separately. The family provides subgroup and portable reductions for the same tiled algorithm.
-- `quant_int8_decode_splitk_nosg` — Appends int8-quantized key/value rows, partitions cached decode across the key axis, and merges partial online-softmax results. Used when one workgroup per head exposes too little independent work.
-- `qkv_present_flash_splitk_nosg` — Partitions direct Q/K/V attention across the key axis, merges partial online-softmax results, and writes the present cache separately. It serves short-query shapes needing more key-axis parallelism.
-- `qkv_present_flash_cluster_nosg` — Computes clustered online-softmax prefill directly from Q/K/V and writes the present cache separately. The family provides subgroup and portable reductions for the same tiled algorithm.
 - `past_kv_decode_splitk` — Updates the float cache in copy, merge, or window-shift mode, then partitions cached decode across the key axis and combines partial online-softmax results. The family includes subgroup and portable forms.
 - `new_kv_past_decode_splitk` — Updates the float cache in copy, merge, or window-shift mode, then partitions cached decode across the key axis and combines partial online-softmax results. The family includes subgroup and portable forms.
 - `window_shift_decode_splitk` — Updates the float cache in copy, merge, or window-shift mode, then partitions cached decode across the key axis and combines partial online-softmax results. The family includes subgroup and portable forms.
-- `past_kv_decode_splitk_nosg` — Updates the float cache in copy, merge, or window-shift mode, then partitions cached decode across the key axis and combines partial online-softmax results. The family includes subgroup and portable forms.
-- `new_kv_past_decode_splitk_nosg` — Updates the float cache in copy, merge, or window-shift mode, then partitions cached decode across the key axis and combines partial online-softmax results. The family includes subgroup and portable forms.
-- `window_shift_decode_splitk_nosg` — Updates the float cache in copy, merge, or window-shift mode, then partitions cached decode across the key axis and combines partial online-softmax results. The family includes subgroup and portable forms.
 - `past_kv_flash_prefill` — Updates the float cache in the selected mode and applies clustered online-softmax causal prefill with the configured cache, mask, rotary, softcap, bias, or head-sink features.
 - `new_kv_past_flash_prefill` — Updates the float cache in the selected mode and applies clustered online-softmax causal prefill with the configured cache, mask, rotary, softcap, bias, or head-sink features.
 - `past_kv_rotary_flash_prefill` — Updates the float cache in the selected mode and applies clustered online-softmax causal prefill with the configured cache, mask, rotary, softcap, bias, or head-sink features.
@@ -107,21 +104,14 @@ One implementation is selected per call from the device capabilities, the reques
 - `past_kv_headsink_flash_prefill` — Updates the float cache in the selected mode and applies clustered online-softmax causal prefill with the configured cache, mask, rotary, softcap, bias, or head-sink features.
 - `past_kv_bias_headsink_flash_prefill` — Updates the float cache in the selected mode and applies clustered online-softmax causal prefill with the configured cache, mask, rotary, softcap, bias, or head-sink features.
 - `window_shift_flash_prefill` — Updates the float cache in the selected mode and applies clustered online-softmax causal prefill with the configured cache, mask, rotary, softcap, bias, or head-sink features.
-- `past_kv_flash_prefill_nosg` — Updates the float cache in the selected mode and applies clustered online-softmax causal prefill with the configured cache, mask, rotary, softcap, bias, or head-sink features.
-- `new_kv_past_flash_prefill_nosg` — Updates the float cache in the selected mode and applies clustered online-softmax causal prefill with the configured cache, mask, rotary, softcap, bias, or head-sink features.
-- `past_kv_rotary_flash_prefill_nosg` — Updates the float cache in the selected mode and applies clustered online-softmax causal prefill with the configured cache, mask, rotary, softcap, bias, or head-sink features.
-- `past_kv_softcap_flash_prefill_nosg` — Updates the float cache in the selected mode and applies clustered online-softmax causal prefill with the configured cache, mask, rotary, softcap, bias, or head-sink features.
-- `past_kv_headsink_flash_prefill_nosg` — Updates the float cache in the selected mode and applies clustered online-softmax causal prefill with the configured cache, mask, rotary, softcap, bias, or head-sink features.
-- `past_kv_bias_headsink_flash_prefill_nosg` — Updates the float cache in the selected mode and applies clustered online-softmax causal prefill with the configured cache, mask, rotary, softcap, bias, or head-sink features.
-- `window_shift_flash_prefill_nosg` — Updates the float cache in the selected mode and applies clustered online-softmax causal prefill with the configured cache, mask, rotary, softcap, bias, or head-sink features.
 - `quant_int8_flash_prefill` — Builds an int8 or int4 cache and applies clustered online-softmax prefill directly from the packed values. The family supplies subgroup and portable reductions.
 - `quant_int4_flash_prefill` — Builds an int8 or int4 cache and applies clustered online-softmax prefill directly from the packed values. The family supplies subgroup and portable reductions.
-- `quant_int8_flash_prefill_nosg` — Builds an int8 or int4 cache and applies clustered online-softmax prefill directly from the packed values. The family supplies subgroup and portable reductions.
-- `quant_int4_flash_prefill_nosg` — Builds an int8 or int4 cache and applies clustered online-softmax prefill directly from the packed values. The family supplies subgroup and portable reductions.
 - `share_append_split_decode_splitk` — Retains shared-capacity cache rows, appends new rows separately, then partitions decode across the key axis. It avoids rebuilding unchanged cache positions while exposing split-key parallelism.
-- `share_append_split_decode_splitk_nosg` — Retains shared-capacity cache rows, appends new rows separately, then partitions decode across the key axis. It avoids rebuilding unchanged cache positions while exposing split-key parallelism.
 - `share_append_split_flash_prefill` — Retains shared-capacity cache rows, appends new rows separately, then applies clustered causal prefill. It avoids rebuilding unchanged cache positions and includes subgroup and portable forms.
-- `share_append_split_flash_prefill_nosg` — Retains shared-capacity cache rows, appends new rows separately, then applies clustered causal prefill. It avoids rebuilding unchanged cache positions and includes subgroup and portable forms.
 ## Device requirements
@@ -132,7 +122,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)
@@ -149,7 +139,7 @@ Some implementation variants require `subgroup-matrix`, `shader-f16`, and `subgr
 ## Use with `@huggingface/kernels`
 ```sh
-npm install --save-exact @huggingface/kernels@0.0.1-preview.2
 ```
 Outputs with inferable metadata are allocated automatically. Explicit `outputs` entries request optional results or provide metadata that cannot be inferred from the supplied inputs and attributes.
@@ -170,18 +160,16 @@ import { getKernel } from "@huggingface/kernels";
 const kernel = await getKernel("webgpu-kernels/com.microsoft.GroupQueryAttention", { version: 1 });
 // Explicit destinations request optional results or supply metadata that cannot be inferred.
 const { outputT, presentKeyT, presentValueT } = await kernel({
-  queryT: { data: queryTData, shape: [2, 1, 16] },
-  keyT: { data: keyTData, shape: [2, 1, 8] },
-  valueT: { data: valueTData, shape: [2, 1, 8] },
-  pastKeyT: { data: pastKeyTData, shape: [2, 1, 8, 8] },
-  pastValueT: { data: pastValueTData, shape: [2, 1, 8, 8] },
-  seqlensKT: { data: seqlensKTData, shape: [2] },
   totalSequenceLengthT: { data: totalSequenceLengthTData, shape: [1] },
 }, {
-  attrs: { num_heads: 2, kv_num_heads: 1 },
   outputs: {
-    presentKeyT: { shape: [2, 1, 8, 8], dtype: "float32" },
-    presentValueT: { shape: [2, 1, 8, 8], dtype: "float32" },
   },
 });
 ```

 - `window_shift_materialized_sgmat_f32` — Materializes causal chunked-prefill attention with a windowed cache. The shift pass compacts surviving rows and appends the chunk; score and apply passes enforce both the sliding-window floor and causal bound.
 - `share_append_materialized_sgmat_f32` — Materializes causal chunked-prefill attention while updating a shared-capacity cache in place. Score and apply passes bound every causal tile by the live length from `seqlens_k`, not the buffer capacity, so right-padded batches remain left-aligned.
 - `new_kv_share_append_split` — Retains the existing cache and appends new key/value rows in separate passes before portable attention. It avoids rebuilding unchanged cache positions when the past and present allocations share capacity.
+- `new_kv_share_append_bidirectional_split` — Bidirectional (`causal = 0`) attention over a shared-capacity cache: the new key/value rows are appended exactly as in the causal sibling, then every query attends to the whole live range `[0, seqlens_k[b] + 1)` instead of stopping at its own position. ONNX Runtime forbids `causal = 0` with a local window, so the window floor never applies here.
 - `qkv_present_tiled_nosg` — Portable tiled prefill route that computes attention online and writes the present cache separately. It is used when the flash shape is valid but no suitable subgroup route is admissible.
+- `qkv_present_flash_causal` — Causal prefill without a past cache and without a sliding window (`causal=1`, `local_window_size=-1`): the spelling onnxruntime's GroupQueryAttention tests use for a first prompt. The twin of `qkv_present_flash` with the causal ceiling enabled; the key sequence equals the query sequence, so the upper-left ceiling is the right one.
+- `qkv_present_causal` — Causal prefill without a past cache and without a sliding window (`causal=1`, `local_window_size=-1`): the spelling onnxruntime's GroupQueryAttention tests use for a first prompt. The twin of `qkv_present` with the causal ceiling enabled; the key sequence equals the query sequence, so the upper-left ceiling is the right one.
 - `quant_int8_decode_splitk` — Appends int8-quantized key/value rows, partitions cached decode across the key axis, and merges partial online-softmax results. Used when one workgroup per head exposes too little independent work.
 - `qkv_present_flash_splitk` — Partitions direct Q/K/V attention across the key axis, merges partial online-softmax results, and writes the present cache separately. It serves short-query shapes needing more key-axis parallelism.
 - `qkv_present_flash_cluster` — Computes clustered online-softmax prefill directly from Q/K/V and writes the present cache separately. The family provides subgroup and portable reductions for the same tiled algorithm.
 - `past_kv_decode_splitk` — Updates the float cache in copy, merge, or window-shift mode, then partitions cached decode across the key axis and combines partial online-softmax results. The family includes subgroup and portable forms.
 - `new_kv_past_decode_splitk` — Updates the float cache in copy, merge, or window-shift mode, then partitions cached decode across the key axis and combines partial online-softmax results. The family includes subgroup and portable forms.
 - `window_shift_decode_splitk` — Updates the float cache in copy, merge, or window-shift mode, then partitions cached decode across the key axis and combines partial online-softmax results. The family includes subgroup and portable forms.
 - `past_kv_flash_prefill` — Updates the float cache in the selected mode and applies clustered online-softmax causal prefill with the configured cache, mask, rotary, softcap, bias, or head-sink features.
 - `new_kv_past_flash_prefill` — Updates the float cache in the selected mode and applies clustered online-softmax causal prefill with the configured cache, mask, rotary, softcap, bias, or head-sink features.
 - `past_kv_rotary_flash_prefill` — Updates the float cache in the selected mode and applies clustered online-softmax causal prefill with the configured cache, mask, rotary, softcap, bias, or head-sink features.
 - `past_kv_headsink_flash_prefill` — Updates the float cache in the selected mode and applies clustered online-softmax causal prefill with the configured cache, mask, rotary, softcap, bias, or head-sink features.
 - `past_kv_bias_headsink_flash_prefill` — Updates the float cache in the selected mode and applies clustered online-softmax causal prefill with the configured cache, mask, rotary, softcap, bias, or head-sink features.
 - `window_shift_flash_prefill` — Updates the float cache in the selected mode and applies clustered online-softmax causal prefill with the configured cache, mask, rotary, softcap, bias, or head-sink features.
+- `window_shift_rotary_append` — Windowed cache with rotary embeddings. The shift pass slides surviving rows down by the eviction count and clears what left the window; the append pass writes the new key/value rows and rotates K at its ABSOLUTE position (window origin plus row, not the cache row); attention rotates Q at the same absolute position and masks cache-relative.
+- `window_shift_rotary_headsink_append` — Windowed rotary cache with a per-head softmax sink (the gpt-oss attention layer). The sink contributes `exp(sink - m)` to the denominator once per query row and nothing to the numerator.
+- `window_shift_rotary_bias_append` — Windowed rotary cache with an additive attention bias. The bias row is indexed by ABSOLUTE key position with the tensor's own last dimension as the stride (ONNX spells it `total_sequence_length`, which keeps counting past the cache capacity), so resident cache row j reads bias column `origin + j`.
 - `quant_int8_flash_prefill` — Builds an int8 or int4 cache and applies clustered online-softmax prefill directly from the packed values. The family supplies subgroup and portable reductions.
 - `quant_int4_flash_prefill` — Builds an int8 or int4 cache and applies clustered online-softmax prefill directly from the packed values. The family supplies subgroup and portable reductions.
 - `share_append_split_decode_splitk` — Retains shared-capacity cache rows, appends new rows separately, then partitions decode across the key axis. It avoids rebuilding unchanged cache positions while exposing split-key parallelism.
 - `share_append_split_flash_prefill` — Retains shared-capacity cache rows, appends new rows separately, then applies clustered causal prefill. It avoids rebuilding unchanged cache positions and includes subgroup and portable forms.
+- `share_append_bidirectional_flash_prefill` — Retains and appends a shared-capacity cache, then applies clustered bidirectional prefill to the full live key range.
 ## Device requirements
 - [`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)
 ## Use with `@huggingface/kernels`
 ```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.3
 ```
 Outputs with inferable metadata are allocated automatically. Explicit `outputs` entries request optional results or provide metadata that cannot be inferred from the supplied inputs and attributes.
 const kernel = await getKernel("webgpu-kernels/com.microsoft.GroupQueryAttention", { version: 1 });
 // Explicit destinations request optional results or supply metadata that cannot be inferred.
 const { outputT, presentKeyT, presentValueT } = await kernel({
+  queryT: { data: queryTData, shape: [1, 2, 8] },
+  keyT: { data: keyTData, shape: [1, 2, 8] },
+  valueT: { data: valueTData, shape: [1, 2, 8] },
+  seqlensKT: { data: seqlensKTData, shape: [1] },
   totalSequenceLengthT: { data: totalSequenceLengthTData, shape: [1] },
 }, {
+  attrs: { num_heads: 1, kv_num_heads: 1 },
   outputs: {
+    presentKeyT: { shape: [1, 1, 2, 8], dtype: "float32" },
+    presentValueT: { shape: [1, 1, 2, 8], dtype: "float32" },
   },
 });
 ```

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
@@ -25,9 +22,6 @@ enable f16;
 const HEAD_DIM_V4: u32 = {{ headDimV4 }}u;
 const Q_HEADS: u32 = {{ qNumHeads }}u;
 const NUM_SPLITS: u32 = {{ numSplits }}u;
-{% if splitQueries %}
-const Q_SEQ: u32 = {{ qSeq }}u;
-{% endif %}
 {% if layout == "bsh" %}
 const Q_HIDDEN_V4: u32 = {{ qHiddenV4 }}u;
 {% endif %}
@@ -57,6 +51,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));
 }
@@ -65,21 +60,14 @@ fn main(
   @builtin(workgroup_id) wg: vec3<u32>,
   @builtin(local_invocation_id) lid: vec3<u32>
 ) {
-{% if splitQueries %}
-  let queryToken = wg.x;
-{% endif %}
   let h = wg.y;
   let b = wg.z;
   let d4 = lid.x;
-  if (h >= Q_HEADS || d4 >= HEAD_DIM_V4{% if splitQueries %} || queryToken >= Q_SEQ{% endif %}) {
     return;
   }
-{% if splitQueries %}
-  let mdBase = ((b * Q_SEQ + queryToken) * Q_HEADS + h) * NUM_SPLITS;
-{% else %}
   let mdBase = (b * Q_HEADS + h) * NUM_SPLITS;
-{% endif %}
   var globalMax = -FLT_MAX;
   for (var s: u32 = 0u; s < NUM_SPLITS; s = s + 1u) {
     globalMax = max(globalMax, partial_stats[mdBase + s].x);
@@ -90,28 +78,16 @@ fn main(
     let stats = partial_stats[mdBase + s];
     let w = exp_shift(stats.x, globalMax);
     globalDenom = globalDenom + stats.y * w;
-{% if splitQueries %}
-    let pBase = (((b * Q_SEQ + queryToken) * Q_HEADS + h) * NUM_SPLITS + s) * HEAD_DIM_V4;
-{% else %}
     let pBase = ((b * Q_HEADS + h) * NUM_SPLITS + s) * HEAD_DIM_V4;
-{% endif %}
     outV = outV + partial_out[pBase + d4] * w;
   }
 {% if layout == "bsh" %}
-{% if splitQueries %}
-  let qBaseV4 = (b * Q_SEQ + queryToken) * Q_HIDDEN_V4 + h * HEAD_DIM_V4;
-{% else %}
   let qBaseV4 = b * Q_HIDDEN_V4 + h * HEAD_DIM_V4;
-{% endif %}
 {% elif layout == "layer_cache" %}
   let qBaseV4 = h * HEAD_DIM_V4;
-{% else %}
-{% if splitQueries %}
-  let qBaseV4 = ((b * Q_HEADS + h) * Q_SEQ + queryToken) * HEAD_DIM_V4;
 {% else %}
   let qBaseV4 = (b * Q_HEADS + h) * HEAD_DIM_V4;
-{% endif %}
 {% endif %}
   // A split that covers no keys writes vec2(-FLT_MAX, 0), so when every split is empty —
   // a query whose window admits nothing — globalMax stays -FLT_MAX, exp_shift(x, x) is 1,
@@ -122,11 +98,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);
 }

+{% set qHiddenV4 = qHiddenV4 | default(0) %}
+{% set qHidden = qHidden | default(0) %}
+{% set hasBias = hasBias is defined and hasBias %}
 {{ env.wgsl.resourceDeclarations }}
 // Split-K flash decode, pass 2 of 2. Combines the per-split un-normalized online
 const HEAD_DIM_V4: u32 = {{ headDimV4 }}u;
 const Q_HEADS: u32 = {{ qNumHeads }}u;
 const NUM_SPLITS: u32 = {{ numSplits }}u;
 {% if layout == "bsh" %}
 const Q_HIDDEN_V4: u32 = {{ qHiddenV4 }}u;
 {% 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));
 }
   @builtin(workgroup_id) wg: vec3<u32>,
   @builtin(local_invocation_id) lid: vec3<u32>
 ) {
   let h = wg.y;
   let b = wg.z;
   let d4 = lid.x;
+  if (h >= Q_HEADS || d4 >= HEAD_DIM_V4) {
     return;
   }
   let mdBase = (b * Q_HEADS + h) * NUM_SPLITS;
   var globalMax = -FLT_MAX;
   for (var s: u32 = 0u; s < NUM_SPLITS; s = s + 1u) {
     globalMax = max(globalMax, partial_stats[mdBase + s].x);
     let stats = partial_stats[mdBase + s];
     let w = exp_shift(stats.x, globalMax);
     globalDenom = globalDenom + stats.y * w;
     let pBase = ((b * Q_HEADS + h) * NUM_SPLITS + s) * HEAD_DIM_V4;
     outV = outV + partial_out[pBase + d4] * w;
   }
 {% if layout == "bsh" %}
   let qBaseV4 = b * Q_HIDDEN_V4 + h * HEAD_DIM_V4;
 {% elif layout == "layer_cache" %}
   let qBaseV4 = h * HEAD_DIM_V4;
 {% else %}
   let qBaseV4 = (b * Q_HEADS + h) * HEAD_DIM_V4;
 {% endif %}
   // A split that covers no keys writes vec2(-FLT_MAX, 0), so when every split is empty —
   // a query whose window admits nothing — globalMax stays -FLT_MAX, exp_shift(x, x) is 1,
   // 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,21 +1,22 @@
 {% 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 layer = layer | default(0) %}
-{% set cacheLen = cacheLen | default(0) %}
-{% 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 %}
 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
@@ -50,9 +51,6 @@ const ATTN_SCALE: f32 = {{ scale }};
 {% endif %}
 const WG: u32 = {{ splitKWorkgroupSize }}u;
 const NUM_SPLITS: u32 = {{ numSplits }}u;
-{% if splitQueries %}
-const Q_SEQ: u32 = {{ qSeq }}u;
-{% endif %}
 // FLT_MAX, not -inf, as the online (m, d) accumulator init: merges must keep
 // `m - m` finite so an empty lane / all--inf row contributes the exact
 // accumulator identity (m, d) = (-FLT_MAX, 0). Operator epilogues interpret
@@ -72,6 +70,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));
 }
@@ -79,7 +78,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 %}
@@ -148,41 +147,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;
@@ -212,12 +179,9 @@ fn combine_partials(m: f32, d: f32, lidx: u32) -> vec2<f32> {
   return merged;
 }
 {% endif %}
-{% endif %}
 {% if layout == "layer_cache" %}{% set ATTN_SCALE_OVERRIDE = "ATTN_SCALE" %}{% 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 %}
@@ -226,7 +190,6 @@ fn scale_value() -> f32 {
 {% endif %}
 }
 {% if quantizedCache %}
 {% macro emit_quant_scale4(kind, scaleBuffer) %}
 fn {{ kind }}scale4(d4: u32, hk: u32) -> vec4<f32> {
@@ -240,28 +203,12 @@ fn {{ kind }}scale4(d4: u32, hk: u32) -> vec4<f32> {
     {{ 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 %}
@@ -279,9 +226,11 @@ fn load_value4(indexV4: u32) -> vec4<f32> {
 // 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 %}
@@ -297,15 +246,10 @@ fn main(
   // choose an intermediate width at pipeline execution time.
   if (sgSize == 0u || sgSize > WG || WG % sgSize != 0u) { return; }
 {% endif %}
-{% if splitQueries %}
-  let queryToken = wg.x / NUM_SPLITS;
-  let split = wg.x % NUM_SPLITS;
-{% else %}
   let split = wg.x;
-{% endif %}
   let h = wg.y;
   let b = wg.z;
-  if (h >= Q_HEADS || split >= NUM_SPLITS{% if splitQueries %} || queryToken >= Q_SEQ{% endif %}{% if layout == "layer_cache" %} || params.past_len >= CACHE_LEN{% endif %}) {
     return;
   }
   let tid = lid.x;
@@ -317,7 +261,7 @@ fn main(
 {% 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 %}
@@ -325,23 +269,15 @@ fn main(
   // Query row (decode uses token zero; short-query prefill folds the token into wg.x).
 {% if layout == "bsh" %}
-{% if splitQueries %}
-  let qBaseV4 = (b * Q_SEQ + queryToken) * Q_HIDDEN_V4 + h * HEAD_DIM_V4;
-{% else %}
   let qBaseV4 = b * Q_HIDDEN_V4 + h * HEAD_DIM_V4;
-{% endif %}
   let kvBaseV4 = b * kvSeq * KV_HIDDEN_V4 + hKv * HEAD_DIM_V4;
   let kvTokenStrideV4 = KV_HIDDEN_V4;
 {% elif layout == "layer_cache" %}
   let qBaseV4 = h * HEAD_DIM_V4;
   let kvBaseV4 = (LAYER * CACHE_LEN * KV_HEADS + hKv) * HEAD_DIM_V4;
   let kvTokenStrideV4 = KV_HEADS * HEAD_DIM_V4;
-{% else %}
-{% if splitQueries %}
-  let qBaseV4 = ((b * Q_HEADS + h) * Q_SEQ + queryToken) * HEAD_DIM_V4;
 {% else %}
   let qBaseV4 = (b * Q_HEADS + h) * HEAD_DIM_V4;
-{% endif %}
   let kvBaseV4 = (b * KV_HEADS + hKv) * cacheSeq * HEAD_DIM_V4;
   let kvTokenStrideV4 = HEAD_DIM_V4;
 {% endif %}
@@ -437,27 +373,12 @@ fn main(
 {% endif %}
 {% if hasMask %}
     if (keyAllowed) {
-{% if splitQueries %}
-      let maskQuery = queryToken;
-{% else %}
       let maskQuery = 0u;
-{% endif %}
       let maskIndex = b * params.maskBatchStride + h * params.maskHeadStride + maskQuery * params.maskSeqStride + kj;
-{% 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.
@@ -473,7 +394,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
@@ -521,20 +441,12 @@ fn main(
   }
   // Emit un-normalized partials for (b, h, split): the merge pass divides.
-{% if splitQueries %}
-  let partialBase = (((b * Q_SEQ + queryToken) * Q_HEADS + h) * NUM_SPLITS + split) * HEAD_DIM_V4;
-{% else %}
   let partialBase = ((b * Q_HEADS + h) * NUM_SPLITS + split) * HEAD_DIM_V4;
-{% endif %}
   for (var d4: u32 = tid; d4 < HEAD_DIM_V4; d4 = d4 + WG) {
     partial_out[partialBase + d4] = running_out[d4];
   }
   if (tid == 0u) {
-{% if splitQueries %}
-    let mdBase = ((b * Q_SEQ + queryToken) * Q_HEADS + h) * NUM_SPLITS + split;
-{% else %}
     let mdBase = (b * Q_HEADS + h) * NUM_SPLITS + split;
-{% endif %}
     // (max, denom) travel together to the merge, so they share one buffer as an
     // interleaved vec2 rather than costing two bindings. Interleaved, not two
     // halves, so the index needs no region size — and the merge reads both

 {% if useSubgroups is not defined %}{% set useSubgroups = true %}{% 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 %}
+{% set cacheSeqlensIndex = "b" %}{% set layer = 0 %}
+{% set cacheLen = 0 %}
+{% set scale = "0.0" %}
 {% if hasRotary is not defined %}{% set hasRotary = false %}{% endif %}
+{% set qHiddenV4 = qHiddenV4 | default(0) %}
+{% set kvHiddenV4 = kvHiddenV4 | default(0) %}
+{% set hasBias = hasBias is defined and hasBias %}
 {% set splitKWorkgroupSize = workgroupSizeSpec if workgroupSizeSpec is defined else tunables.WORKGROUP_SIZE %}
 {% if useSubgroups %}
 enable subgroups;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
+// Split-K flash attention. Single-partition direct output skips the merge pass.
+// Otherwise this is the first of two passes. This geometry
 // handles decode and short-query, long-context prefill inputs.
 //
 // The non-split flash decode launches only `batch * numHeads` workgroups, each
 {% endif %}
 const WG: u32 = {{ splitKWorkgroupSize }}u;
 const NUM_SPLITS: u32 = {{ numSplits }}u;
 // FLT_MAX, not -inf, as the online (m, d) accumulator init: merges must keep
 // `m - m` finite so an empty lane / all--inf row contributes the exact
 // accumulator identity (m, d) = (-FLT_MAX, 0). Operator epilogues interpret
   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 %}
 {% if layout == "layer_cache" %}{% set ATTN_SCALE_OVERRIDE = "ATTN_SCALE" %}{% endif %}
+{% set ATTN_SCALE_DIM = "HEAD_DIM" %}fn scale_value() -> f32 {
 {% if ATTN_SCALE_OVERRIDE is defined %}
   return {{ ATTN_SCALE_OVERRIDE }};
 {% else %}
 {% endif %}
 }
 {% if quantizedCache %}
 {% macro emit_quant_scale4(kind, scaleBuffer) %}
 fn {{ kind }}scale4(d4: u32, hk: u32) -> vec4<f32> {
     {{ 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> {
   return vec4<f32>({{ buffer }}[indexV4]) * {{ kind }}scale4(d4, hk);
+}{% endmacro %}
 {{ emit_quant_load4("int8", "key", "key", "k_scale") }}
 {{ emit_quant_load4("int8", "value", "value", "v_scale") }}
 {% else %}
 // 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 %}
   // choose an intermediate width at pipeline execution time.
   if (sgSize == 0u || sgSize > WG || WG % sgSize != 0u) { return; }
 {% endif %}
   let split = wg.x;
   let h = wg.y;
   let b = wg.z;
+  if (h >= Q_HEADS || split >= NUM_SPLITS{% if layout == "layer_cache" %} || params.past_len >= CACHE_LEN{% endif %}) {
     return;
   }
   let tid = lid.x;
 {% if cacheSeqlens %}
   // Buffer-sharing caches retain their capacity in the physical BNSH stride;
   // seqlens_k supplies the active end independently for each batch.
+  let kvSeq = min(cacheSeq, u32(seqlens_k[{{ cacheSeqlensIndex }}]) + 1u);
 {% else %}
   let kvSeq = cacheSeq;
 {% endif %}
   // Query row (decode uses token zero; short-query prefill folds the token into wg.x).
 {% if layout == "bsh" %}
   let qBaseV4 = b * Q_HIDDEN_V4 + h * HEAD_DIM_V4;
   let kvBaseV4 = b * kvSeq * KV_HIDDEN_V4 + hKv * HEAD_DIM_V4;
   let kvTokenStrideV4 = KV_HIDDEN_V4;
 {% elif layout == "layer_cache" %}
   let qBaseV4 = h * HEAD_DIM_V4;
   let kvBaseV4 = (LAYER * CACHE_LEN * KV_HEADS + hKv) * HEAD_DIM_V4;
   let kvTokenStrideV4 = KV_HEADS * HEAD_DIM_V4;
 {% else %}
   let qBaseV4 = (b * Q_HEADS + h) * HEAD_DIM_V4;
   let kvBaseV4 = (b * KV_HEADS + hKv) * cacheSeq * HEAD_DIM_V4;
   let kvTokenStrideV4 = HEAD_DIM_V4;
 {% endif %}
 {% endif %}
 {% if hasMask %}
     if (keyAllowed) {
       let maskQuery = 0u;
       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
   }
   // Emit un-normalized partials for (b, h, split): the merge pass divides.
   let partialBase = ((b * Q_HEADS + h) * NUM_SPLITS + split) * HEAD_DIM_V4;
   for (var d4: u32 = tid; d4 < HEAD_DIM_V4; d4 = d4 + WG) {
     partial_out[partialBase + d4] = running_out[d4];
   }
   if (tid == 0u) {
     let mdBase = (b * Q_HEADS + h) * NUM_SPLITS + split;
     // (max, denom) travel together to the merge, so they share one buffer as an
     // interleaved vec2 rather than costing two bindings. Interleaved, not two
     // halves, so the index needs no region size — and the merge reads both

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

@@ -1,6 +1,4 @@
-{% if combineSubgroups %}
 enable subgroups;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 // Flash-style tiled online-softmax attention for vec4-aligned head dimensions.
@@ -28,8 +26,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 +48,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,8 +62,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
 // publishes the subgroup pair in the slot at its OWN invocation index and sets
@@ -115,96 +112,28 @@ fn combine_partials(m: f32, d: f32, lidx: u32, sgSize: u32) -> vec2<f32> {
   workgroupBarrier();
   return combinedMD;
 }
-{% else %}
-{% set mdStreamed = mdStreams is defined %}
-{% set mdStreams = mdStreams if mdStreams is defined else 1 %}
-{% set mdExtent = "WG" if mdStreams == 1 else "WG * " ~ mdStreams ~ "u" %}
-var<workgroup> partialM: array<f32, {{ mdExtent }}>;
-var<workgroup> partialD: array<f32, {{ mdExtent }}>;
-{% if mdStreamed %}
-// In-place fold of {{ mdStreams }} streams. 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;
-  partialD[lidx] = d;
-  workgroupBarrier();
-  var stride = WG / 2u;
-  loop {
-    if (stride == 0u) {
-      break;
-    }
-    if (lidx < stride) {
-      let m1 = partialM[lidx];
-      let d1 = partialD[lidx];
-      let m2 = partialM[lidx + stride];
-      let d2 = partialD[lidx + stride];
-      let mNew = max(m1, m2);
-      partialD[lidx] = d1 * exp_shift(m1, mNew) + d2 * exp_shift(m2, mNew);
-      partialM[lidx] = mNew;
-    }
-    workgroupBarrier();
-    stride = stride / 2u;
-  }
-  let merged = vec2<f32>(partialM[0], partialD[0]);
-  // Trailing barrier so back-to-back calls cannot race a next call's partial
-  // stores against this call's reads of slot 0.
-  workgroupBarrier();
-  return merged;
-}
-{% endif %}
-{% endif %}
 // 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 %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(
   @builtin(workgroup_id) wg: vec3<u32>,
-  @builtin(local_invocation_id) lid: vec3<u32>{% if combineSubgroups %},
-  @builtin(subgroup_size) sgSize: u32{% endif %}
-) {
   let qi = wg.x;
   let h = wg.y;
   let b = wg.z;
@@ -243,7 +172,11 @@ fn main(
   // Causal upper bound: query qi attends only keys 0..qi, so stop after the tile
   // containing qi and skip the unattended tail. Non-causal keeps the full kvSeq
   // sweep.
   var keyBoundV = params.kvSeq;
   var keyFloor: u32 = 0u;
 {% if hasWindow %}
   // Sliding window: query qi sits at absolute position p = kvSeq - qSeq
@@ -269,7 +202,7 @@ fn main(
     var score = -FLT_MAX;
     var m = -FLT_MAX;
     var dPart = 0.0;
-    var keyAllowed = kj < keyBound{% if hasWindow %} && kj >= keyFloor{% endif %};
     if (keyAllowed) {
       let kRowV4 = kvBaseV4 + kj * kvTokenStrideV4;
 {% if hasMask %}
@@ -277,17 +210,6 @@ fn main(
       // [q, k] masks set batch/head strides to 0).
       let maskIndex = b * params.maskBatchStride + h * params.maskHeadStride + qi * params.maskSeqStride + kj;
 {% endif %}
-{% if hasMask and maskIsBool %}
-      if (attn_mask[maskIndex] == 0u) {
-        keyAllowed = false;
-      } else {
-        var acc: f32 = 0.0;
-        for (var d4: u32 = 0u; d4 < HEAD_DIM_V4; d4 = d4 + 1u) {
-          acc = acc + dot(q_shared[d4], vec4<f32>(key[kRowV4 + d4]));
-        }
-        score = acc * scale;
-      }
-{% else %}
       var acc: f32 = 0.0;
       for (var d4: u32 = 0u; d4 < HEAD_DIM_V4; d4 = d4 + 1u) {
         acc = acc + dot(q_shared[d4], vec4<f32>(key[kRowV4 + d4]));
@@ -295,12 +217,11 @@ fn main(
       score = acc * scale;
 {% if hasMask %}
       score = score + f32(attn_mask[maskIndex]);
-{% endif %}
 {% endif %}
       m = score;
       dPart = select(0.0, 1.0, keyAllowed);
     }
-    let tile = combine_partials(m, dPart, tid{% if combineSubgroups %}, sgSize{% endif %});
     // Online merge of the tile into the running state (softmax-online rule).
 // Merge one key tile's online-softmax (maximum, denominator) partial into the
@@ -317,7 +238,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);

 enable subgroups;
 {{ env.wgsl.resourceDeclarations }}
 // Flash-style tiled online-softmax attention for vec4-aligned head dimensions.
 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.
 // Cross-subgroup merge that assumes nothing about which invocations share a
 // subgroup or how many subgroups there are: each subgroup's elected lane
 // publishes the subgroup pair in the slot at its OWN invocation index and sets
   workgroupBarrier();
   return combinedMD;
 }
 // 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 %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(
   @builtin(workgroup_id) wg: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>,
+  @builtin(subgroup_size) sgSize: u32) {
   let qi = wg.x;
   let h = wg.y;
   let b = wg.z;
   // Causal upper bound: query qi attends only keys 0..qi, so stop after the tile
   // containing qi and skip the unattended tail. Non-causal keeps the full kvSeq
   // sweep.
+{% if hasCausal %}
+  var keyBoundV = select(params.kvSeq, min(params.kvSeq, qi + 1u), params.isCausal != 0u);
+{% else %}
   var keyBoundV = params.kvSeq;
+{% endif %}
   var keyFloor: u32 = 0u;
 {% if hasWindow %}
   // Sliding window: query qi sits at absolute position p = kvSeq - qSeq
     var score = -FLT_MAX;
     var m = -FLT_MAX;
     var dPart = 0.0;
+    var keyAllowed = kj < keyBound{% if hasCausal %} && (params.isCausal == 0u || kj <= qi){% endif %}{% if hasWindow %} && kj >= keyFloor{% endif %};
     if (keyAllowed) {
       let kRowV4 = kvBaseV4 + kj * kvTokenStrideV4;
 {% if hasMask %}
       // [q, k] masks set batch/head strides to 0).
       let maskIndex = b * params.maskBatchStride + h * params.maskHeadStride + qi * params.maskSeqStride + kj;
 {% endif %}
       var acc: f32 = 0.0;
       for (var d4: u32 = 0u; d4 < HEAD_DIM_V4; d4 = d4 + 1u) {
         acc = acc + dot(q_shared[d4], vec4<f32>(key[kRowV4 + d4]));
       score = acc * scale;
 {% if hasMask %}
       score = score + f32(attn_mask[maskIndex]);
 {% endif %}
       m = score;
       dPart = select(0.0, 1.0, keyAllowed);
     }
+    let tile = combine_partials(m, dPart, tid, sgSize);
     // Online merge of the tile into the running state (softmax-online rule).
 // Merge one key tile's online-softmax (maximum, denominator) partial into the
     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,22 +1,18 @@
-{% 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 %}
 {% 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
@@ -24,7 +20,14 @@
 // 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;
@@ -40,11 +43,30 @@ enable subgroups;
 {% 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 %}
 {% set ROPE_LANE_XOR = ((LPQ / 2) | int) %}
 {% set QL = qLayout if qLayout is defined else layout %}
 {% set KL = kvLayout if kvLayout is defined else layout %}
@@ -69,8 +91,7 @@ const HEAD_DIM: u32 = {{ headDim }}u;
 const HEAD_DIM_V4: u32 = {{ headDimV4 }}u;
 {% if QL != "bhsd" %}
 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 %}{% endif %}
 {% if KL != "bhsd" %}
 const KV_HIDDEN_V4: u32 = {{ kvHiddenV4 }}u;
 {% endif %}
@@ -91,16 +112,7 @@ 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)>;
@@ -121,16 +133,10 @@ 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
@@ -147,14 +153,13 @@ fn {{ kind }}scale4(d4: u32, hk: u32) -> vec4<f32> {
     {{ 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];
@@ -165,14 +170,21 @@ fn load_{{ kind }}4(indexV4: u32, d4: u32, hk: u32) -> vec4<f32> {
   );
   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 %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(
   @builtin(workgroup_id) wg: vec3<u32>,
@@ -211,6 +223,9 @@ fn main(
 {% endif %}
 {% for c in range(SLICE_COUNT) %}
   var {{ qn("qr" ~ c, qi) }} = vec4<f32>({{ QUERY }}[{{ qn("qBase4", qi) }} + {{ c }}u]);
   var {{ qn("o" ~ c, qi) }} = vec4<f32>(0.0);
 {% endfor %}
 {% endfor %}
@@ -390,9 +405,13 @@ fn main(
     // TILE_K is a small shader constant; the loop updates the named q/o slices in place.
 {% if not useSubgroups and batchNoSgReduction %}
     // First publish every key's partial dot without intervening barriers.
     var s: array<f32, TILE_K>;
     for (var kk: u32 = 0u; kk < TILE_K; kk = kk + 1u) {
       s[kk] = NEG_INF;
       var part: f32 = 0.0;
       let kb = kk * HEAD_DIM_V4 + lane8 * SLICE;
 {% for c in range(SLICE_COUNT) %}
@@ -422,26 +441,41 @@ fn main(
     workgroupBarrier();
     // The barrier at the start of the next K tile protects this scratch before
     // it is overwritten.
     for (var kk: u32 = 0u; kk < TILE_K; kk = kk + 1u) {
       let kj = kStart + kk;
       let part = red[kk * WG + tid - lane8];
       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 %}
       }
     }
 {% elif QPL > 1 %}
 {% for qi in range(QPL) %}
     var {{ qn("s", qi) }}: array<f32, TILE_K>;
@@ -481,21 +515,7 @@ fn main(
           // 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 %}
@@ -541,21 +561,7 @@ fn main(
         // 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 %}
@@ -563,30 +569,36 @@ 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.
     // Interior tiles use the unguarded FMA chain. Mask exclusion applies to
     // materialized V rows and does not require this range guard.
-{% if QPL > 1 %}
     let tileInterior = {% for qi in range(QPL) %}(kStart >= {{ qn("minKj", qi) }} && kStart + TILE_K <= {{ qn("maxKj", qi) }}){% if not loop.last %} && {% endif %}{% endfor %};
 {% for c in range(SLICE_COUNT) %}
     {
@@ -635,7 +647,7 @@ fn main(
     kStart = kStart + TILE_K;
   }
-{% macro attention_value(c, qi) %}{{ qn("o" ~ c, qi) }} * inv{% endmacro %}
 {% for qi in range(QPL) %}
   if ({{ qn("qValid", qi) }}) {
 {% if QL == "bhsd" %}

+{% 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 %}
+{% set SHARE_PROB = shareNoSgProb is defined and shareNoSgProb %}
 {% if hasMask is not defined %}{% set hasMask = false %}{% endif %}
+{% set maskIsKeyKeep = false %}{% set maskIsBool = false %}{% 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
 // 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 useSeqlens %}
+{% set KVA = "kvActive" %}
+{% elif hasKeyLimit %}
+{% set KVA = "select(params.kvSeq, min(params.keyLimit, params.kvSeq), params.keyLimit > 0u)" %}
+{% else %}
+{% set KVA = KVSEQ %}
+{% endif %}
 {% 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;
 {% 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_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 %}
 {% set ROPE_LANE_XOR = ((LPQ / 2) | int) %}
 {% set QL = qLayout if qLayout is defined else layout %}
 {% set KL = kvLayout if kvLayout is defined else layout %}
 const HEAD_DIM_V4: u32 = {{ headDimV4 }}u;
 {% if QL != "bhsd" %}
 const Q_HIDDEN_V4: u32 = {{ qHiddenV4 }}u;   // bsh token-major hidden stride
+{% endif %}
 {% if KL != "bhsd" %}
 const KV_HIDDEN_V4: u32 = {{ kvHiddenV4 }}u;
 {% endif %}
 {% 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)>;
 {% 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 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
     {{ 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 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 %}
+{% 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(WG, 1, 1)
 fn main(
   @builtin(workgroup_id) wg: vec3<u32>,
 {% endif %}
 {% for c in range(SLICE_COUNT) %}
   var {{ qn("qr" ~ c, qi) }} = vec4<f32>({{ QUERY }}[{{ qn("qBase4", qi) }} + {{ c }}u]);
+{% if hasBias %}
+  {{ qn("qr" ~ c, qi) }} = {{ qn("qr" ~ c, qi) }} + load_bias4(h * HEAD_DIM, lane8 * SLICE + {{ c }}u);
+{% endif %}
   var {{ qn("o" ~ c, qi) }} = vec4<f32>(0.0);
 {% endfor %}
 {% endfor %}
     // TILE_K is a small shader constant; the loop updates the named q/o slices in place.
 {% if not useSubgroups and batchNoSgReduction %}
     // First publish every key's partial dot without intervening barriers.
+{% if not SHARE_PROB %}
     var s: array<f32, TILE_K>;
+{% endif %}
     for (var kk: u32 = 0u; kk < TILE_K; kk = kk + 1u) {
+{% if not SHARE_PROB %}
       s[kk] = NEG_INF;
+{% endif %}
       var part: f32 = 0.0;
       let kb = kk * HEAD_DIM_V4 + lane8 * SLICE;
 {% for c in range(SLICE_COUNT) %}
     workgroupBarrier();
     // The barrier at the start of the next K tile protects this scratch before
     // it is overwritten.
+{% if SHARE_PROB %}
+    // One lane computes the query softmax, reusing the reduced-score slots for
+    // probabilities. The discarded partial lanes hold the three online states.
+    if (lane8 == 0u) {
+      var s: array<f32, TILE_K>;
+{% endif %}
     for (var kk: u32 = 0u; kk < TILE_K; kk = kk + 1u) {
       let kj = kStart + kk;
       let part = red[kk * WG + tid - lane8];
+{% if SHARE_PROB %}
+      s[kk] = NEG_INF;
+{% endif %}
       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 %}
       }
     }
+{% if SHARE_PROB %}
+{{ emit_tile_softmax() }}
+      for (var kk: u32 = 0u; kk < TILE_K; kk = kk + 1u) {
+        red[kk * WG + cbase] = s[kk];
+      }
+      red[cbase + 1u] = corr;
+      red[cbase + 2u] = m;
+      red[cbase + 3u] = l;
+    }
+    workgroupBarrier();
+    let corr = red[cbase + 1u];
+    m = red[cbase + 2u];
+    l = red[cbase + 3u];
+{% endif %}
 {% elif QPL > 1 %}
 {% for qi in range(QPL) %}
     var {{ qn("s", qi) }}: array<f32, TILE_K>;
           // 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 %}
           {{ qn("s", qi) }}[kk] = {{ score_expr("sc") }} + {{ MASK_ADDITIVE }};
   {% else %}
           {{ qn("s", qi) }}[kk] = {{ score_expr("sc") }};
   {% endif %}
         // 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 %}
         s[kk] = {{ score_expr("part") }} + {{ MASK_ADDITIVE }};
 {% else %}
         s[kk] = {{ score_expr("part") }};
 {% endif %}
     }
 {% endif %}
+{% if not SHARE_PROB %}
+{{ emit_tile_softmax() }}
+{% endif %}
     // 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.
     // Interior tiles use the unguarded FMA chain. Mask exclusion applies to
     // materialized V rows and does not require this range guard.
+{% if SHARE_PROB %}
+{% macro shared_pv(masked) %}
+        for (var kk: u32 = 0u; kk < TILE_K; kk = kk + 1u) {
+          let weight = red[kk * WG + cbase];
+          let vb = kk * HEAD_DIM_V4 + lane8 * SLICE;
+{% for c in range(SLICE_COUNT) %}
+          o{{ c }} = o{{ c }} + weight * {% if masked %}select(vec4<f32>(), {% endif %}vec4<f32>(v_tile[vb + {{ c }}u]){% if masked %}, weight != 0.0){% endif %};
+{% endfor %}
+        }
+{% endmacro %}
+    // Consume each shared probability once across all output slices. This
+    // retains the key order of every output accumulator without a private tile.
+{% for c in range(SLICE_COUNT) %}
+    o{{ c }} = o{{ c }} * corr;
+{% endfor %}
+    let tileInterior = kStart >= minKj && kStart + TILE_K <= maxKj;
+    if (tileInterior) {
+{{ shared_pv(false) }}
+    } else {
+{{ shared_pv(true) }}
+    }
+{% elif QPL > 1 %}
     let tileInterior = {% for qi in range(QPL) %}(kStart >= {{ qn("minKj", qi) }} && kStart + TILE_K <= {{ qn("maxKj", qi) }}){% if not loop.last %} && {% endif %}{% endfor %};
 {% for c in range(SLICE_COUNT) %}
     {
     kStart = kStart + TILE_K;
   }
+{% macro attention_value(c, qi) %}{{ qn("o" ~ c, qi) }} * inv{% if hasBias %} + load_bias4(2u * {{ Q_HIDDEN }}u + h * HEAD_DIM, lane8 * SLICE + {{ c }}u){% endif %}{% endmacro %}
 {% for qi in range(QPL) %}
   if ({{ qn("qValid", qi) }}) {
 {% if QL == "bhsd" %}

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,10 +11,15 @@
 {% 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 %}
 {% if USE_SUBGROUPS %}
 enable subgroups;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
 const HEAD_DIM_V4: u32 = {{ headDimV4 }}u;
@@ -43,8 +49,7 @@ fn scale_value() -> f32 {
 }
-@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 %},
@@ -85,8 +90,19 @@ fn main(
   // Causal key ceiling. The key-loop bound (kvEnd) uses the workgroup's LAST query
   // so every lane shares a uniform trip count (subgroup ops stay reconverged);
   // each lane masks its own keys past myMaxKj to NEG_INF.
   let kvEnd = params.kvSeq;
   let myMaxKj = params.kvSeq;
   let kvBase = b * params.kvSeq * KV_HIDDEN_V4 + hKv * HEAD_DIM_V4;
   let kvKeyStride = KV_HIDDEN_V4;
@@ -174,6 +190,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 %}
@@ -190,21 +213,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 %}
 {% 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;
 }
+@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 %},
   // Causal key ceiling. The key-loop bound (kvEnd) uses the workgroup's LAST query
   // 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 = fixedCausal is defined and fixedCausal %}
+{% 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;
+{% endif %}
   let kvBase = b * params.kvSeq * KV_HIDDEN_V4 + hKv * HEAD_DIM_V4;
   let kvKeyStride = KV_HIDDEN_V4;
 {% 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-rowstats-combine-f32.wgsl.jinja CHANGED Viewed

@@ -12,6 +12,12 @@
 // `maxOnly` means the producer published only a row max per slot because
 // 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.
 {{ env.wgsl.resourceDeclarations }}
 const SLOTS: u32 = {{ statSlots }}u;
@@ -37,6 +43,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));
 }
@@ -45,8 +52,7 @@ fn exp_shift(value: f32, maxValue: f32) -> f32 {
 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

 // `maxOnly` means the producer published only a row max per slot because
 // 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.
+{% 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);
 }
 fn exp_shift(value: f32, maxValue: f32) -> f32 {
   return exp(shifted_value(value, maxValue));
 }
 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-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,18 +7,16 @@ enable subgroup_size_control;
 enable chromium_experimental_subgroup_matrix;
 diagnostic(off, chromium.subgroup_matrix_uniformity);
 {{ env.wgsl.resourceDeclarations }}
-{% set layout = layout | default("bsh") %}
-{% set headMajor = layout == "bhsd" %}
 {% set kvHeadMajor = (kvLayout | default(layout)) == "bhsd" %}
-{% set CAUSAL_UPPER_LEFT = causalUpperLeft is defined and causalUpperLeft %}
 {% set CAUSAL = (causalRightAlign is defined and causalRightAlign) or CAUSAL_UPPER_LEFT %}
-{% macro q_index(row, d) %}{% if headMajor %}((b * HEADS + h) * params.qSeq + {{ row }}) * HEAD_DIM + {{ d }}{% else %}(b * params.qSeq + {{ row }}) * HIDDEN + h * HEAD_DIM + {{ d }}{% endif %}{% endmacro %}
 {% 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 SCORE_BIAS = scorePhase and scoreBias is defined and scoreBias %}
 {% set SCORE_WINDOW = CAUSAL and scoreWindow is defined and scoreWindow %}
@@ -28,45 +24,38 @@ diagnostic(off, chromium.subgroup_matrix_uniformity);
 {% 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 %}
 {% set WORKGROUP_THREADS = SUBGROUP_COUNT * 32 %}
 {% 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 %}
-{% macro q_tile_value(index) %}{% if hasBias %}(query[{{ index }}] + bias[h * HEAD_DIM + k]){% else %}query[{{ index }}]{% endif %}{% endmacro %}
 const HEADS: u32 = {{ qNumHeads }}u;
 const KV_HEADS: u32 = {{ kvNumHeads }}u;
 const HEAD_DIM: u32 = {{ headDim }}u;
-{% if layout != "bhsd" or hasBias %}
 const HIDDEN: u32 = {{ qHidden }}u;
-{% endif %}
 {% if hasBias %}
 /* Packed [Q; K; V] bias. Q folds into the query tile before the GEMM. K is
  * omitted: expanding (q + bq).(k + bk) leaves a (q + bq).bk term that is
@@ -108,14 +97,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 +116,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 "" }}
@@ -177,7 +160,7 @@ fn main(
       let stat_row = m_base + idx / {{ SUBGROUP_COLS }}u;
       if (stat_row < params.qSeq) {
         let slot = wg.x * {{ SUBGROUP_COLS }}u + idx % {{ SUBGROUP_COLS }}u;
-        let out_index = (((b * HEADS + h) * STAT_SLOTS + slot) * params.qSeq + stat_row) * 2u;
         scorePartials[out_index] = -FLT_MAX;
         scorePartials[out_index + 1u] = 0.0;
       }
@@ -211,30 +194,18 @@ fn main(
 {% endif %}
 {% 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 = {% if SCORE_WINDOW and not scorePhase %}inner_start{% else %}0u{% endif %}; k_base < {% if CAUSAL and not scorePhase %}inner_bound{% else %}inner{% endif %}; 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 (
@@ -244,7 +215,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>(
@@ -294,7 +265,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 %}
@@ -309,20 +280,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 %}
@@ -341,13 +312,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 %}
@@ -362,12 +333,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;
@@ -410,7 +381,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>(
@@ -467,7 +438,9 @@ fn main(
           + row_in_block * 8u + col_in_block + pair
       ];
 {% if scorePhase %}
       let scale = {{ attentionScaleExpression }};
 {% if CAUSAL %}
       var scored = result * scale;
 {% if SCORE_BIAS %}
@@ -485,10 +458,10 @@ fn main(
       if (i32(col) + i32(params.windowSize) <= kv_causal_off + i32(row)) { scored = -FLT_MAX; }
 {% endif %}
 {% else %}
-      let scored = result * 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
@@ -502,7 +475,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 %}
     }
   }
@@ -530,7 +503,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 kvHeadMajor = (kvLayout | default(layout)) == "bhsd" %}
+{% set CAUSAL_UPPER_LEFT = false %}
 {% set CAUSAL = (causalRightAlign is defined and causalRightAlign) or CAUSAL_UPPER_LEFT %}
+{% macro q_index(row, d) %}(b * params.qSeq + {{ row }}) * HIDDEN + h * HEAD_DIM + {{ d }}{% endmacro %}
 {% 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 = "HIDDEN" %}
 {% set KV_ROW_STRIDE = "HEAD_DIM" if kvHeadMajor else "KV_HIDDEN" %}
+{% set SCRATCH_HEAD = "(b * HEADS + h)" %}
 {% set scorePhase = phase == "score" %}
 {% set SCORE_BIAS = scorePhase and scoreBias is defined and scoreBias %}
 {% set SCORE_WINDOW = CAUSAL and scoreWindow is defined and scoreWindow %}
 {% 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 %}
 {% set WORKGROUP_THREADS = SUBGROUP_COUNT * 32 %}
 {% if hasBias is not defined %}{% set hasBias = false %}{% endif %}
 {% set EMIT_ROW_STATS = emitRowStats is defined and emitRowStats %}
+{% set DIRECT_SCORE_STORE = false %}
+{% set DIRECT_APPLY_STORE = false %}
 {% 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;
 const KV_HEADS: u32 = {{ kvNumHeads }}u;
 const HEAD_DIM: u32 = {{ headDim }}u;
 const HIDDEN: u32 = {{ qHidden }}u;
 {% if hasBias %}
 /* Packed [Q; K; V] bias. Q folds into the query tile before the GEMM. K is
  * omitted: expanding (q + bq).(k + bk) leaves a (q + bq).bk term that is
 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 stat_row = m_base + idx / {{ SUBGROUP_COLS }}u;
       if (stat_row < params.qSeq) {
         let slot = wg.x * {{ SUBGROUP_COLS }}u + idx % {{ SUBGROUP_COLS }}u;
+        let out_index = (({{ SCRATCH_HEAD }} * STAT_SLOTS + slot) * params.qSeq + stat_row) * 2u;
         scorePartials[out_index] = -FLT_MAX;
         scorePartials[out_index + 1u] = 0.0;
       }
 {% endif %}
 {% 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 = {% if SCORE_WINDOW and not scorePhase %}inner_start{% else %}0u{% endif %}; k_base < {% if CAUSAL and not scorePhase %}inner_bound{% else %}inner{% endif %}; 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>(
           + row_in_block * 8u + col_in_block + pair
       ];
 {% if scorePhase %}
+{% if not SCALE_IN_Q %}
       let scale = {{ attentionScaleExpression }};
+{% endif %}
 {% if CAUSAL %}
       var scored = result * scale;
 {% if SCORE_BIAS %}
       if (i32(col) + i32(params.windowSize) <= kv_causal_off + i32(row)) { scored = -FLT_MAX; }
 {% endif %}
 {% else %}
+      let scored = result{% if not SCALE_IN_Q %} * scale{% endif %};
 {% 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-online-scalar.wgsl.jinja CHANGED Viewed

@@ -1,14 +1,17 @@
 {{ 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.
 {% if hasMask %}
 // Additive per-score attention bias with broadcast strides: a stride of 0
 // collapses that axis (batch and/or head broadcast).
@@ -18,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
@@ -39,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;
@@ -55,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>,
@@ -110,6 +124,9 @@ fn main(
   var maxKj = params.kvSeq;
   var minKj: u32 = 0u;
 {% if hasWindow %}
   // local_window_size: the query at relative index
   // query_token sits at absolute position p = kvSeq - qSeq + query_token, so it
@@ -124,21 +141,26 @@ 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]);
-      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);
@@ -153,6 +175,9 @@ fn main(
     let probability_numerator = exp(score - running_max);
     for (var d: u32 = tid; d < HEAD_DIM; d = d + WG) {
       var v_value = f32(value[kRow + d]);
       running_out[d] = running_out[d] * previous_scale + probability_numerator * v_value;
     }
     workgroupBarrier();

 {{ 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.
+{% if hasBias %}
+// Packed [Q; K; V] bias rows are applied during the serial key walk.
+{% endif %}
 {% if hasMask %}
 // Additive per-score attention bias with broadcast strides: a stride of 0
 // collapses that axis (batch and/or head broadcast).
 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>,
   var maxKj = params.kvSeq;
   var minKj: u32 = 0u;
+{% if hasCausal %}
+  maxKj = min(maxKj, select(params.kvSeq, query_token + 1u, params.isCausal != 0u));
+{% endif %}
 {% if hasWindow %}
   // local_window_size: the query at relative index
   // query_token sits at absolute position p = kvSeq - qSeq + query_token, so it
   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]);
+{% if hasBias %}
+      let channel = h * HEAD_DIM + 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);
     let probability_numerator = exp(score - running_max);
     for (var d: u32 = tid; d < HEAD_DIM; d = d + WG) {
       var v_value = f32(value[kRow + d]);
+{% if hasBias %}
+      v_value = v_value + f32(bias[2u * Q_HIDDEN + h * HEAD_DIM + d]);
+{% endif %}
       running_out[d] = running_out[d] * previous_scale + probability_numerator * v_value;
     }
     workgroupBarrier();

build/webgpu/bench.json CHANGED Viewed

@@ -710,7 +710,9 @@
     {
       "name": "quant-int8-prefill-h32kv8-d128-s512-pathology",
       "preset": "stress",
-      "provenance": { "notes": "A production-sized prefill exercises the int8 KV-cache attention contract." },
       "attrs": {
         "num_heads": 32,
         "kv_num_heads": 8,
@@ -765,7 +767,9 @@
     {
       "name": "quant-int4-prefill-h32kv8-d128-s512-pathology",
       "preset": "stress",
-      "provenance": { "notes": "A production-sized prefill exercises the packed int4 KV-cache attention contract." },
       "attrs": {
         "num_heads": 32,
         "kv_num_heads": 8,
@@ -1043,7 +1047,7 @@
       "name": "headsink-prefill-h32kv8-d128-s512-generic-pathology",
       "preset": "stress",
       "provenance": {
-        "notes": "A Llama-sized GQA prefill exercises smooth-softmax head sinks on a production attention shape."
       },
       "attrs": { "num_heads": 32, "kv_num_heads": 8 },
       "inputs": {
@@ -1086,7 +1090,7 @@
       "name": "bias-headsink-prefill-h32kv8-d128-s512-generic-pathology",
       "preset": "stress",
       "provenance": {
-        "notes": "A production prefill combining additive bias with smooth-softmax head sinks exercises the generic attention route."
       },
       "attrs": { "num_heads": 32, "kv_num_heads": 8 },
       "inputs": {
@@ -1133,7 +1137,9 @@
     {
       "name": "softcap-prefill-h32kv8-d128-s512-generic-pathology",
       "preset": "stress",
-      "provenance": { "notes": "A production prefill exercises model-used score soft-capping on the past-KV contract." },
       "attrs": { "num_heads": 32, "kv_num_heads": 8, "softcap": 30 },
       "inputs": {
         "queryT": {
@@ -1460,6 +1466,74 @@
         ]
       }
     },
     {
       "name": "sharedkv-chunk-append-32h8kv-d128-cap1024-q128",
       "preset": "smoke",
@@ -1551,7 +1625,97 @@
       },
       "provenance": {
         "source": "register-geometry gate asymmetry",
-        "notes": "Cached f32 prefill at headDim 256, above the shared-memory cluster's register-geometry boundary. This guards consistent route admission between no-past and cached prefill families."
       }
     },
     {
@@ -1576,6 +1740,439 @@
           { "type": "gflops", "value": "4 * args.batch * args.qSeq * args.kvSeq * args.heads * args.headDim" }
         ]
       }
     }
   ]
 }

     {
       "name": "quant-int8-prefill-h32kv8-d128-s512-pathology",
       "preset": "stress",
+      "provenance": {
+        "notes": "Measures GroupQueryAttention over a 512-token Llama-sized prefill (32 query heads grouped 4:1 over 8 key/value heads, head width 128) with an int8 quantized key/value cache."
+      },
       "attrs": {
         "num_heads": 32,
         "kv_num_heads": 8,
     {
       "name": "quant-int4-prefill-h32kv8-d128-s512-pathology",
       "preset": "stress",
+      "provenance": {
+        "notes": "Measures GroupQueryAttention over a 512-token Llama-sized prefill (32 query heads grouped 4:1 over 8 key/value heads, head width 128) with a packed int4 quantized key/value cache."
+      },
       "attrs": {
         "num_heads": 32,
         "kv_num_heads": 8,
       "name": "headsink-prefill-h32kv8-d128-s512-generic-pathology",
       "preset": "stress",
       "provenance": {
+        "notes": "Measures GroupQueryAttention over a 512-token Llama-sized prefill (32 query heads grouped 4:1 over 8 key/value heads, head width 128) with per-head attention sink logits."
       },
       "attrs": { "num_heads": 32, "kv_num_heads": 8 },
       "inputs": {
       "name": "bias-headsink-prefill-h32kv8-d128-s512-generic-pathology",
       "preset": "stress",
       "provenance": {
+        "notes": "Measures GroupQueryAttention over a 512-token Llama-sized prefill (32 query heads grouped 4:1 over 8 key/value heads, head width 128) with an additive attention bias and per-head attention sink logits."
       },
       "attrs": { "num_heads": 32, "kv_num_heads": 8 },
       "inputs": {
     {
       "name": "softcap-prefill-h32kv8-d128-s512-generic-pathology",
       "preset": "stress",
+      "provenance": {
+        "notes": "Measures GroupQueryAttention over a 512-token Llama-sized prefill (32 query heads grouped 4:1 over 8 key/value heads, head width 128) with a logit softcap of 30."
+      },
       "attrs": { "num_heads": 32, "kv_num_heads": 8, "softcap": 30 },
       "inputs": {
         "queryT": {
         ]
       }
     },
+    {
+      "name": "window-headsink-rotary-decode-gptoss-64h8kv-d64-cap384-w128",
+      "preset": "smoke",
+      "attrs": {
+        "num_heads": 64,
+        "kv_num_heads": 8,
+        "do_rotary": 1,
+        "sliding_window_cache": 1,
+        "local_window_size": 128
+      },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 4096],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 512],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.031, "cosStep": 0.007 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 512],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.037, "cosStep": 0.005 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 8, 384, 64],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.017, "cosStep": 0.019 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 8, 384, 64],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [4095] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [4096] } },
+        "cosCacheT": {
+          "dtype": "float32",
+          "shape": [4096, 32],
+          "data": { "kind": "rotaryCos", "thetaStart": 0.0, "thetaStep": 0.017 }
+        },
+        "sinCacheT": {
+          "dtype": "float32",
+          "shape": [4096, 32],
+          "data": { "kind": "rotarySin", "thetaStart": 0.0, "thetaStep": 0.017 }
+        },
+        "headSinkT": {
+          "dtype": "float32",
+          "shape": [64],
+          "data": { "kind": "fillFloat32", "scale": 0.5, "sinStep": 0.041, "cosStep": 0.011 }
+        }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 1, 4096] },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 8, 384, 64] },
+        "presentValueT": { "dtype": "float32", "shape": [1, 8, 384, 64] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * 2 * numel(shapes.pastKeyT) * 4 + numel(shapes.queryT) * 4 + numel(shapes.outputT) * 4"
+          }
+        ]
+      }
+    },
     {
       "name": "sharedkv-chunk-append-32h8kv-d128-cap1024-q128",
       "preset": "smoke",
       },
       "provenance": {
         "source": "register-geometry gate asymmetry",
+        "notes": "Measures cached float32 prefill that appends 128 tokens after 896 cached ones, filling a 1024-slot key/value buffer, with 8 query heads over 2 key/value heads of width 256."
+      }
+    },
+    {
+      "name": "bidirectional-share-append-32h8kv-d128-cap1024-q128",
+      "preset": "smoke",
+      "vars": { "batch": 1, "qSeq": 128, "kvSeq": 1024, "heads": 32, "headDim": 128 },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "causal": 0 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 128, 4096],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 128, 1024],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.031, "cosStep": 0.007 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 128, 1024],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.037, "cosStep": 0.005 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 8, 1024, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.017, "cosStep": 0.019 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 8, 1024, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1023] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1024] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 128, 4096] },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 8, 1024, 128] },
+        "presentValueT": { "dtype": "float32", "shape": [1, 8, 1024, 128] }
+      },
+      "bench": {
+        "metrics": [
+          { "type": "gflops", "value": "4 * args.batch * args.qSeq * args.kvSeq * args.heads * args.headDim" }
+        ]
+      }
+    },
+    {
+      "name": "bidirectional-share-append-f16-32h8kv-d128-cap512-q64",
+      "preset": "smoke",
+      "vars": { "batch": 1, "qSeq": 64, "kvSeq": 512, "heads": 32, "headDim": 128 },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "causal": 0 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [1, 64, 4096],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [1, 64, 1024],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.031, "cosStep": 0.007 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [1, 64, 1024],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.037, "cosStep": 0.005 }
+        },
+        "pastKeyT": {
+          "dtype": "float16",
+          "shape": [1, 8, 512, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.017, "cosStep": 0.019 }
+        },
+        "pastValueT": {
+          "dtype": "float16",
+          "shape": [1, 8, 512, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [512] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float16", "shape": [1, 64, 4096] },
+        "presentKeyT": { "dtype": "float16", "shape": [1, 8, 512, 128] },
+        "presentValueT": { "dtype": "float16", "shape": [1, 8, 512, 128] }
+      },
+      "bench": {
+        "metrics": [
+          { "type": "gflops", "value": "4 * args.batch * args.qSeq * args.kvSeq * args.heads * args.headDim" }
+        ]
       }
     },
     {
           { "type": "gflops", "value": "4 * args.batch * args.qSeq * args.kvSeq * args.heads * args.headDim" }
         ]
       }
+    },
+    {
+      "name": "causal-prefill-f16-h4kv2-d64-q64",
+      "provenance": {
+        "source": "synthetic",
+        "test": "Causal half-precision prompt prefill",
+        "notes": "Each query attends only itself and earlier key positions. Verify the complete attention output and exact retention of both key/value inputs in the present caches."
+      },
+      "attrs": { "num_heads": 4, "kv_num_heads": 2 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [1, 64, 256],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.0625 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [1, 64, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.023, "scale": 0.0625 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [1, 64, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.031, "scale": 0.0625 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [63] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [64] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float16", "shape": [1, 64, 256] },
+        "presentKeyT": { "dtype": "float16", "shape": [1, 2, 64, 64] },
+        "presentValueT": { "dtype": "float16", "shape": [1, 2, 64, 64] }
+      },
+      "preset": "model",
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * dim(shapes.queryT, 0) * dim(shapes.queryT, 1) * (dim(shapes.queryT, 1) + 1) * dim(shapes.queryT, 2)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "causal-prefill-f16-h4kv2-d128-q65",
+      "provenance": {
+        "source": "synthetic",
+        "test": "Causal half-precision prompt prefill",
+        "notes": "Each query attends only itself and earlier key positions. Verify the complete attention output and exact retention of both key/value inputs in the present caches."
+      },
+      "attrs": { "num_heads": 4, "kv_num_heads": 2 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [1, 65, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.0625 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [1, 65, 256],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.023, "scale": 0.0625 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [1, 65, 256],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.031, "scale": 0.0625 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [64] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [65] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float16", "shape": [1, 65, 512] },
+        "presentKeyT": { "dtype": "float16", "shape": [1, 2, 65, 128] },
+        "presentValueT": { "dtype": "float16", "shape": [1, 2, 65, 128] }
+      },
+      "preset": "model",
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * dim(shapes.queryT, 0) * dim(shapes.queryT, 1) * (dim(shapes.queryT, 1) + 1) * dim(shapes.queryT, 2)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "causal-prefill-f16-h8kv2-d64-q512",
+      "provenance": {
+        "source": "synthetic",
+        "test": "Causal half-precision prompt prefill",
+        "notes": "Each query attends only itself and earlier key positions. Verify the complete attention output and exact retention of both key/value inputs in the present caches."
+      },
+      "attrs": { "num_heads": 8, "kv_num_heads": 2 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [1, 512, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.0625 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [1, 512, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.023, "scale": 0.0625 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [1, 512, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.031, "scale": 0.0625 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [512] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float16", "shape": [1, 512, 512] },
+        "presentKeyT": { "dtype": "float16", "shape": [1, 2, 512, 64] },
+        "presentValueT": { "dtype": "float16", "shape": [1, 2, 512, 64] }
+      },
+      "preset": "model",
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * dim(shapes.queryT, 0) * dim(shapes.queryT, 1) * (dim(shapes.queryT, 1) + 1) * dim(shapes.queryT, 2)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "prefill-f32-h8kv1-d32-q33",
+      "preset": "model",
+      "provenance": { "notes": "Small bidirectional prefill with query tails and varying head widths." },
+      "vars": { "batch": 1, "qSeq": 33, "kvSeq": 33, "heads": 8, "kvHeads": 1, "headDim": 32 },
+      "attrs": { "num_heads": 8, "kv_num_heads": 1, "scale": 0.125, "causal": 0 },
+      "inputs": {
+        "queryT": { "shape": [1, 33, 256], "dtype": "float32", "dist": "normal", "seed": 1075, "scale": 0.2 },
+        "keyT": { "shape": [1, 33, 32], "dtype": "float32", "dist": "normal", "seed": 1076, "scale": 0.2 },
+        "valueT": { "shape": [1, 33, 32], "dtype": "float32", "dist": "normal", "seed": 1077, "scale": 0.2 },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [32] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [33] } }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 33, 256], "dtype": "float32" },
+        "presentKeyT": { "shape": [1, 1, 33, 32], "dtype": "float32" },
+        "presentValueT": { "shape": [1, 1, 33, 32], "dtype": "float32" }
+      },
+      "bench": {
+        "metrics": [
+          { "type": "gflops", "value": "4 * args.batch * args.qSeq * args.kvSeq * args.heads * args.headDim" }
+        ]
+      }
+    },
+    {
+      "name": "prefill-f32-h8kv1-d128-q33",
+      "preset": "model",
+      "provenance": { "notes": "Small bidirectional prefill with query tails and varying head widths." },
+      "vars": { "batch": 1, "qSeq": 33, "kvSeq": 33, "heads": 8, "kvHeads": 1, "headDim": 128 },
+      "attrs": { "num_heads": 8, "kv_num_heads": 1, "scale": 0.125, "causal": 0 },
+      "inputs": {
+        "queryT": { "shape": [1, 33, 1024], "dtype": "float32", "dist": "normal", "seed": 1075, "scale": 0.2 },
+        "keyT": { "shape": [1, 33, 128], "dtype": "float32", "dist": "normal", "seed": 1076, "scale": 0.2 },
+        "valueT": { "shape": [1, 33, 128], "dtype": "float32", "dist": "normal", "seed": 1077, "scale": 0.2 },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [32] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [33] } }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 33, 1024], "dtype": "float32" },
+        "presentKeyT": { "shape": [1, 1, 33, 128], "dtype": "float32" },
+        "presentValueT": { "shape": [1, 1, 33, 128], "dtype": "float32" }
+      },
+      "bench": {
+        "metrics": [
+          { "type": "gflops", "value": "4 * args.batch * args.qSeq * args.kvSeq * args.heads * args.headDim" }
+        ]
+      }
+    },
+    {
+      "name": "prefill-f32-h8kv1-d256-q32",
+      "preset": "model",
+      "provenance": { "notes": "Small bidirectional prefill with query tails and varying head widths." },
+      "vars": { "batch": 1, "qSeq": 32, "kvSeq": 32, "heads": 8, "kvHeads": 1, "headDim": 256 },
+      "attrs": { "num_heads": 8, "kv_num_heads": 1, "scale": 0.125, "causal": 0 },
+      "inputs": {
+        "queryT": { "shape": [1, 32, 2048], "dtype": "float32", "dist": "normal", "seed": 1075, "scale": 0.2 },
+        "keyT": { "shape": [1, 32, 256], "dtype": "float32", "dist": "normal", "seed": 1076, "scale": 0.2 },
+        "valueT": { "shape": [1, 32, 256], "dtype": "float32", "dist": "normal", "seed": 1077, "scale": 0.2 },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [31] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [32] } }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 32, 2048], "dtype": "float32" },
+        "presentKeyT": { "shape": [1, 1, 32, 256], "dtype": "float32" },
+        "presentValueT": { "shape": [1, 1, 32, 256], "dtype": "float32" }
+      },
+      "bench": {
+        "metrics": [
+          { "type": "gflops", "value": "4 * args.batch * args.qSeq * args.kvSeq * args.heads * args.headDim" }
+        ]
+      }
+    },
+    {
+      "name": "cooperative-window-h4kv2-d122-cap256",
+      "preset": "smoke",
+      "attrs": { "num_heads": 4, "kv_num_heads": 2, "sliding_window_cache": 1, "local_window_size": 256 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 488],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 244],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.031, "cosStep": 0.007 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 244],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.037, "cosStep": 0.005 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 2, 256, 122],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.017, "cosStep": 0.019 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 2, 256, 122],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [65535] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [65536] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 1, 488] },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 2, 256, 122] },
+        "presentValueT": { "dtype": "float32", "shape": [1, 2, 256, 122] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * 2 * numel(shapes.pastKeyT) * 4 + numel(shapes.queryT) * 4 + numel(shapes.outputT) * 4"
+          }
+        ]
+      }
+    },
+    {
+      "name": "cooperative-window-h4kv2-d124-cap256",
+      "preset": "smoke",
+      "attrs": { "num_heads": 4, "kv_num_heads": 2, "sliding_window_cache": 1, "local_window_size": 256 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 496],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 248],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.031, "cosStep": 0.007 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 248],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.037, "cosStep": 0.005 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 2, 256, 124],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.017, "cosStep": 0.019 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 2, 256, 124],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [65535] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [65536] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 1, 496] },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 2, 256, 124] },
+        "presentValueT": { "dtype": "float32", "shape": [1, 2, 256, 124] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * 2 * numel(shapes.pastKeyT) * 4 + numel(shapes.queryT) * 4 + numel(shapes.outputT) * 4"
+          }
+        ]
+      }
+    },
+    {
+      "name": "cooperative-window-h4kv2-d246-cap256",
+      "preset": "smoke",
+      "attrs": { "num_heads": 4, "kv_num_heads": 2, "sliding_window_cache": 1, "local_window_size": 256 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 984],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 492],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.031, "cosStep": 0.007 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 492],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.037, "cosStep": 0.005 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 2, 256, 246],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.017, "cosStep": 0.019 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 2, 256, 246],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [65535] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [65536] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 1, 984] },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 2, 256, 246] },
+        "presentValueT": { "dtype": "float32", "shape": [1, 2, 256, 246] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * 2 * numel(shapes.pastKeyT) * 4 + numel(shapes.queryT) * 4 + numel(shapes.outputT) * 4"
+          }
+        ]
+      }
+    },
+    {
+      "name": "cooperative-window-h4kv2-d248-cap256",
+      "preset": "smoke",
+      "attrs": { "num_heads": 4, "kv_num_heads": 2, "sliding_window_cache": 1, "local_window_size": 256 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 992],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 496],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.031, "cosStep": 0.007 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 496],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.037, "cosStep": 0.005 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 2, 256, 248],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.017, "cosStep": 0.019 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 2, 256, 248],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [65535] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [65536] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 1, 992] },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 2, 256, 248] },
+        "presentValueT": { "dtype": "float32", "shape": [1, 2, 256, 248] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * 2 * numel(shapes.pastKeyT) * 4 + numel(shapes.queryT) * 4 + numel(shapes.outputT) * 4"
+          }
+        ]
+      }
+    },
+    {
+      "name": "cooperative-window-h4kv2-d256-cap256",
+      "preset": "smoke",
+      "attrs": { "num_heads": 4, "kv_num_heads": 2, "sliding_window_cache": 1, "local_window_size": 256 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 1024],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 512],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.031, "cosStep": 0.007 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 512],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.037, "cosStep": 0.005 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 2, 256, 256],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.017, "cosStep": 0.019 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 2, 256, 256],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [65535] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [65536] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 1, 1024] },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 2, 256, 256] },
+        "presentValueT": { "dtype": "float32", "shape": [1, 2, 256, 256] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * 2 * numel(shapes.pastKeyT) * 4 + numel(shapes.queryT) * 4 + numel(shapes.outputT) * 4"
+          }
+        ]
+      }
     }
   ]
 }

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

@@ -2,8 +2,12 @@
 {% set quantized = quantized | default(false) %}
 {% set bits = bits | default(0) %}
 {% if useSeqlens is not defined %}{% set useSeqlens = false %}{% endif %}
-{% if usesF16 %}enable f16;
-{% endif %}{{ env.wgsl.resourceDeclarations }}
 // f16 inputs widen to f32 on load; computation and shared memory remain f32,
 // and results narrow only on store. The casts are identities for f32.
 {% set IO = "f16" if usesF16 else "f32" %}
@@ -29,15 +33,13 @@ const Q_HIDDEN: u32 = {{ qHidden }}u;
 const GROUP: u32 = {{ qHeads }}u / {{ kvHeads }}u;
 const PACKED: u32 = {{ packed }}u;
 {% if hasQNorm %}const QK_EPS: f32 = {{ qkEps }};{% endif %}
-{% if cooperative %}const WG: u32 = {{ tunables.COOPERATIVE_WORKGROUP_SIZE }}u;{% else %}const WG: u32 = {{ scalarWorkgroupSize }}u;{% endif %}
 const NEG_INF: 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 }}));
 }
 {% if quantized %}
 fn kscale(d: u32, hk: u32) -> f32 { return k_scale[select(0u, hk * HEAD_DIM + d, params.perChannel != 0u)]; }
 fn vscale(d: u32, hk: u32) -> f32 { return v_scale[select(0u, hk * HEAD_DIM + d, params.perChannel != 0u)]; }
@@ -67,10 +69,15 @@ fn read_v(base: u32, d: u32{% if quantized %}, hk: u32{% endif %}) -> f32 {
 }
 {% if cooperative %}
-var<workgroup> qsh: array<f32, HEAD_DIM>;       // query head-vector (post norm/rotary)
-var<workgroup> acc_sh: array<f32, HEAD_DIM * 32u>; // per-lane V accumulator [lane*HEAD_DIM + d]
-var<workgroup> m_sh: array<f32, 32u>;
-var<workgroup> l_sh: array<f32, 32u>;
 @compute @workgroup_size(WG)
 fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
@@ -92,42 +99,47 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid
 {% else %}
   let activeEnd = totalSeq;
 {% endif %}
-  let absPos = activeEnd - params.qSeq + s;
   // Cooperative query-vector prep (post norm/rotary) into shared memory.
   let qBase = (b * params.qSeq + s) * Q_HIDDEN + h * HEAD_DIM;
-  for (var d = lane; d < HEAD_DIM; d = d + WG) { qsh[d] = f32(query[qBase + d]); }
   workgroupBarrier();
 {% if hasQNorm %}
   var part = 0.0;
-  for (var d = lane; d < HEAD_DIM; d = d + WG) { part = part + qsh[d] * qsh[d]; }
-  m_sh[lane] = part;
   workgroupBarrier();
   var ms = 0.0;
-  for (var L = 0u; L < WG; L = L + 1u) { ms = ms + m_sh[L]; }
   let invRms = inverseSqrt(ms / f32(HEAD_DIM) + QK_EPS);
-  for (var d = lane; d < HEAD_DIM; d = d + WG) { qsh[d] = qsh[d] * invRms * f32(q_norm_weight[d]); }
   workgroupBarrier();
 {% endif %}
 {% if hasRotary %}
   // Each lane owns disjoint pairs (d, d+HALF), so the in-place rotate is safe.
   for (var d = lane; d < HALF; d = d + WG) {
-    let cs = f32(cos_cache[absPos * HALF + d]);
-    let sn = f32(sin_cache[absPos * HALF + d]);
-    let x0 = qsh[d];
-    let x1 = qsh[d + HALF];
-    qsh[d] = x0 * cs - x1 * sn;
-    qsh[d + HALF] = x1 * cs + x0 * sn;
   }
   workgroupBarrier();
 {% endif %}
   let scale = scale_value();
-  let maxKj = absPos + 1u;
   var minKj = 0u;
   if (params.windowSize > 0u && maxKj > params.windowSize) { minKj = maxKj - params.windowSize; }
   let accBase = lane * HEAD_DIM;
-  for (var d = 0u; d < HEAD_DIM; d = d + 1u) { acc_sh[accBase + d] = 0.0; }
   var m = NEG_INF;
   var l = 0.0;
@@ -135,34 +147,34 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid
   for (var j = minKj + lane; j < maxKj; j = j + WG) {
     let base = ((b * KV_HEADS + hk) * totalSeq + j) * PACKED;
     var dot = 0.0;
-    for (var d = 0u; d < HEAD_DIM; d = d + 1u) { dot = dot + qsh[d] * read_k(base, d{% if quantized %}, hk{% endif %}); }
     var score = dot * scale;
     if (params.softcap > 0.0) { score = params.softcap * tanh(clamp(score / params.softcap, -30.0, 30.0)); }
 {% if hasBias %}
     let bb = select(b, 0u, params.biasBatch == 1u);
     let bh = select(h, 0u, params.biasHeads == 1u);
-    score = score + f32(attn_bias[((bb * params.biasHeads + bh) * params.qSeq + s) * totalSeq + j]);
 {% endif %}
     let newM = max(m, score);
     let corr = exp(m - newM);
     let p = exp(score - newM);
     l = l * corr + p;
-    for (var d = 0u; d < HEAD_DIM; d = d + 1u) { acc_sh[accBase + d] = acc_sh[accBase + d] * corr + p * read_v(base, d{% if quantized %}, hk{% endif %}); }
     m = newM;
   }
   // Flash merge across the WG lanes: global max, rescale, summed denom.
-  m_sh[lane] = m;
   workgroupBarrier();
   var gm = NEG_INF;
-  for (var L = 0u; L < WG; L = L + 1u) { gm = max(gm, m_sh[L]); }
   let fctr = exp(m - gm);
   l = l * fctr;
-  for (var d = 0u; d < HEAD_DIM; d = d + 1u) { acc_sh[accBase + d] = acc_sh[accBase + d] * fctr; }
-  l_sh[lane] = l;
   workgroupBarrier();
   var glsum = 0.0;
-  for (var L = 0u; L < WG; L = L + 1u) { glsum = glsum + l_sh[L]; }
   var sink = 0.0;
 {% if hasHeadSink %}
@@ -187,7 +199,7 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid
   let oBase = (b * params.qSeq + s) * Q_HIDDEN + h * HEAD_DIM;
   for (var d = lane; d < HEAD_DIM; d = d + WG) {
     var sumacc = 0.0;
-    for (var L = 0u; L < WG; L = L + 1u) { sumacc = sumacc + acc_sh[L * HEAD_DIM + d]; }
     output[oBase + d] = {{ IO }}(sumacc * invDenom);
   }
 }
@@ -215,7 +227,12 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid
 {% else %}
   let activeEnd = totalSeq;
 {% endif %}
-  let absPos = activeEnd - params.qSeq + s;
   // Private query-vector prep (post norm/rotary).
   var q: array<f32, HEAD_DIM>;
@@ -229,8 +246,8 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid
 {% endif %}
 {% if hasRotary %}
   for (var d = 0u; d < HALF; d = d + 1u) {
-    let cs = f32(cos_cache[absPos * HALF + d]);
-    let sn = f32(sin_cache[absPos * HALF + d]);
     let x0 = q[d];
     let x1 = q[d + HALF];
     q[d] = x0 * cs - x1 * sn;
@@ -238,7 +255,7 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid
   }
 {% endif %}
   let scale = scale_value();
-  let maxKj = absPos + 1u;
   var minKj = 0u;
   if (params.windowSize > 0u && maxKj > params.windowSize) { minKj = maxKj - params.windowSize; }
@@ -256,7 +273,7 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid
 {% if hasBias %}
     let bb = select(b, 0u, params.biasBatch == 1u);
     let bh = select(h, 0u, params.biasHeads == 1u);
-    score = score + f32(attn_bias[((bb * params.biasHeads + bh) * params.qSeq + s) * totalSeq + j]);
 {% endif %}
     let newM = max(m, score);
     let corr = exp(m - newM);

 {% set quantized = quantized | default(false) %}
 {% set bits = bits | default(0) %}
 {% if useSeqlens is not defined %}{% set useSeqlens = false %}{% endif %}
+{% if windowOrigin is not defined %}{% set windowOrigin = false %}{% endif %}
+{% set needsOrigin = windowOrigin and (hasRotary or hasBias) %}
+{% set ropePos = "(absPos + origin)" if needsOrigin else "absPos" %}
+{% set biasStride = "params.biasCols" if needsOrigin else "totalSeq" %}
+{% set biasCol = "(origin + j)" if needsOrigin else "j" %}
+{{ env.wgsl.resourceDeclarations }}
 // f16 inputs widen to f32 on load; computation and shared memory remain f32,
 // and results narrow only on store. The casts are identities for f32.
 {% set IO = "f16" if usesF16 else "f32" %}
 const GROUP: u32 = {{ qHeads }}u / {{ kvHeads }}u;
 const PACKED: u32 = {{ packed }}u;
 {% if hasQNorm %}const QK_EPS: f32 = {{ qkEps }};{% endif %}
+{% if cooperative %}const WG: u32 = {{ gqaCooperativeWorkgroupSize }}u;{% else %}const WG: u32 = {{ scalarWorkgroupSize }}u;{% endif %}
 const NEG_INF: 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 }}));
 }
 {% if quantized %}
 fn kscale(d: u32, hk: u32) -> f32 { return k_scale[select(0u, hk * HEAD_DIM + d, params.perChannel != 0u)]; }
 fn vscale(d: u32, hk: u32) -> f32 { return v_scale[select(0u, hk * HEAD_DIM + d, params.perChannel != 0u)]; }
 }
 {% if cooperative %}
+// One allocation avoids the per-variable 16-byte padding between arrays.
+// The manifest budgets the structure's complete size, rounded to 16 bytes.
+struct CooperativeState {
+  q: array<f32, HEAD_DIM>,
+  acc: array<f32, HEAD_DIM * WG>,
+  m: array<f32, WG>,
+  l: array<f32, WG>,
+}
+var<workgroup> state: CooperativeState;
 @compute @workgroup_size(WG)
 fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
 {% else %}
   let activeEnd = totalSeq;
 {% endif %}
+{% if not bidirectional or hasRotary %}  let absPos = activeEnd - params.qSeq + s;
+{% endif %}{% if needsOrigin %}
+  // Absolute position of cache row 0 (windowed cache): the query end is the
+  // unclamped total while the cache holds only its last `activeEnd` rows.
+  let origin = max(params.qSeq, u32(seqlens_k[b]) + 1u) - activeEnd;
+{% endif %}
   // Cooperative query-vector prep (post norm/rotary) into shared memory.
   let qBase = (b * params.qSeq + s) * Q_HIDDEN + h * HEAD_DIM;
+  for (var d = lane; d < HEAD_DIM; d = d + WG) { state.q[d] = f32(query[qBase + d]); }
   workgroupBarrier();
 {% if hasQNorm %}
   var part = 0.0;
+  for (var d = lane; d < HEAD_DIM; d = d + WG) { part = part + state.q[d] * state.q[d]; }
+  state.m[lane] = part;
   workgroupBarrier();
   var ms = 0.0;
+  for (var L = 0u; L < WG; L = L + 1u) { ms = ms + state.m[L]; }
   let invRms = inverseSqrt(ms / f32(HEAD_DIM) + QK_EPS);
+  for (var d = lane; d < HEAD_DIM; d = d + WG) { state.q[d] = state.q[d] * invRms * f32(q_norm_weight[d]); }
   workgroupBarrier();
 {% endif %}
 {% if hasRotary %}
   // Each lane owns disjoint pairs (d, d+HALF), so the in-place rotate is safe.
   for (var d = lane; d < HALF; d = d + WG) {
+    let cs = f32(cos_cache[{{ ropePos }} * HALF + d]);
+    let sn = f32(sin_cache[{{ ropePos }} * HALF + d]);
+    let x0 = state.q[d];
+    let x1 = state.q[d + HALF];
+    state.q[d] = x0 * cs - x1 * sn;
+    state.q[d + HALF] = x1 * cs + x0 * sn;
   }
   workgroupBarrier();
 {% endif %}
   let scale = scale_value();
+  let maxKj = {% if bidirectional %}activeEnd{% else %}absPos + 1u{% endif %};
   var minKj = 0u;
   if (params.windowSize > 0u && maxKj > params.windowSize) { minKj = maxKj - params.windowSize; }
   let accBase = lane * HEAD_DIM;
+  for (var d = 0u; d < HEAD_DIM; d = d + 1u) { state.acc[accBase + d] = 0.0; }
   var m = NEG_INF;
   var l = 0.0;
   for (var j = minKj + lane; j < maxKj; j = j + WG) {
     let base = ((b * KV_HEADS + hk) * totalSeq + j) * PACKED;
     var dot = 0.0;
+    for (var d = 0u; d < HEAD_DIM; d = d + 1u) { dot = dot + state.q[d] * read_k(base, d{% if quantized %}, hk{% endif %}); }
     var score = dot * scale;
     if (params.softcap > 0.0) { score = params.softcap * tanh(clamp(score / params.softcap, -30.0, 30.0)); }
 {% if hasBias %}
     let bb = select(b, 0u, params.biasBatch == 1u);
     let bh = select(h, 0u, params.biasHeads == 1u);
+    score = score + f32(attn_bias[((bb * params.biasHeads + bh) * params.qSeq + s) * {{ biasStride }} + {{ biasCol }}]);
 {% endif %}
     let newM = max(m, score);
     let corr = exp(m - newM);
     let p = exp(score - newM);
     l = l * corr + p;
+    for (var d = 0u; d < HEAD_DIM; d = d + 1u) { state.acc[accBase + d] = state.acc[accBase + d] * corr + p * read_v(base, d{% if quantized %}, hk{% endif %}); }
     m = newM;
   }
   // Flash merge across the WG lanes: global max, rescale, summed denom.
+  state.m[lane] = m;
   workgroupBarrier();
   var gm = NEG_INF;
+  for (var L = 0u; L < WG; L = L + 1u) { gm = max(gm, state.m[L]); }
   let fctr = exp(m - gm);
   l = l * fctr;
+  for (var d = 0u; d < HEAD_DIM; d = d + 1u) { state.acc[accBase + d] = state.acc[accBase + d] * fctr; }
+  state.l[lane] = l;
   workgroupBarrier();
   var glsum = 0.0;
+  for (var L = 0u; L < WG; L = L + 1u) { glsum = glsum + state.l[L]; }
   var sink = 0.0;
 {% if hasHeadSink %}
   let oBase = (b * params.qSeq + s) * Q_HIDDEN + h * HEAD_DIM;
   for (var d = lane; d < HEAD_DIM; d = d + WG) {
     var sumacc = 0.0;
+    for (var L = 0u; L < WG; L = L + 1u) { sumacc = sumacc + state.acc[L * HEAD_DIM + d]; }
     output[oBase + d] = {{ IO }}(sumacc * invDenom);
   }
 }
 {% else %}
   let activeEnd = totalSeq;
 {% endif %}
+{% if not bidirectional or hasRotary %}  let absPos = activeEnd - params.qSeq + s;
+{% endif %}{% if needsOrigin %}
+  // Absolute position of cache row 0 (windowed cache): the query end is the
+  // unclamped total while the cache holds only its last `activeEnd` rows.
+  let origin = max(params.qSeq, u32(seqlens_k[b]) + 1u) - activeEnd;
+{% endif %}
   // Private query-vector prep (post norm/rotary).
   var q: array<f32, HEAD_DIM>;
 {% endif %}
 {% if hasRotary %}
   for (var d = 0u; d < HALF; d = d + 1u) {
+    let cs = f32(cos_cache[{{ ropePos }} * HALF + d]);
+    let sn = f32(sin_cache[{{ ropePos }} * HALF + d]);
     let x0 = q[d];
     let x1 = q[d + HALF];
     q[d] = x0 * cs - x1 * sn;
   }
 {% endif %}
   let scale = scale_value();
+  let maxKj = {% if bidirectional %}activeEnd{% else %}absPos + 1u{% endif %};
   var minKj = 0u;
   if (params.windowSize > 0u && maxKj > params.windowSize) { minKj = maxKj - params.windowSize; }
 {% if hasBias %}
     let bb = select(b, 0u, params.biasBatch == 1u);
     let bh = select(h, 0u, params.biasHeads == 1u);
+    score = score + f32(attn_bias[((bb * params.biasHeads + bh) * params.qSeq + s) * {{ biasStride }} + {{ biasCol }}]);
 {% endif %}
     let newM = max(m, score);
     let corr = exp(m - newM);

build/webgpu/gqa-present.wgsl.jinja CHANGED Viewed

@@ -9,8 +9,15 @@
 {% set bits = bits | default(0) %}
 {% set qmax = qmax | default(0) %}
 {% set qmin = qmin | default(0) %}
-{% if usesF16 is defined and usesF16 %}enable f16;
-{% endif %}{{ env.wgsl.resourceDeclarations }}
 {% if mode == "transpose" %}
 const KV_HEADS: u32 = {{ kvNumHeads }}u;
@@ -31,9 +38,7 @@ const KV_HIDDEN: u32 = {{ kvHidden }}u;
 fn main(
   @builtin(global_invocation_id) gid: vec3<u32>
 ) {
-  // The dispatch folds oversized one-dimensional grids into x/y. Rebuild the
-  // flat invocation index; this reduces to gid.x when no fold is needed.
-  let index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
 {% if presentVec4 %}
   let total = params.batchSize * KV_HEADS * params.kvSeq * HEAD_DIM_V4;
   if (index >= total) {
@@ -60,9 +65,11 @@ fn main(
   present_value[index] = {{ presentScalar }}(value[packed_index]);
 {% endif %}
 }
-{%- else %}
 {% set shareAppend = mode == "merge_share" and shareRegion == "append" %}
 {% set cooperativeCopy = device.adapterInfo.vendor != "apple" %}
 {% set cooperativeMerge = cooperativeCopy or shareAppend %}
 {% set cooperativeMode = cooperativeMerge and (mode == "copy" or mode == "merge" or mode == "merge_share") %}
@@ -95,7 +102,7 @@ const HEAD_DIM: u32 = {{ headDim }}u;
 {% if mode != "copy" %}const KV_HEADS: u32 = {{ kvHeads }}u;
 {% endif %}
 const WG: u32 = {{ tunables.COPY_WORKGROUP_SIZE }}u;
-{% if mode != "copy" and (mode != "merge_share" or shareAppend) %}const KV_HIDDEN: u32 = {{ kvHeads }}u * {{ headDim }}u;
 {% endif %}
 {% if mode == "build_quant" or mode == "append_quant" %}const PACKED: u32 = {{ packed }}u;
 {% endif %}
@@ -153,8 +160,7 @@ fn main({% if needsGid %}@builtin(global_invocation_id) gid: vec3<u32>{% if coop
 {% endmacro %}
 {% macro append_per_row_walk() %}
   // Per-row append walk: one thread streams its row's contiguous bytes.
-  let i = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
-  if (i >= params.count) { return; }
   let j = i % params.keySeq;
   let tmp = i / params.keySeq;
   let hk = tmp % KV_HEADS;
@@ -263,8 +269,7 @@ fn main({% if needsGid %}@builtin(global_invocation_id) gid: vec3<u32>{% if coop
 {% endif %}
 {% elif mode == "copy" %}
   // Thread i streams one contiguous (batch, kvHead, token) row.
-  let i = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
-  if (i >= params.count) { return; }
   let base = i * HEAD_DIM;
   for (var d = 0u; d < HEAD_DIM; d = d + 1u) {
     present_key[base + d] = src_k[base + d];
@@ -272,8 +277,7 @@ fn main({% if needsGid %}@builtin(global_invocation_id) gid: vec3<u32>{% if coop
   }
 {% elif mode == "merge" %}
   // Per-row merge walk.
-  let i = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
-  if (i >= params.count) { return; }
   let t = i % params.seq;
   let tmp = i / params.seq;
   let hk = tmp % KV_HEADS;
@@ -296,8 +300,7 @@ fn main({% if needsGid %}@builtin(global_invocation_id) gid: vec3<u32>{% if coop
 {% elif mode == "merge_share" %}
   // Past and present share the full-capacity stride, so outside-window rows
   // copy at the same index.
-  let i = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
-  if (i >= params.count) { return; }
   let t = i % params.seq;
   let tmp = i / params.seq;
   let hk = tmp % KV_HEADS;
@@ -316,28 +319,51 @@ fn main({% if needsGid %}@builtin(global_invocation_id) gid: vec3<u32>{% if coop
   }
   }
 {% elif mode == "window_shift" %}
-  // Windowed cache: `params.seq` is a fixed capacity C, and the invariant is that
-  // the L = min(T, C) most recent tokens live contiguously at rows [0, L).
-  //
-  //   residentBefore = min(T - S, C)      E = max(0, residentBefore + S - C)
-  //   rows [0, appendStart)          <- past rows shifted down by E
-  //   rows [appendStart, append+S)   <- the new tokens
-  //   rows beyond that               <- outside the window, cleared
-  //
-  // The shift reads past and writes present, which must be distinct buffers: an
-  // in-place compaction would have one invocation overwrite row t while another
-  // still needs it as the source for row t-E, with no ordering between them.
-  //
-  // Attention needs no change for this layout. It derives its key range from
-  // `min(capacity, seqlens_k[b] + 1)`, which is exactly L, and both the causal
-  // and local-window masks depend only on the query/key distance:
-  //   q_abs - k_abs = (T - qSeq + s) - (origin + t) = (L - qSeq + s) - t
-  // so scoring a windowed cache as if it were a full L-length one is the same
-  // arithmetic. Only RoPE would need the true absolute position, which is why
-  // `windowShiftOk` refuses a rotary request outright rather than silently
-  // rotating at the cache row.
-  let i = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
-  if (i >= params.count) { return; }
   let t = i % params.seq;
   let tmp = i / params.seq;
   let hk = tmp % KV_HEADS;
@@ -358,6 +384,9 @@ fn main({% if needsGid %}@builtin(global_invocation_id) gid: vec3<u32>{% if coop
       present_key[dstBase + d] = past_k[srcBase + d];
       present_value[dstBase + d] = past_v[srcBase + d];
     }
   } else if (t < appendStart + params.keySeq) {
     let nSrc = (b * params.keySeq + (t - appendStart)) * KV_HIDDEN + hk * HEAD_DIM;
     for (var d = 0u; d < HEAD_DIM; d = d + 1u) {
@@ -365,6 +394,7 @@ fn main({% if needsGid %}@builtin(global_invocation_id) gid: vec3<u32>{% if coop
       present_value[dstBase + d] = new_v[nSrc + d];
     }
   } else {
     // Outside the resident window. Cleared rather than left stale so the present
     // buffer does not expose stale contents from its distinct allocation.
     for (var d = 0u; d < HEAD_DIM; d = d + 1u) {
@@ -372,11 +402,9 @@ fn main({% if needsGid %}@builtin(global_invocation_id) gid: vec3<u32>{% if coop
       present_value[dstBase + d] = {{ zeroScalar }}(0.0);
     }
   }
 {% else %}
-  // The dispatch folds oversized one-dimensional grids into x/y. Rebuild the
-  // flat invocation index; this reduces to gid.x when no fold is needed.
-  let i = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
-  if (i >= params.count) { return; }
   let t = i % params.seq;
   let tmp = i / params.seq;
   let hk = tmp % KV_HEADS;
@@ -435,4 +463,4 @@ fn main({% if needsGid %}@builtin(global_invocation_id) gid: vec3<u32>{% if coop
 {% endif %}
 {% endif %}
 }
-{%- endif %}

 {% set bits = bits | default(0) %}
 {% set qmax = qmax | default(0) %}
 {% set qmin = qmin | default(0) %}
+{% 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.
+{% if bound == "" %}
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};{% else %}
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
+  if ({{ name }} >= {{ bound }}) { return; }{% endif %}{% endmacro %}
+{{ env.wgsl.resourceDeclarations }}
 {% if mode == "transpose" %}
 const KV_HEADS: u32 = {{ kvNumHeads }}u;
 fn main(
   @builtin(global_invocation_id) gid: vec3<u32>
 ) {
+{{ flat_index_2d("WG", "index", "") }}
 {% if presentVec4 %}
   let total = params.batchSize * KV_HEADS * params.kvSeq * HEAD_DIM_V4;
   if (index >= total) {
   present_value[index] = {{ presentScalar }}(value[packed_index]);
 {% endif %}
 }
+{% else %}
 {% set shareAppend = mode == "merge_share" and shareRegion == "append" %}
+{% set shiftRegion = shiftRegion | default("all") %}
+{% set shiftRetainOnly = mode == "window_shift" and shiftRegion == "shift" %}
 {% set cooperativeCopy = device.adapterInfo.vendor != "apple" %}
 {% set cooperativeMerge = cooperativeCopy or shareAppend %}
 {% set cooperativeMode = cooperativeMerge and (mode == "copy" or mode == "merge" or mode == "merge_share") %}
 {% if mode != "copy" %}const KV_HEADS: u32 = {{ kvHeads }}u;
 {% endif %}
 const WG: u32 = {{ tunables.COPY_WORKGROUP_SIZE }}u;
+{% if mode != "copy" and (mode != "merge_share" or shareAppend) and not shiftRetainOnly %}const KV_HIDDEN: u32 = {{ kvHeads }}u * {{ headDim }}u;
 {% endif %}
 {% if mode == "build_quant" or mode == "append_quant" %}const PACKED: u32 = {{ packed }}u;
 {% endif %}
 {% endmacro %}
 {% macro append_per_row_walk() %}
   // Per-row append walk: one thread streams its row's contiguous bytes.
+{{ flat_index_2d("WG", guardInline=true) }}
   let j = i % params.keySeq;
   let tmp = i / params.keySeq;
   let hk = tmp % KV_HEADS;
 {% endif %}
 {% elif mode == "copy" %}
   // Thread i streams one contiguous (batch, kvHead, token) row.
+{{ flat_index_2d("WG", guardInline=true) }}
   let base = i * HEAD_DIM;
   for (var d = 0u; d < HEAD_DIM; d = d + 1u) {
     present_key[base + d] = src_k[base + d];
   }
 {% elif mode == "merge" %}
   // Per-row merge walk.
+{{ flat_index_2d("WG", guardInline=true) }}
   let t = i % params.seq;
   let tmp = i / params.seq;
   let hk = tmp % KV_HEADS;
 {% elif mode == "merge_share" %}
   // Past and present share the full-capacity stride, so outside-window rows
   // copy at the same index.
+{{ flat_index_2d("WG", guardInline=true) }}
   let t = i % params.seq;
   let tmp = i / params.seq;
   let hk = tmp % KV_HEADS;
   }
   }
 {% elif mode == "window_shift" %}
+  // Keep the most recent min(total, capacity) tokens in contiguous cache rows.
+  // Read past and write a distinct present buffer so compaction cannot race.
+  // Rotary append uses the absolute token position; retained keys are already rotated.
+{% if shiftRegion == "append" %}
+  // Append-only pass: the thread range spans just the appended window (batch x
+  // kvHead x keySeq rows), while `params.seq` stays the cache capacity. The shift
+  // pass owns every other row and the two ranges are disjoint, so their order is
+  // immaterial.
+{{ flat_index_2d("WG", guardInline=true) }}
+  let j = i % params.keySeq;
+  let tmp = i / params.keySeq;
+  let hk = tmp % KV_HEADS;
+  let b = tmp / KV_HEADS;
+  let absTotal = max(params.keySeq, u32(seqlens_k[b]) + 1u);
+  let residentBefore = min(absTotal - params.keySeq, params.seq);
+  // `max(0u, a + b - c)` does not clamp in u32: the subtraction wraps first, so a
+  // step that evicts nothing reads back ~2^32 instead of 0. Compare before
+  // subtracting.
+  let filled = residentBefore + params.keySeq;
+  let evicted = select(0u, filled - params.seq, filled > params.seq);
+  let appendStart = residentBefore - evicted;
+  let t = appendStart + j;
+  let dstBase = ((b * KV_HEADS + hk) * params.seq + t) * HEAD_DIM;
+  let nSrc = (b * params.keySeq + j) * KV_HIDDEN + hk * HEAD_DIM;
+  // Row appendStart + j carries absolute position absTotal - keySeq + j: the cache
+  // holds only the last L positions, so the row index is short of the absolute one
+  // by the window origin. The reference rotates the same rows at
+  // `appendStart + cacheOrigin`, which is the same number.
+  let pos = absTotal - params.keySeq + j;
+  var k: array<f32, HEAD_DIM>;
+  for (var d = 0u; d < HEAD_DIM; d = d + 1u) { k[d] = new_k[nSrc + d]; }
+  for (var d = 0u; d < HALF; d = d + 1u) {
+    let cs = cos_cache[pos * HALF + d];
+    let sn = sin_cache[pos * HALF + d];
+    let x0 = k[d];
+    let x1 = k[d + HALF];
+    k[d] = x0 * cs - x1 * sn;
+    k[d + HALF] = x1 * cs + x0 * sn;
+  }
+  for (var d = 0u; d < HEAD_DIM; d = d + 1u) {
+    present_key[dstBase + d] = k[d];
+    present_value[dstBase + d] = new_v[nSrc + d];
+  }
+{% else %}
+{{ flat_index_2d("WG", guardInline=true) }}
   let t = i % params.seq;
   let tmp = i / params.seq;
   let hk = tmp % KV_HEADS;
       present_key[dstBase + d] = past_k[srcBase + d];
       present_value[dstBase + d] = past_v[srcBase + d];
     }
+{% if shiftRegion == "shift" %}
+  } else if (t >= appendStart + params.keySeq) {
+{% else %}
   } else if (t < appendStart + params.keySeq) {
     let nSrc = (b * params.keySeq + (t - appendStart)) * KV_HIDDEN + hk * HEAD_DIM;
     for (var d = 0u; d < HEAD_DIM; d = d + 1u) {
       present_value[dstBase + d] = new_v[nSrc + d];
     }
   } else {
+{% endif %}
     // Outside the resident window. Cleared rather than left stale so the present
     // buffer does not expose stale contents from its distinct allocation.
     for (var d = 0u; d < HEAD_DIM; d = d + 1u) {
       present_value[dstBase + d] = {{ zeroScalar }}(0.0);
     }
   }
+{% endif %}
 {% else %}
+{{ flat_index_2d("WG", guardInline=true) }}
   let t = i % params.seq;
   let tmp = i / params.seq;
   let hk = tmp % KV_HEADS;
 {% endif %}
 {% endif %}
 }
+{% endif %}

build/webgpu/gqa-qprep.wgsl.jinja CHANGED Viewed

@@ -1,5 +1,9 @@
-{% if usesF16 is defined and usesF16 %}enable f16;
-{% endif %}{{ env.wgsl.resourceDeclarations }}
 // One invocation per (batch, query head, query token) applies optional per-head
 // RMS normalization followed by NeoX half-split rotary embedding at the query's
@@ -21,9 +25,7 @@ const QK_EPS: f32 = {{ qkEps }};
 @compute @workgroup_size(WG)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-  // 2D-folded flat index: gid.y carries the high bits past the per-axis dispatch fold width.
-  // Reduces to gid.x when the dispatch does not fold.
-  let qi = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
   let total = params.batch * Q_HEADS * params.qSeq;
   if (qi >= total) { return; }
   let s = qi % params.qSeq;

+{% 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 }};{% endmacro %}
+{{ env.wgsl.resourceDeclarations }}
 // One invocation per (batch, query head, query token) applies optional per-head
 // RMS normalization followed by NeoX half-split rotary embedding at the query's
 @compute @workgroup_size(WG)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d("WG", "qi", "") }}
   let total = params.batch * Q_HEADS * params.qSeq;
   if (qi >= total) { return; }
   let s = qi % params.qSeq;

build/webgpu/manifest.json CHANGED Viewed

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

build/webgpu/metadata.json CHANGED Viewed

@@ -1,32 +1,32 @@
 {
   "name": "com.microsoft.GroupQueryAttention",
-  "id": "_com_microsoft_groupqueryattention_webgpu_1ea0022",
   "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": "y3ubiijNd5rGJa2KVdluBdaoKYAGeqrPVd6KSppJPHw=",
-      "attn-flash-decode-splitk.wgsl.jinja": "xRRcDB3IGuE25xq3sG1qVPHRsESDmfA15KFVKlq5l4w=",
-      "attn-flash-online.wgsl.jinja": "ontnzJw9RHN7DklqqyavsAiRuKrsnWuEiU+vllbKfGE=",
-      "attn-flash-prefill-cluster.wgsl.jinja": "UHb5IdDeGHB/0gsQdAHcfg9FI73nyzk3zoZiIhaWbgI=",
-      "attn-flash-q32-broadcast.wgsl.jinja": "U/O9TNr3pXvFP+VZ2EAJLuXaWSz62De0TQEmEiU7TLc=",
-      "attn-materialized-rowstats-combine-f32.wgsl.jinja": "zcR02XUVlPyYeLeWH7gLOX65BcozZ1qje9hSIZSOaBA=",
-      "attn-materialized-sgmat-f32.wgsl.jinja": "rNruz2BmIcsQeu5eUnQZh41ftfDt/PjqnZiy2XT+9bk=",
-      "attn-online-scalar.wgsl.jinja": "AJGUYoMkrTHnwCPxu1gVp2CIrdHlca5PMgr9mgqukz4=",
-      "bench.json": "bUGLPxdrsT20K/2ptNbqVs+28i4d2W/9vOyPld+joKs=",
-      "gqa-attention.wgsl.jinja": "9JBHgllUq4dekL6Y1uXkyi2NoEYKjHTkHqnhktbgdzY=",
-      "gqa-present.wgsl.jinja": "/63mkauF/TKCsH4dxgIVlTZnREzSd/j64WGkvu6i2ls=",
-      "gqa-qprep.wgsl.jinja": "kRFlYxr5SdQJlbjmcF/Ii+0w56q0xhF8jyAOKPNx47I=",
-      "manifest.json": "U9Xgwr8xQCIMv/BNUNjiaMDyUjEVXkawKtl5Bsnn4UA=",
-      "test.json": "GGjXGSZThloU2mCpXKJf9lJVgddpxc8ouDX+JRD5zTs="
     }
   },
-  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
   "webgpu": {
-    "manifestSpec": "2.0",
     "variants": {
       "qkv_present_materialized_sgmat_f32": ["attn-materialized-rowstats-combine-f32.wgsl.jinja", "attn-materialized-sgmat-f32.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "past_kv_materialized_sgmat_f32": ["attn-materialized-rowstats-combine-f32.wgsl.jinja", "attn-materialized-sgmat-f32.wgsl.jinja", "gqa-present.wgsl.jinja"],
@@ -40,9 +40,12 @@
       "new_kv_share_append_split": ["gqa-attention.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "new_kv_share_append_headsink_split": ["gqa-attention.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "new_kv_share_append_rotary_split": ["gqa-attention.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "qkv_present_tiled_nosg": ["attention-rank4-tiled.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "qkv_present_flash": ["attn-flash-online.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "qkv_present": ["attn-online-scalar.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "quant_int8": ["gqa-attention.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "quant_int4": ["gqa-attention.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "quant_int8_decode_splitk": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja", "gqa-present.wgsl.jinja"],
@@ -50,17 +53,9 @@
       "qkv_present_flash_cluster": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "past_kv_bias_flash_prefill": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "past_kv_qnorm_rotary_flash_prefill": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja", "gqa-qprep.wgsl.jinja"],
-      "quant_int8_decode_splitk_nosg": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja", "gqa-present.wgsl.jinja"],
-      "qkv_present_flash_splitk_nosg": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja", "gqa-present.wgsl.jinja"],
-      "qkv_present_flash_cluster_nosg": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja"],
-      "past_kv_bias_flash_prefill_nosg": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja"],
-      "past_kv_qnorm_rotary_flash_prefill_nosg": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja", "gqa-qprep.wgsl.jinja"],
       "past_kv_decode_splitk": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "new_kv_past_decode_splitk": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "window_shift_decode_splitk": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja", "gqa-present.wgsl.jinja"],
-      "past_kv_decode_splitk_nosg": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja", "gqa-present.wgsl.jinja"],
-      "new_kv_past_decode_splitk_nosg": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja", "gqa-present.wgsl.jinja"],
-      "window_shift_decode_splitk_nosg": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "past_kv_flash_prefill": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "new_kv_past_flash_prefill": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "past_kv_rotary_flash_prefill": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja"],
@@ -68,32 +63,27 @@
       "past_kv_headsink_flash_prefill": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "past_kv_bias_headsink_flash_prefill": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "window_shift_flash_prefill": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja"],
-      "past_kv_flash_prefill_nosg": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja"],
-      "new_kv_past_flash_prefill_nosg": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja"],
-      "past_kv_rotary_flash_prefill_nosg": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja"],
-      "past_kv_softcap_flash_prefill_nosg": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja"],
-      "past_kv_headsink_flash_prefill_nosg": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja"],
-      "past_kv_bias_headsink_flash_prefill_nosg": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja"],
-      "window_shift_flash_prefill_nosg": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "qkv_present_flash_q32_broadcast": ["attn-flash-q32-broadcast.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "qkv_present_flash_q32_shared": ["attn-flash-q32-broadcast.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "past_kv": ["gqa-attention.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "past_kv_rotary": ["gqa-attention.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "past_kv_qnorm_rotary": ["gqa-attention.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "new_kv_past": ["gqa-attention.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "window_shift_append": ["gqa-attention.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "new_kv_qnorm_rotary": ["gqa-attention.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "past_kv_bias": ["gqa-attention.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "past_kv_headsink": ["gqa-attention.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "past_kv_bias_headsink": ["gqa-attention.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "quant_int8_flash_prefill": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "quant_int4_flash_prefill": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja"],
-      "quant_int8_flash_prefill_nosg": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja"],
-      "quant_int4_flash_prefill_nosg": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "share_append_split_decode_splitk": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja", "gqa-present.wgsl.jinja"],
-      "share_append_split_decode_splitk_nosg": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "share_append_split_flash_prefill": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja"],
-      "share_append_split_flash_prefill_nosg": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja"]
     }
   }
 }

 {
   "name": "com.microsoft.GroupQueryAttention",
+  "id": "_com_microsoft_groupqueryattention_webgpu_0c84e3e",
   "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": "nKMUcaMP7mcKykLHoGsn/KnsIaJiJRHDsceE7dKVcwo=",
+      "attn-flash-decode-splitk.wgsl.jinja": "ZD6C2fn5I7xZvvpTlsHdLtai8JBK2Ar+xc0SrLn3SAE=",
+      "attn-flash-online.wgsl.jinja": "QGaUiZyIJuQ9xS00I8cuZWYQI8bGldV+UJ0Mp/5kj6w=",
+      "attn-flash-prefill-cluster.wgsl.jinja": "Q7DUmulv66HOVBMkAannNVhEtOWkD+uhDn5muWo4p+E=",
+      "attn-flash-q32-broadcast.wgsl.jinja": "vogUJggIxK8tRg2tDORhIXZLheesY3Y221iVFFyHVY4=",
+      "attn-materialized-rowstats-combine-f32.wgsl.jinja": "SG+V1o5KhJ1LvkUQL0jcq62C+ETRGpy0XV9tvvOQivk=",
+      "attn-materialized-sgmat-f32.wgsl.jinja": "R0+zNbbnfRKpA+R7lcgY5FPlkmLNzNBJwJe1PmaSOvE=",
+      "attn-online-scalar.wgsl.jinja": "ysj3CuuKeu+MPlnaTwnNvCu3sSf8egjHE/uhE/Essmo=",
+      "bench.json": "v06hV3rz9UPzelx7YbIHmyOWgbj/0Y9ci7iz9mQS6m0=",
+      "gqa-attention.wgsl.jinja": "Wy1oR/7g/EzxQZizt/zGtOh4qiD60SVN4pfLSjmfsDo=",
+      "gqa-present.wgsl.jinja": "r0ELXXscqaVYqTYzA9IYr/xl9jWAe20ZFPScyiw34UQ=",
+      "gqa-qprep.wgsl.jinja": "VI/IHqwEMTpvrhFza3kKRDuLO3Eau91C4aMZRW0nAjg=",
+      "manifest.json": "8m8RL1h3mBW9P8voV+gr/KpJiMheynqdWWhmcrOoi8E=",
+      "test.json": "I771sUbwb/9MXIc27bUlvNl3UsiiFxP3H0/wRMt/B9M="
     }
   },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
   "webgpu": {
+    "manifestSpec": "2.1",
     "variants": {
       "qkv_present_materialized_sgmat_f32": ["attn-materialized-rowstats-combine-f32.wgsl.jinja", "attn-materialized-sgmat-f32.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "past_kv_materialized_sgmat_f32": ["attn-materialized-rowstats-combine-f32.wgsl.jinja", "attn-materialized-sgmat-f32.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "new_kv_share_append_split": ["gqa-attention.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "new_kv_share_append_headsink_split": ["gqa-attention.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "new_kv_share_append_rotary_split": ["gqa-attention.wgsl.jinja", "gqa-present.wgsl.jinja"],
+      "new_kv_share_append_bidirectional_split": ["gqa-attention.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "qkv_present_tiled_nosg": ["attention-rank4-tiled.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "qkv_present_flash": ["attn-flash-online.wgsl.jinja", "gqa-present.wgsl.jinja"],
+      "qkv_present_flash_causal": ["attn-flash-online.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "qkv_present": ["attn-online-scalar.wgsl.jinja", "gqa-present.wgsl.jinja"],
+      "qkv_present_causal": ["attn-online-scalar.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "quant_int8": ["gqa-attention.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "quant_int4": ["gqa-attention.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "quant_int8_decode_splitk": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "qkv_present_flash_cluster": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "past_kv_bias_flash_prefill": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "past_kv_qnorm_rotary_flash_prefill": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja", "gqa-qprep.wgsl.jinja"],
       "past_kv_decode_splitk": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "new_kv_past_decode_splitk": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "window_shift_decode_splitk": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "past_kv_flash_prefill": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "new_kv_past_flash_prefill": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "past_kv_rotary_flash_prefill": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "past_kv_headsink_flash_prefill": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "past_kv_bias_headsink_flash_prefill": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "window_shift_flash_prefill": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "qkv_present_flash_q32_broadcast": ["attn-flash-q32-broadcast.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "qkv_present_flash_q32_shared": ["attn-flash-q32-broadcast.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "past_kv": ["gqa-attention.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "past_kv_rotary": ["gqa-attention.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "past_kv_qnorm_rotary": ["gqa-attention.wgsl.jinja", "gqa-present.wgsl.jinja"],
+      "window_shift_rotary_append": ["gqa-attention.wgsl.jinja", "gqa-present.wgsl.jinja"],
+      "window_shift_rotary_headsink_append": ["gqa-attention.wgsl.jinja", "gqa-present.wgsl.jinja"],
+      "window_shift_rotary_bias_append": ["gqa-attention.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "new_kv_past": ["gqa-attention.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "window_shift_append": ["gqa-attention.wgsl.jinja", "gqa-present.wgsl.jinja"],
+      "window_shift_headsink_append": ["gqa-attention.wgsl.jinja", "gqa-present.wgsl.jinja"],
+      "window_shift_bias_append": ["gqa-attention.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "new_kv_qnorm_rotary": ["gqa-attention.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "past_kv_bias": ["gqa-attention.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "past_kv_headsink": ["gqa-attention.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "past_kv_bias_headsink": ["gqa-attention.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "quant_int8_flash_prefill": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "quant_int4_flash_prefill": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "share_append_split_decode_splitk": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja", "gqa-present.wgsl.jinja"],
       "share_append_split_flash_prefill": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja"],
+      "share_append_bidirectional_flash_prefill": ["attn-flash-prefill-cluster.wgsl.jinja", "gqa-present.wgsl.jinja"]
     }
   }
 }

build/webgpu/test.json CHANGED Viewed

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