Xenova HF Staff commited on 3 days ago

Commit

78d4fcb

verified ·

1 Parent(s): 2fad9e5

sync 91d990483a17

Browse files

Files changed (17) hide show

README.md +82 -29
build/webgpu/attention-rank4-tiled.wgsl.jinja +26 -13
build/webgpu/attn-flash-decode-splitk-merge.wgsl.jinja +8 -7
build/webgpu/attn-flash-decode-splitk.wgsl.jinja +44 -49
build/webgpu/attn-flash-online.wgsl.jinja +30 -48
build/webgpu/attn-flash-prefill-cluster.wgsl.jinja +209 -96
build/webgpu/attn-flash-q32-broadcast.wgsl.jinja +1 -27
build/webgpu/attn-materialized-rowstats-combine-f32.wgsl.jinja +2 -3
build/webgpu/attn-materialized-sgmat-f32.wgsl.jinja +30 -51
build/webgpu/attn-online-scalar.wgsl.jinja +6 -40
build/webgpu/bench.json +6 -13
build/webgpu/gqa-attention.wgsl.jinja +15 -10
build/webgpu/gqa-present.wgsl.jinja +16 -18
build/webgpu/gqa-qprep.wgsl.jinja +10 -17
build/webgpu/manifest.json +0 -0
build/webgpu/metadata.json +88 -18
build/webgpu/test.json +438 -77

README.md CHANGED Viewed

@@ -18,31 +18,31 @@ See the [ONNX Runtime `GroupQueryAttention` contrib-operator spec](https://githu
 ## Inputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `query` | `queryT` | `T` | `3` | — | Query tensor of shape `(batch_size, sequence_length, num_heads * head_size)`. | required |
-| `key` | `keyT` | `T` | `3` | — | Key tensor of shape `(batch_size, kv_sequence_length, kv_num_heads * head_size)`. | required |
-| `value` | `valueT` | `T` | `3` | — | Value tensor of shape `(batch_size, kv_sequence_length, kv_num_heads * head_size)`. | required |
-| `past_key` | `pastKeyT` | `T_CACHE` | `4` | — | Optional cached key state in BNSH format. Its sequence axis is `max_sequence_length` when the past and present buffers are shared, otherwise `past_sequence_length`; int4 stores each signed value as a +8-biased nibble, with the even head coordinate low, packing two values per logical uint8 element and widening each byte to one u32 WebGPU buffer word. | optional |
-| `past_value` | `pastValueT` | `T_CACHE` | `4` | — | Optional cached value state in BNSH format with the same length and packing semantics as `past_key`. | optional |
-| `seqlens_k` | `seqlensKT` | `M` | `1` | — | Int32 tensor of shape `(batch_size)` containing each sample's total sequence length minus one. | required |
-| `total_sequence_length` | `totalSequenceLengthT` | `M` | `1` | — | Length-one int32 tensor containing the maximum total sequence length (past plus new) in the batch. | required |
-| `cos_cache` | `cosCacheT` | `T` | `2` | — | Optional cosine cache for rotary embeddings with shape `(max_sequence_length, head_size / 2)`. | optional |
-| `sin_cache` | `sinCacheT` | `T` | `2` | — | Optional sine cache for rotary embeddings with shape `(max_sequence_length, head_size / 2)`. | optional |
-| `attention_bias` | `attentionBiasT` | `T` | `4` | — | Optional additive term for QK scores with shape `(batch_size or 1, num_heads or 1, sequence_length, total_sequence_length)`; the first two dimensions broadcast. | optional |
-| `head_sink` | `headSinkT` | `T` | `1` | — | Optional per-head smooth factor of shape `(num_heads)` added to the softmax denominator. | optional |
-| `k_scale` | `kScaleT` | `T_KV_SCALE` | `1` | — | Optional float32 key-cache scale: one value for `PER_TENSOR`, or `kv_num_heads * head_size` values for `PER_CHANNEL`. | optional |
-| `v_scale` | `vScaleT` | `T_KV_SCALE` | `1` | — | Optional float32 value-cache scale with the same shape convention as `k_scale`. | optional |
-| `q_norm_weight` | `qNormWeightT` | `T` | `1` | — | Optional per-head RMS-normalization weight of shape `(head_size)` applied to queries before rotary embedding. It must be provided together with `k_norm_weight`. | optional |
-| `k_norm_weight` | `kNormWeightT` | `T` | `1` | — | Optional per-head RMS-normalization weight of shape `(head_size)` applied to keys before rotary embedding. It must be provided together with `q_norm_weight`. | optional |
 ## Outputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `output` | `outputT` | `T` | `3` | same as `query` | Attention output of shape `(batch_size, sequence_length, hidden_size)`. | required |
-| `present_key` | `presentKeyT` | `T_CACHE` | `4` | — | Updated key cache in BNSH format. Its sequence axis is `max_sequence_length` for a shared buffer, otherwise `past_sequence_length + kv_sequence_length`; int4 stores each signed value as a +8-biased nibble, with the even head coordinate low, packing two values per logical uint8 element and widening each byte to one u32 WebGPU buffer word. | required |
-| `present_value` | `presentValueT` | `T_CACHE` | `4` | — | Updated value cache in BNSH format with the same length and packing semantics as `present_key`. | required |
 ## Attributes
@@ -53,16 +53,16 @@ Attributes and default values (overridable per request):
 | `causal` | `1` | Whether to apply a causal mask. Set to 0 for bidirectional attention; `local_window_size` must then be -1. |
 | `do_rotary` | `0` | Set to 1 to apply rotary position embeddings. The default 0 disables them. |
 | `k_quant_type` | `"NONE"` | Key-cache quantization mode: `NONE`, `PER_TENSOR`, or `PER_CHANNEL`. |
 | `local_window_size` | `-1` | Left window size for causal local attention. The default -1 disables local attention, and the value must be -1 when `causal` is 0. |
 | `qk_norm_epsilon` | `0.000001` | Epsilon for the per-head Q/K RMS normalization applied when both normalization weights are provided. |
 | `sliding_window_cache` | `0` | Set to 1 when past/present caches are fixed-size window buffers that evict old tokens from the front. Requires `local_window_size > 0` and enough cache capacity. |
 | `smooth_softmax` | `-1` | Set to 1 to enable the smooth-softmax denominator term. |
 | `softcap` | `0` | Positive softcap applied to attention scores. The default 0 disables soft-capping. |
 | `v_quant_type` | `"NONE"` | Value-cache quantization mode: `NONE`, `PER_TENSOR`, or `PER_CHANNEL`. |
-| `kv_cache_bit_width` | — | Quantized cache bit width, either 8 or 4. Four-bit values are packed two per uint8 element. |
-| `kv_num_heads` | — | Number of key/value attention heads. |
-| `num_heads` | — | Number of query attention heads. |
-| `scale` | — | Optional QK score scale; zero or omission selects `1 / sqrt(head_size)`. |
 ## Type constraints
@@ -73,13 +73,63 @@ Attributes and default values (overridable per request):
 | `T_KV_SCALE` | `float32` |
 | `M` | `int32` |
 ## Device requirements
 Some implementation variants require `subgroup-matrix`, `shader-f16`, and `subgroups`. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.
 ## Files
-- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, 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
@@ -98,16 +148,19 @@ Some implementation variants require `subgroup-matrix`, `shader-f16`, and `subgr
 ## Use with `@huggingface/kernels`
-The loader automatically allocates outputs whose metadata it can derive from the manifest contract and this call.
-The explicit `outputs` entries provide shape and logical dtype metadata for the results listed below:
 - `presentKeyT`
 - `presentValueT`
-Each entry either requests an optional result or supplies metadata that cannot be inferred from the inputs.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

 ## Inputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `queryT` | `query` | `T` | `3` | — | Query tensor of shape `(batch_size, sequence_length, num_heads * head_size)`. | required |
+| `keyT` | `key` | `T` | `3` | — | Key tensor of shape `(batch_size, kv_sequence_length, kv_num_heads * head_size)`. | required |
+| `valueT` | `value` | `T` | `3` | — | Value tensor of shape `(batch_size, kv_sequence_length, kv_num_heads * head_size)`. | required |
+| `pastKeyT` | `past_key` | `T_CACHE` | `4` | — | Optional cached key state in BNSH format. Its sequence axis is `max_sequence_length` when the past and present buffers are shared, otherwise `past_sequence_length`; int4 stores each signed value as a +8-biased nibble, with the even head coordinate low, packing two values per logical uint8 element and widening each byte to one u32 WebGPU buffer word. | optional |
+| `pastValueT` | `past_value` | `T_CACHE` | `4` | — | Optional cached value state in BNSH format with the same length and packing semantics as `past_key`. | optional |
+| `seqlensKT` | `seqlens_k` | `M` | `1` | — | Int32 tensor of shape `(batch_size)` containing each sample's total sequence length minus one. | required |
+| `totalSequenceLengthT` | `total_sequence_length` | `M` | `1` | — | Length-one int32 tensor containing the maximum total sequence length (past plus new) in the batch. | required |
+| `cosCacheT` | `cos_cache` | `T` | `2` | — | Optional cosine cache for rotary embeddings with shape `(max_sequence_length, head_size / 2)`. | optional |
+| `sinCacheT` | `sin_cache` | `T` | `2` | — | Optional sine cache for rotary embeddings with shape `(max_sequence_length, head_size / 2)`. | optional |
+| `attentionBiasT` | `attention_bias` | `T` | `4` | — | Optional additive term for QK scores with shape `(batch_size or 1, num_heads or 1, sequence_length, total_sequence_length)`; the first two dimensions broadcast. | optional |
+| `headSinkT` | `head_sink` | `T` | `1` | — | Optional per-head smooth factor of shape `(num_heads)` added to the softmax denominator. | optional |
+| `kScaleT` | `k_scale` | `T_KV_SCALE` | `1` | — | Optional float32 key-cache scale: one value for `PER_TENSOR`, or `kv_num_heads * head_size` values for `PER_CHANNEL`. | optional |
+| `vScaleT` | `v_scale` | `T_KV_SCALE` | `1` | — | Optional float32 value-cache scale with the same shape convention as `k_scale`. | optional |
+| `qNormWeightT` | `q_norm_weight` | `T` | `1` | — | Optional per-head RMS-normalization weight of shape `(head_size)` applied to queries before rotary embedding. It must be provided together with `k_norm_weight`. | optional |
+| `kNormWeightT` | `k_norm_weight` | `T` | `1` | — | Optional per-head RMS-normalization weight of shape `(head_size)` applied to keys before rotary embedding. It must be provided together with `q_norm_weight`. | optional |
 ## Outputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `outputT` | `output` | `T` | `3` | same as `queryT` | Attention output of shape `(batch_size, sequence_length, hidden_size)`. | required |
+| `presentKeyT` | `present_key` | `T_CACHE` | `4` | — | Updated key cache in BNSH format. Its sequence axis is `max_sequence_length` for a shared buffer, otherwise `past_sequence_length + kv_sequence_length`; int4 stores each signed value as a +8-biased nibble, with the even head coordinate low, packing two values per logical uint8 element and widening each byte to one u32 WebGPU buffer word. | required |
+| `presentValueT` | `present_value` | `T_CACHE` | `4` | — | Updated value cache in BNSH format with the same length and packing semantics as `present_key`. | required |
 ## Attributes
 | `causal` | `1` | Whether to apply a causal mask. Set to 0 for bidirectional attention; `local_window_size` must then be -1. |
 | `do_rotary` | `0` | Set to 1 to apply rotary position embeddings. The default 0 disables them. |
 | `k_quant_type` | `"NONE"` | Key-cache quantization mode: `NONE`, `PER_TENSOR`, or `PER_CHANNEL`. |
+| `kv_cache_bit_width` | — | Quantized cache bit width, either 8 or 4. Four-bit values are packed two per uint8 element. |
+| `kv_num_heads` | — | Number of key/value attention heads. |
 | `local_window_size` | `-1` | Left window size for causal local attention. The default -1 disables local attention, and the value must be -1 when `causal` is 0. |
+| `num_heads` | — | Number of query attention heads. |
 | `qk_norm_epsilon` | `0.000001` | Epsilon for the per-head Q/K RMS normalization applied when both normalization weights are provided. |
+| `scale` | — | Optional QK score scale; zero or omission selects `1 / sqrt(head_size)`. |
 | `sliding_window_cache` | `0` | Set to 1 when past/present caches are fixed-size window buffers that evict old tokens from the front. Requires `local_window_size > 0` and enough cache capacity. |
 | `smooth_softmax` | `-1` | Set to 1 to enable the smooth-softmax denominator term. |
 | `softcap` | `0` | Positive softcap applied to attention scores. The default 0 disables soft-capping. |
 | `v_quant_type` | `"NONE"` | Value-cache quantization mode: `NONE`, `PER_TENSOR`, or `PER_CHANNEL`. |
 ## Type constraints
 | `T_KV_SCALE` | `float32` |
 | `M` | `int32` |
+## Implementation variants
+One implementation is selected per call from the device capabilities, the request shapes and the dtypes; these notes say what each one covers.
+- `qkv_present_materialized_sgmat_f32` — Materialized float32 subgroup-matrix prefill for bidirectional QKV without past state. `seqlens_k` is metadata-only; the score pass emits per-row softmax statistics, the apply pass normalizes and applies the scores, and the present copy preserves the supplied key and value tensors.
+- `past_kv_materialized_sgmat_f32` — Materializes the causal score matrix with float32 subgroup-matrix tiles over the shared float cache, then applies the softmax-normalized weights with subgroup-matrix tiles. Both passes skip key tiles beyond each query tile's causal bound.
+- `past_kv_bias_materialized_sgmat_f32` — Materializes causal scores with float32 subgroup-matrix tiles, broadcasts and adds the attention bias across batch and head dimensions, folds the biased scores into row statistics, and applies the normalized weights with subgroup-matrix tiles.
+- `past_kv_materialized_sgmat_f16` — Materializes the causal score matrix with `f16` operand tiles feeding `f32`-accumulating subgroup matrices over the shared cache and applies the softmax-normalized weights the same way; scores and row statistics stay `f32`. The score pass skips key tiles past each query tile's causal bound and the apply pass stops its reduction there.
+- `past_kv_rotary_materialized_sgmat_f32` — Rotates each query block at its absolute positions, then materializes causal scores against rotary-transformed cached keys and applies the normalized weights. The score and apply passes skip tiles beyond each query tile's causal bound.
+- `past_kv_rotary_materialized_sgmat_f16` — Materializes the causal score matrix with `f16` operand tiles feeding `f32`-accumulating subgroup matrices over the shared cache and applies the softmax-normalized weights the same way; scores and row statistics stay `f32`. The score pass skips key tiles past each query tile's causal bound and the apply pass stops its reduction there.
+- `new_kv_past_materialized_sgmat_f32` — For chunked prefill, merges the past cache with new key/value rows, materializes causal scores over the merged cache, and applies the normalized weights using the same right-aligned causal bound.
+- `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.
+- `past_kv_softcap_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.
+- `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
 Some implementation variants require `subgroup-matrix`, `shader-f16`, and `subgroups`. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.
 ## Files
+- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
 - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 ## 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.
+This example supplies explicit metadata for:
 - `presentKeyT`
 - `presentValueT`
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

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

@@ -1,19 +1,15 @@
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 // Thread-per-query online-softmax flash fallback for rank-4 (BNSH) attention with
 // no optional feature requirements. Each invocation owns one full query row and walks the keys
 // independently, avoiding per-key workgroup barriers and tree reductions.
 //
-// BLOCK_M consecutive queries of one (batch, q_head) are tiled into a workgroup so
-// they stream the same K/V rows at the same key index, so those global loads hit the
-// L2 cache (read once per head, broadcast across the tile). The per-query online
 // accumulator lives in shared memory laid out d-major (d*BLOCK_M + tid) so the
-// BLOCK_M threads touch consecutive addresses for a fixed d (no bank conflict) and
-// nothing spills to registers. The value-head width is capped so acc fits in
-// the workgroup-storage budget;
 // larger heads use the scalar online-attention fallback.
 const BLOCK_M: u32 = {{ blockM }}u;
@@ -34,13 +30,12 @@ fn kv_head(q_head: u32) -> u32 {
 @compute @workgroup_size(BLOCK_M, 1, 1)
 fn main(
   @builtin(workgroup_id) wg: vec3<u32>,
-  @builtin(num_workgroups) nwg: vec3<u32>,
   @builtin(local_invocation_id) lid: vec3<u32>
 ) {
   let tid = lid.x;
   // 2D-folded tile index over (batch, q_head, q_seq tile); wg.y carries the high
-  // bits past the maxComputeWorkgroupsPerDimension dispatch limit.
-  let tile = wg.x + wg.y * nwg.x;
   let qTiles = (params.qSeq + BLOCK_M - 1u) / BLOCK_M;
   // count = numel(Y) = batch·qHeads·qSeq·vHeadSize, so count/(vHeadSize·qSeq) =
   // batch·qHeads. Drop the over-dispatched tail from the 2D fold.
@@ -56,8 +51,14 @@ fn main(
   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) {
@@ -82,10 +83,14 @@ fn main(
     // 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; }
-    // The K/V row is shared across the tile, so these loads broadcast from L2.
     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]);
@@ -105,14 +110,22 @@ fn main(
     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);
   }

+{% 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
 // no optional feature requirements. Each invocation owns one full query row and walks the keys
 // independently, avoiding per-key workgroup barriers and tree reductions.
 //
+// BLOCK_M consecutive queries of one (batch, q_head) are tiled into a workgroup,
+// so adjacent query threads reuse the same K/V row at each key index. The per-query online
 // accumulator lives in shared memory laid out d-major (d*BLOCK_M + tid) so the
+// BLOCK_M threads touch consecutive addresses for a fixed d. The value-head
+// width is capped so the accumulator fits in the workgroup-storage budget;
 // larger heads use the scalar online-attention fallback.
 const BLOCK_M: u32 = {{ blockM }}u;
 @compute @workgroup_size(BLOCK_M, 1, 1)
 fn main(
   @builtin(workgroup_id) wg: vec3<u32>,
   @builtin(local_invocation_id) lid: vec3<u32>
 ) {
   let tid = lid.x;
   // 2D-folded tile index over (batch, q_head, q_seq tile); wg.y carries the high
+  // bits past the per-axis dispatch fold width.
+  let tile = wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u;
   let qTiles = (params.qSeq + BLOCK_M - 1u) / BLOCK_M;
   // count = numel(Y) = batch·qHeads·qSeq·vHeadSize, so count/(vHeadSize·qSeq) =
   // batch·qHeads. Drop the over-dispatched tail from the 2D fold.
   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) {
     // 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]);
     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);
   }

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

@@ -3,6 +3,7 @@ enable f16;
 {% endif %}
 {% if splitQueries is not defined %}{% set splitQueries = false %}{% endif %}
 {% if hasGate is not defined %}{% set hasGate = false %}{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 // Split-K flash decode, pass 2 of 2. Combines the per-split un-normalized online
@@ -14,10 +15,10 @@ enable f16;
 // One workgroup per (batch, head); one thread per output vec4 channel. The
 // per-split (max, denom) scan is tiny (NUM_SPLITS is small), so every thread
 // recomputes it locally and keeps the merge state in registers.
-{% if source.layout == "bhsd" %}
 // Output layout: rank-4 [batch, heads, seq, headDim], matching the split pass.
-{% elif source.layout == "layer_cache" %}
-// Output layout: flat [heads, headDim], matching Qwen3.5 decode.
 {% else %}
 // Output layout: token-major [batch, seq, heads * headDim], matching the split pass.
 {% endif %}
@@ -27,7 +28,7 @@ const NUM_SPLITS: u32 = {{ numSplits }}u;
 {% if splitQueries %}
 const Q_SEQ: u32 = {{ qSeq }}u;
 {% endif %}
-{% if source.layout == "bsh" %}
 const Q_HIDDEN_V4: u32 = {{ qHiddenV4 }}u;
 {% endif %}
 {% if hasBias %}
@@ -97,13 +98,13 @@ fn main(
     outV = outV + partial_out[pBase + d4] * w;
   }
-{% if source.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 source.layout == "layer_cache" %}
   let qBaseV4 = h * HEAD_DIM_V4;
 {% else %}
 {% if splitQueries %}
@@ -123,7 +124,7 @@ fn main(
   outValue = outValue + vec4<f32>(bias[vBiasBase], bias[vBiasBase + 1u], bias[vBiasBase + 2u], bias[vBiasBase + 3u]);
 {% endif %}
 {% if hasGate %}
-  // Qwen3.5 full attention gates the normalized attention output elementwise.
   let gateV = vec4<f32>(gate[qBaseV4 + d4]);
   outValue = outValue * (vec4<f32>(1.0) / (vec4<f32>(1.0) + exp(-gateV)));
 {% endif %}

 {% endif %}
 {% if splitQueries is not defined %}{% set splitQueries = false %}{% endif %}
 {% if hasGate is not defined %}{% set hasGate = false %}{% endif %}
+{% set qSeq = qSeq | default(0) %}
 {{ env.wgsl.resourceDeclarations }}
 // Split-K flash decode, pass 2 of 2. Combines the per-split un-normalized online
 // One workgroup per (batch, head); one thread per output vec4 channel. The
 // per-split (max, denom) scan is tiny (NUM_SPLITS is small), so every thread
 // recomputes it locally and keeps the merge state in registers.
+{% if layout == "bhsd" %}
 // Output layout: rank-4 [batch, heads, seq, headDim], matching the split pass.
+{% elif layout == "layer_cache" %}
+// Output layout: flat [heads, headDim], matching the cached-decode split pass.
 {% else %}
 // Output layout: token-major [batch, seq, heads * headDim], matching the split pass.
 {% endif %}
 {% if splitQueries %}
 const Q_SEQ: u32 = {{ qSeq }}u;
 {% endif %}
+{% if layout == "bsh" %}
 const Q_HIDDEN_V4: u32 = {{ qHiddenV4 }}u;
 {% endif %}
 {% if hasBias %}
     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 %}
   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 %}

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

@@ -4,17 +4,19 @@
 {% 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 splitKWorkgroupSize = source.workgroupSize if source.workgroupSize is defined else tunables.WORKGROUP_SIZE %}
 {% if useSubgroups %}
 enable subgroups;
 {% endif %}
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
-// Split-K flash attention, pass 1 of 2; the merge pass follows. Shared by
-// dense-attention decode and short-query/long-context prefill paths.
 //
 // The non-split flash decode launches only `batch * numHeads` workgroups, each
 // sweeping the whole KV sequence serially in WG-key tiles. This pass splits the
@@ -24,25 +26,24 @@ enable f16;
 // its range — the running (max, denom) and the softmax-weighted V sum before the
 // final divide — and the merge pass combines the per-split states with the online
 // rule.
-{% if source.layout == "bhsd" %}
 // Layout: rank-4 [batch, heads, seq, headDim] for Q/K/V.
-{% elif source.layout == "layer_cache" %}
 // Layout: flat query [heads, headDim] plus a persistent KV cache laid out
-// [layer, cacheLen, kvHeads, headDim]. This is the Qwen3.5 decode layout; the
-// dispatch has a single implicit batch.
 {% else %}
 // Layout: token-major [batch, seq, heads * headDim]; Q and KV hidden strides
 // are compiled constants.
 {% endif %}
-{% if (fusedQNormRope is defined and fusedQNormRope) or source.layout != "layer_cache" %}const HEAD_DIM: u32 = {{ headDim }}u;
 {% endif %}
 const HEAD_DIM_V4: u32 = {{ headDimV4 }}u;
 const Q_HEADS: u32 = {{ qNumHeads }}u;
 const KV_HEADS: u32 = {{ kvNumHeads }}u;
-{% if source.layout == "bsh" %}
 const Q_HIDDEN_V4: u32 = {{ qHiddenV4 }}u;
 const KV_HIDDEN_V4: u32 = {{ kvHiddenV4 }}u;
-{% elif source.layout == "layer_cache" %}
 const LAYER: u32 = {{ layer }}u;
 const CACHE_LEN: u32 = {{ cacheLen }}u;
 const ATTN_SCALE: f32 = {{ scale }};
@@ -96,20 +97,22 @@ var<workgroup> vacc_sh: array<vec4<f32>, WG>;
 // calls before their shared partial storage is reused.
 {% set combineSubgroups = combineSubgroups is defined and combineSubgroups %}
 {% if combineSubgroups %}
-// Per-subgroup partials are published into a deterministic slot: the subgroup's
-// ordinal index within the workgroup (lidx / sgSize). The online (m, d) merge
-// is not float-associative, so thread 0 must fold partials in a fixed order.
-// Subgroups partition a workgroup into contiguous ordinal ranges on supported
-// backends, so the ordinal slot is unique per subgroup and every slot in
-// [0, subgroupCount) is written (each subgroup elects one leader).
-// Sized for the worst case of one partial per invocation.
 var<workgroup> partialM: array<f32, WG>;
 var<workgroup> partialD: array<f32, WG>;
 var<workgroup> combinedMD: vec2<f32>;
 // When the whole workgroup is one subgroup the subgroup reduce already covers
-// it (no barriers, no shared state); otherwise subgroup leaders publish
-// partials through shared memory and thread 0 folds them in ordinal order.
 fn combine_partials(m: f32, d: f32, lidx: u32, sgSize: u32) -> vec2<f32> {
   let sgM = subgroupMax(m);
   // A lane with no elements contributes d == 0 (exact identity). A +inf
@@ -120,31 +123,24 @@ fn combine_partials(m: f32, d: f32, lidx: u32, sgSize: u32) -> vec2<f32> {
   if (sgSize == WG) {
     return vec2<f32>(sgM, sgD);
   }
-  let subgroupCount = (WG + sgSize - 1u) / sgSize;
-  // Pre-seed every fold slot with the (max, denom) identity. The fold below reads a
-  // fixed subgroupCount slots in ordinal order (for determinism), but a slot whose
-  // subgroup elects no leader this call — e.g. a fully out-of-window key tile in the
-  // flash-attention loop that re-uses this shared memory each iteration — would
-  // otherwise read stale shared memory. Identity makes such a slot a no-op.
-  // (max identity = -FLT_MAX, denom identity = 0.)
-  if (lidx < subgroupCount) {
-    partialM[lidx] = -FLT_MAX;
-    partialD[lidx] = 0.0;
-  }
-  workgroupBarrier();
   if (subgroupElect()) {
-    let slot = lidx / sgSize;
-    partialM[slot] = sgM;
-    partialD[slot] = sgD;
   }
   workgroupBarrier();
   if (lidx == 0u) {
     var accM = -FLT_MAX;
     var accD = 0.0;
-    for (var i = 0u; i < subgroupCount; i = i + 1u) {
-      let mNew = max(accM, partialM[i]);
-      accD = accD * exp_shift(accM, mNew) + partialD[i] * exp_shift(partialM[i], mNew);
-      accM = mNew;
     }
     combinedMD = vec2<f32>(accM, accD);
   }
@@ -159,9 +155,8 @@ var<workgroup> partialM: array<f32, {{ mdExtent }}>;
 var<workgroup> partialD: array<f32, {{ mdExtent }}>;
 {% if mdStreamed %}
-// In-place fold of {{ mdStreams }} streams. The caller stores its per-thread
-// partials into partialM/partialD first and reads the merged pair of stream s
-// from slot s * WG afterwards.
 fn combine_partials_streams(lidx: u32) {
   workgroupBarrier();
   var stride = WG / 2u;
@@ -220,7 +215,7 @@ fn combine_partials(m: f32, d: f32, lidx: u32) -> vec2<f32> {
 {% endif %}
-{% if source.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 %}
@@ -310,12 +305,12 @@ fn main(
 {% endif %}
   let h = wg.y;
   let b = wg.z;
-  if (h >= Q_HEADS || split >= NUM_SPLITS{% if splitQueries %} || queryToken >= Q_SEQ{% endif %}{% if source.layout == "layer_cache" %} || params.past_len >= CACHE_LEN{% endif %}) {
     return;
   }
   let tid = lid.x;
   let hKv = h / (Q_HEADS / KV_HEADS);
-{% if source.layout == "layer_cache" %}
   let kvSeq = params.past_len + 1u;
 {% else %}
   let cacheSeq = params.kvSeq;
@@ -329,7 +324,7 @@ fn main(
 {% endif %}
   // Query row (decode uses token zero; short-query prefill folds the token into wg.x).
-{% if source.layout == "bsh" %}
 {% if splitQueries %}
   let qBaseV4 = (b * Q_SEQ + queryToken) * Q_HIDDEN_V4 + h * HEAD_DIM_V4;
 {% else %}
@@ -337,7 +332,7 @@ fn main(
 {% endif %}
   let kvBaseV4 = b * kvSeq * KV_HIDDEN_V4 + hKv * HEAD_DIM_V4;
   let kvTokenStrideV4 = KV_HIDDEN_V4;
-{% elif source.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;

 {% 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
 // sweeping the whole KV sequence serially in WG-key tiles. This pass splits the
 // its range — the running (max, denom) and the softmax-weighted V sum before the
 // final divide — and the merge pass combines the per-split states with the online
 // rule.
+{% if layout == "bhsd" %}
 // Layout: rank-4 [batch, heads, seq, headDim] for Q/K/V.
+{% elif layout == "layer_cache" %}
 // Layout: flat query [heads, headDim] plus a persistent KV cache laid out
+// [layer, cacheLen, kvHeads, headDim]. The dispatch has a single implicit batch.
 {% else %}
 // Layout: token-major [batch, seq, heads * headDim]; Q and KV hidden strides
 // are compiled constants.
 {% endif %}
+{% if (fusedQNormRope is defined and fusedQNormRope) or layout != "layer_cache" %}const HEAD_DIM: u32 = {{ headDim }}u;
 {% endif %}
 const HEAD_DIM_V4: u32 = {{ headDimV4 }}u;
 const Q_HEADS: u32 = {{ qNumHeads }}u;
 const KV_HEADS: u32 = {{ kvNumHeads }}u;
+{% if layout == "bsh" %}
 const Q_HIDDEN_V4: u32 = {{ qHiddenV4 }}u;
 const KV_HIDDEN_V4: u32 = {{ kvHiddenV4 }}u;
+{% elif layout == "layer_cache" %}
 const LAYER: u32 = {{ layer }}u;
 const CACHE_LEN: u32 = {{ cacheLen }}u;
 const ATTN_SCALE: f32 = {{ scale }};
 // calls before their shared partial storage is reused.
 {% set combineSubgroups = combineSubgroups is defined and combineSubgroups %}
 {% if combineSubgroups %}
+// Cross-subgroup merge that assumes nothing about which invocations share a
+// subgroup or how many subgroups there are: each subgroup's elected lane
+// publishes the subgroup pair in the slot at its OWN invocation index and sets
+// that index's bit in a workgroup bitmask; thread 0 then folds exactly the
+// published slots, in ascending index order (the online (m, d) merge is not
+// float-associative, so the order is fixed), and clears the mask for the next
+// call as it reads it. Workgroup memory starts zeroed, so the mask needs no
+// setup. Same three collectives as a single-subgroup reduce, two barriers.
 var<workgroup> partialM: array<f32, WG>;
 var<workgroup> partialD: array<f32, WG>;
+var<workgroup> leaderMask: array<atomic<u32>, (WG + 31u) / 32u>;
 var<workgroup> combinedMD: vec2<f32>;
 // When the whole workgroup is one subgroup the subgroup reduce already covers
+// it (no barriers, no shared state). `subgroup_size` is the size of the current
+// subgroup and uniform, so the test is exact and may guard the barriers below.
 fn combine_partials(m: f32, d: f32, lidx: u32, sgSize: u32) -> vec2<f32> {
   let sgM = subgroupMax(m);
   // A lane with no elements contributes d == 0 (exact identity). A +inf
   if (sgSize == WG) {
     return vec2<f32>(sgM, sgD);
   }
   if (subgroupElect()) {
+    partialM[lidx] = sgM;
+    partialD[lidx] = sgD;
+    atomicOr(&leaderMask[lidx / 32u], 1u << (lidx % 32u));
   }
   workgroupBarrier();
   if (lidx == 0u) {
     var accM = -FLT_MAX;
     var accD = 0.0;
+    for (var w = 0u; w < (WG + 31u) / 32u; w = w + 1u) {
+      var bits = atomicExchange(&leaderMask[w], 0u);
+      while (bits != 0u) {
+        let slot = w * 32u + firstTrailingBit(bits);
+        bits = bits & (bits - 1u);
+        let mNew = max(accM, partialM[slot]);
+        accD = accD * exp_shift(accM, mNew) + partialD[slot] * exp_shift(partialM[slot], mNew);
+        accM = mNew;
+      }
     }
     combinedMD = vec2<f32>(accM, accD);
   }
 var<workgroup> partialD: array<f32, {{ mdExtent }}>;
 {% if mdStreamed %}
+// In-place fold of {{ mdStreams }} streams. Input partials occupy
+// partialM/partialD; stream s returns its merged pair in slot s * WG.
 fn combine_partials_streams(lidx: u32) {
   workgroupBarrier();
   var stride = WG / 2u;
 {% 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 %}
 {% 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;
   let hKv = h / (Q_HEADS / KV_HEADS);
+{% if layout == "layer_cache" %}
   let kvSeq = params.past_len + 1u;
 {% else %}
   let cacheSeq = params.kvSeq;
 {% endif %}
   // 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 %}
 {% 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;

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

@@ -1,13 +1,10 @@
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {% if combineSubgroups %}
 enable subgroups;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
-// Flash-style tiled online-softmax attention, shared by every dense
-// attention op with vec4-aligned head dims. One workgroup per
 // (batch, head, query token); the workgroup sweeps the KV sequence in tiles
 // of WG keys, each thread owning one key of the tile (full q·k dot in
 // registers, vec4 loads), so a tile needs a single fused (max, sum-exp)
@@ -16,9 +13,6 @@ enable subgroups;
 // per tile; the full score matrix is never materialized.
 // Layout: token-major [batch, seq, heads * headDim]; Q and KV hidden strides
 // are compiled constants.
-{% if headsFromParams %}
-// Q/KV head counts are uniforms; the GQA mapping is h / (qHeads / kvHeads).
-{% endif %}
 {% if hasBias %}
 // Packed [Q; K; V] bias rows. The K bias adds the same dot(q, biasK) to every
 // key score, which softmax cancels, so only the Q and V biases are applied.
@@ -32,10 +26,8 @@ const HEAD_DIM: u32 = {{ headDim }}u;
 const HEAD_DIM_V4: u32 = {{ headDimV4 }}u;
 const Q_HIDDEN_V4: u32 = {{ qHiddenV4 }}u;
 const KV_HIDDEN_V4: u32 = {{ kvHiddenV4 }}u;
-{% if not headsFromParams %}
 const Q_HEADS: u32 = {{ qNumHeads }}u;
 const KV_HEADS: u32 = {{ kvNumHeads }}u;
-{% endif %}
 {% 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;
@@ -73,20 +65,22 @@ var<workgroup> probs: array<f32, WG>;
 // calls before their shared partial storage is reused.
 {% set combineSubgroups = combineSubgroups is defined and combineSubgroups %}
 {% if combineSubgroups %}
-// Per-subgroup partials are published into a deterministic slot: the subgroup's
-// ordinal index within the workgroup (lidx / sgSize). The online (m, d) merge
-// is not float-associative, so thread 0 must fold partials in a fixed order.
-// Subgroups partition a workgroup into contiguous ordinal ranges on supported
-// backends, so the ordinal slot is unique per subgroup and every slot in
-// [0, subgroupCount) is written (each subgroup elects one leader).
-// Sized for the worst case of one partial per invocation.
 var<workgroup> partialM: array<f32, WG>;
 var<workgroup> partialD: array<f32, WG>;
 var<workgroup> combinedMD: vec2<f32>;
 // When the whole workgroup is one subgroup the subgroup reduce already covers
-// it (no barriers, no shared state); otherwise subgroup leaders publish
-// partials through shared memory and thread 0 folds them in ordinal order.
 fn combine_partials(m: f32, d: f32, lidx: u32, sgSize: u32) -> vec2<f32> {
   let sgM = subgroupMax(m);
   // A lane with no elements contributes d == 0 (exact identity). A +inf
@@ -97,31 +91,24 @@ fn combine_partials(m: f32, d: f32, lidx: u32, sgSize: u32) -> vec2<f32> {
   if (sgSize == WG) {
     return vec2<f32>(sgM, sgD);
   }
-  let subgroupCount = (WG + sgSize - 1u) / sgSize;
-  // Pre-seed every fold slot with the (max, denom) identity. The fold below reads a
-  // fixed subgroupCount slots in ordinal order (for determinism), but a slot whose
-  // subgroup elects no leader this call — e.g. a fully out-of-window key tile in the
-  // flash-attention loop that re-uses this shared memory each iteration — would
-  // otherwise read stale shared memory. Identity makes such a slot a no-op.
-  // (max identity = -FLT_MAX, denom identity = 0.)
-  if (lidx < subgroupCount) {
-    partialM[lidx] = -FLT_MAX;
-    partialD[lidx] = 0.0;
-  }
-  workgroupBarrier();
   if (subgroupElect()) {
-    let slot = lidx / sgSize;
-    partialM[slot] = sgM;
-    partialD[slot] = sgD;
   }
   workgroupBarrier();
   if (lidx == 0u) {
     var accM = -FLT_MAX;
     var accD = 0.0;
-    for (var i = 0u; i < subgroupCount; i = i + 1u) {
-      let mNew = max(accM, partialM[i]);
-      accD = accD * exp_shift(accM, mNew) + partialD[i] * exp_shift(partialM[i], mNew);
-      accM = mNew;
     }
     combinedMD = vec2<f32>(accM, accD);
   }
@@ -136,9 +123,8 @@ var<workgroup> partialM: array<f32, {{ mdExtent }}>;
 var<workgroup> partialD: array<f32, {{ mdExtent }}>;
 {% if mdStreamed %}
-// In-place fold of {{ mdStreams }} streams. The caller stores its per-thread
-// partials into partialM/partialD first and reads the merged pair of stream s
-// from slot s * WG afterwards.
 fn combine_partials_streams(lidx: u32) {
   workgroupBarrier();
   var stride = WG / 2u;
@@ -197,8 +183,8 @@ fn combine_partials(m: f32, d: f32, lidx: u32) -> vec2<f32> {
 {% endif %}
-// Omitted ONNX scale uses 1/sqrt(headDim). Explicit zero is handled by
-// specialization because zero is the runtime omitted-value sentinel.
 {% 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; }
@@ -257,11 +243,7 @@ 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.
-{% 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
@@ -287,7 +269,7 @@ fn main(
     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 %}

 {% if combineSubgroups %}
 enable subgroups;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
+// Flash-style tiled online-softmax attention for vec4-aligned head dimensions.
+// One workgroup per
 // (batch, head, query token); the workgroup sweeps the KV sequence in tiles
 // of WG keys, each thread owning one key of the tile (full q·k dot in
 // registers, vec4 loads), so a tile needs a single fused (max, sum-exp)
 // per tile; the full score matrix is never materialized.
 // Layout: token-major [batch, seq, heads * headDim]; Q and KV hidden strides
 // are compiled constants.
 {% if hasBias %}
 // Packed [Q; K; V] bias rows. The K bias adds the same dot(q, biasK) to every
 // key score, which softmax cancels, so only the Q and V biases are applied.
 const HEAD_DIM_V4: u32 = {{ headDimV4 }}u;
 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;
 // calls before their shared partial storage is reused.
 {% set combineSubgroups = combineSubgroups is defined and combineSubgroups %}
 {% if combineSubgroups %}
+// Cross-subgroup merge that assumes nothing about which invocations share a
+// subgroup or how many subgroups there are: each subgroup's elected lane
+// publishes the subgroup pair in the slot at its OWN invocation index and sets
+// that index's bit in a workgroup bitmask; thread 0 then folds exactly the
+// published slots, in ascending index order (the online (m, d) merge is not
+// float-associative, so the order is fixed), and clears the mask for the next
+// call as it reads it. Workgroup memory starts zeroed, so the mask needs no
+// setup. Same three collectives as a single-subgroup reduce, two barriers.
 var<workgroup> partialM: array<f32, WG>;
 var<workgroup> partialD: array<f32, WG>;
+var<workgroup> leaderMask: array<atomic<u32>, (WG + 31u) / 32u>;
 var<workgroup> combinedMD: vec2<f32>;
 // When the whole workgroup is one subgroup the subgroup reduce already covers
+// it (no barriers, no shared state). `subgroup_size` is the size of the current
+// subgroup and uniform, so the test is exact and may guard the barriers below.
 fn combine_partials(m: f32, d: f32, lidx: u32, sgSize: u32) -> vec2<f32> {
   let sgM = subgroupMax(m);
   // A lane with no elements contributes d == 0 (exact identity). A +inf
   if (sgSize == WG) {
     return vec2<f32>(sgM, sgD);
   }
   if (subgroupElect()) {
+    partialM[lidx] = sgM;
+    partialD[lidx] = sgD;
+    atomicOr(&leaderMask[lidx / 32u], 1u << (lidx % 32u));
   }
   workgroupBarrier();
   if (lidx == 0u) {
     var accM = -FLT_MAX;
     var accD = 0.0;
+    for (var w = 0u; w < (WG + 31u) / 32u; w = w + 1u) {
+      var bits = atomicExchange(&leaderMask[w], 0u);
+      while (bits != 0u) {
+        let slot = w * 32u + firstTrailingBit(bits);
+        bits = bits & (bits - 1u);
+        let mNew = max(accM, partialM[slot]);
+        accD = accD * exp_shift(accM, mNew) + partialD[slot] * exp_shift(partialM[slot], mNew);
+        accM = mNew;
+      }
     }
     combinedMD = vec2<f32>(accM, accD);
   }
 var<workgroup> partialD: array<f32, {{ mdExtent }}>;
 {% if mdStreamed %}
+// In-place fold of {{ mdStreams }} streams. Input partials occupy
+// partialM/partialD; stream s returns its merged pair in slot s * WG.
 fn combine_partials_streams(lidx: u32) {
   workgroupBarrier();
   var stride = WG / 2u;
 {% 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; }
   // 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
     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 %}

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

@@ -1,4 +1,6 @@
-{% set sourceProfile = source.sourceProfile if source.sourceProfile is defined else 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" %}
@@ -7,6 +9,7 @@
 {% 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 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 %}
@@ -14,7 +17,6 @@
 {% if maskIsBool is not defined %}{% set maskIsBool = false %}{% endif %}
 {% if hasSoftcap is not defined %}{% set hasSoftcap = false %}{% endif %}
 {% if hasHeadSink is not defined %}{% set hasHeadSink = false %}{% endif %}
-{% set Q_HIDDEN = qHidden | default(0) %}
 {% if quantCacheFormat is not defined %}{% set quantCacheFormat = "" %}{% endif %}
 {% if useSeqlens is not defined %}{% set useSeqlens = false %}{% endif %}
 // A windowed cache binds a fixed CAPACITY but keeps only the most recent
@@ -27,9 +29,6 @@
 {% if useSubgroups %}
 enable subgroups;
 {% endif %}
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 {% set ST = "f16" if usesF16 else "f32" %}
 // K/V tiles are staged as f16 when this specialization uses f16, reducing
@@ -44,16 +43,18 @@ enable f16;
 {% 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 ROPE_LANE_XOR = ((LPQ / 2) | int) %}
-{% set QL = source.qLayout if source.qLayout is defined else source.layout %}
-{% set KL = source.kvLayout if source.kvLayout is defined else source.layout %}
 {% set RIGHT = causalRightAlign is defined and causalRightAlign %}
 // Tiled flash prefill attention with configurable-width query clusters for
 // token-major [batch, seq, heads*headDim] attention ops. Each workgroup covers
 // TILE_Q queries with one LPQ-lane cluster per query; K/V are staged in
-// workgroup memory once per TILE_K tile. Within a cluster, each lane holds q/o
-// register slices of HEAD_DIM/LPQ dimensions. Per-lane dot fragments are
 // combined with subgroup shuffles when available, or shared-memory reductions
 // on compatibility variants. Dot products, online softmax, and weighted-value
 // sums accumulate in f32. Optional masks are broadcast-strided; the
@@ -68,7 +69,8 @@ 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
-{% endif %}
 {% if KL != "bhsd" %}
 const KV_HIDDEN_V4: u32 = {{ kvHiddenV4 }}u;
 {% endif %}
@@ -83,7 +85,12 @@ const HALF_DIM: u32 = HEAD_DIM / 2u;
 const HALF_LPQ: u32 = LPQ / 2u;
 {% endif %}
 const TILE_K: u32 = {{ TILE_K }}u;
 const WG: u32 = TILE_Q * LPQ;
 {% 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
@@ -116,8 +123,12 @@ var<workgroup> red: array<f32, 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 }}));
 }
 {% if quantCacheFormat %}
@@ -162,13 +173,6 @@ fn load_{{ kind }}4(indexV4: u32, d4: u32, hk: u32) -> vec4<f32> {
 {{ emit_quant_load4(quantCacheFormat, "value", VALUE, "v_scale") }}
 {% endif %}
-{% if hasBias %}
-fn load_bias4(base: u32, d4: u32) -> vec4<f32> {
-  let offset = base + d4 * 4u;
-  return vec4<f32>(bias[offset], bias[offset + 1u], bias[offset + 2u], bias[offset + 3u]);
-}
-{% endif %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(
   @builtin(workgroup_id) wg: vec3<u32>,
@@ -184,74 +188,85 @@ fn main(
   // portability.
   if (sgSize < LPQ || sgSize % LPQ != 0u) { return; }
 {% endif %}
   let qSub = tid / LPQ;
   let lane8 = tid % LPQ;
-  let qIdx = wg.x * TILE_Q + qSub;
-  let qValid = qIdx < {{ QSEQ }} && h < Q_HEADS;
   let hKv = h / (Q_HEADS / KV_HEADS);
   let SCALE = scale_value();
-  // Per-thread q slice + output accumulator (SLICE vec4s each) in registers.
-  let qClamped = min(qIdx, {{ QSEQ }} - 1u);
 {% if QL == "bhsd" %}
-  let qBase4 = ((b * Q_HEADS + h) * {{ QSEQ }} + qClamped) * HEAD_DIM_V4 + lane8 * SLICE;
 {% else %}
-  let qBase4 = (b * {{ QSEQ }} + qClamped) * {{ Q_STRIDE }} + h * HEAD_DIM_V4 + lane8 * SLICE;
 {% endif %}
 {% for c in range(SLICE_COUNT) %}
-  var qr{{ c }} = vec4<f32>({{ QUERY }}[qBase4 + {{ c }}u]);
-{% if hasBias %}
-  qr{{ c }} = qr{{ c }} + load_bias4(h * HEAD_DIM, lane8 * SLICE + {{ c }}u);
-{% endif %}
-  var o{{ c }} = vec4<f32>(0.0);
 {% endfor %}
 {% if FUSED_ROTARY %}
-  // The present-cache keys are already rotary transformed. Transform Q once in
-  // registers before the K loop, avoiding a separate full-size Q scratch pass.
-  // NeoX half-split pairs lanes separated by LPQ/2; subgroup variants exchange
-  // the paired register directly, while the portable path reloads that vec4.
 {% if not useSubgroups %}
   let qHeadBase4 = qBase4 - lane8 * SLICE;
   let pairedLane = (lane8 + HALF_LPQ) % LPQ;
 {% endif %}
   let ropeChunk = (lane8 % HALF_LPQ) * SLICE;
   let pastLenForRope = select(0u, {{ KVSEQ }} - {{ QSEQ }}, {{ KVSEQ }} >= {{ QSEQ }});
-  let ropePosition = pastLenForRope + qClamped;
 {% for c in range(SLICE_COUNT) %}
 {% if useSubgroups %}
-  let pairedQ{{ c }} = subgroupShuffleXor(qr{{ c }}, {{ ROPE_LANE_XOR }}u);
 {% else %}
-  let pairedQ{{ c }} = vec4<f32>({{ QUERY }}[qHeadBase4 + pairedLane * SLICE + {{ c }}u]);
-{% endif %}
-  let ropeBase{{ c }} = ropePosition * HALF_DIM + (ropeChunk + {{ c }}u) * 4u;
-  let ropeCos{{ c }} = vec4<f32>(
-    f32(cos_cache[ropeBase{{ c }}]),
-    f32(cos_cache[ropeBase{{ c }} + 1u]),
-    f32(cos_cache[ropeBase{{ c }} + 2u]),
-    f32(cos_cache[ropeBase{{ c }} + 3u])
   );
-  let ropeSin{{ c }} = vec4<f32>(
-    f32(sin_cache[ropeBase{{ c }}]),
-    f32(sin_cache[ropeBase{{ c }} + 1u]),
-    f32(sin_cache[ropeBase{{ c }} + 2u]),
-    f32(sin_cache[ropeBase{{ c }} + 3u])
   );
-  qr{{ c }} = select(
-    qr{{ c }} * ropeCos{{ c }} + pairedQ{{ c }} * ropeSin{{ c }},
-    qr{{ c }} * ropeCos{{ c }} - pairedQ{{ c }} * ropeSin{{ c }},
     lane8 < HALF_LPQ
   );
 {% endfor %}
 {% endif %}
-  var m: f32 = NEG_INF;
-  var l: f32 = 0.0;
 {% if useSeqlens %}
   // Resident rows of a windowed cache: the survivors were shifted down to [0, kvActive),
   // so query/key DISTANCE is unchanged and every bound below reads as if the cache were
   // exactly kvActive long. Rotary is excluded from this path (it would need the absolute
   // position, not the cache-relative one), so pastLenForRope keeps the physical length.
-  let kvActive = min({{ KVSEQ }}, u32(seqlens_k[b]) + 1u);
 {% endif %}
   // Causal ceiling per query; the key loop runs over the workgroup's union range
   // (uniform trip count), masking out-of-range (query, key) pairs.
@@ -267,16 +282,18 @@ fn main(
 {% if RIGHT or hasWindow %}
   let pastLen = select(0u, {{ KVA }} - {{ QSEQ }}, {{ KVA }} >= {{ QSEQ }});
 {% endif %}
 {% if hasCausal %}
 {% if RIGHT %}
-  var maxKjV = select({{ KVA }}, min(pastLen + qIdx + 1u, {{ KVA }}), {{ IS_CAUSAL }} != 0u);
 {% else %}
-  var maxKjV = select({{ KVA }}, min(qIdx + 1u, {{ KVA }}), {{ IS_CAUSAL }} != 0u);
 {% endif %}
 {% else %}
-  var maxKjV = {{ KVA }};
 {% endif %}
-  var minKjV: u32 = 0u;
   let lastQ = min(wg.x * TILE_Q + TILE_Q - 1u, {{ QSEQ }} - 1u);
 {% if hasCausal %}
 {% if RIGHT %}
@@ -289,16 +306,20 @@ fn main(
 {% endif %}
   var wgStartV: u32 = 0u;
 {% if hasWindow %}
-  let qAbsP1 = pastLen + qIdx + 1u;
-  maxKjV = min(maxKjV, qAbsP1);
-  if (qAbsP1 > params.windowSize) { minKjV = qAbsP1 - params.windowSize; }
   let lastQAbsP1 = pastLen + lastQ + 1u;
   wgEndV = min(wgEndV, lastQAbsP1);
   let firstQAbsP1 = pastLen + wg.x * TILE_Q + 1u;
   if (firstQAbsP1 > params.windowSize) { wgStartV = firstQAbsP1 - params.windowSize; }
 {% endif %}
-  let maxKj = maxKjV;
-  let minKj = minKjV;
   let wgEnd = wgEndV;
   let wgStart = wgStartV;
 {% else %}
@@ -306,6 +327,10 @@ fn main(
   let minKj: u32 = 0u;
   let wgEnd = {{ KVA }};
   let wgStart: u32 = 0u;
 {% endif %}
 {% if KL == "bhsd" %}
 {% if quantCacheFormat == "int4" %}
@@ -322,7 +347,7 @@ fn main(
   loop {
     if (kStart >= wgEnd) { break; }
-    // Cooperative K/V tile load (vec4-coalesced; OOB keys zero-filled).
     workgroupBarrier();
     for (var i: u32 = tid; i < TILE_K * HEAD_DIM_V4; i = i + WG) {
       let slot = i / HEAD_DIM_V4;
@@ -362,8 +387,7 @@ fn main(
     }
 {% endif %}
     workgroupBarrier();
-    // TILE_K remains a small constant-trip loop: compilers can unroll it without
-    // multiplying source size, while the q/o register slices remain named.
 {% if not useSubgroups and batchNoSgReduction %}
     // First publish every key's partial dot without intervening barriers.
     var s: array<f32, TILE_K>;
@@ -418,6 +442,67 @@ fn main(
 {% endif %}
       }
     }
 {% else %}
     var s: array<f32, TILE_K>;
     for (var kk: u32 = 0u; kk < TILE_K; kk = kk + 1u) {
@@ -480,30 +565,55 @@ fn main(
     // Per-thread online softmax over the tile. s[kk] is reused to hold the
     // exponentiated probabilities for the PV accumulation below.
-    var tileMax: f32 = s[0];
     for (var kk: u32 = 1u; kk < TILE_K; kk = kk + 1u) {
-      tileMax = max(tileMax, s[kk]);
     }
-    let newMax = max(m, tileMax);
-    let corr = select(exp(m - newMax), 0.0, m == NEG_INF);
-    var pSum: f32 = 0.0;
     for (var kk: u32 = 0u; kk < TILE_K; kk = kk + 1u) {
-      let pk = select(0.0, exp(s[kk] - newMax), s[kk] != NEG_INF);
-      s[kk] = pk;
-      pSum = pSum + pk;
     }
-    l = l * corr + pSum;
-    m = newMax;
-    // A boundary tile holds keys outside this query's attended range, and their staged V rows can
-    // be anything — a dynamic-rows prefill legitimately leaves pad rows unwritten, so stale NaN is
-    // reachable. An excluded key's weight is exactly 0.0, but 0.0 * NaN is NaN, so the multiply
-    // cannot be trusted to drop it: the guarded loop selects the V operand away instead. Interior
-    // tiles keep the unguarded FMA chain; the guard changes nothing there anyway (for finite v,
-    // acc + 0.0 * v == acc bit-for-bit), it only costs the extra select.
-    // Scope: only RANGE exclusion (minKj/maxKj) needs this. Every op that binds
-    // a mask supplies materialized K/V, so a mask-rejected key's exact zero
-    // weight multiplies finite data. If a masked op ever grows a producer that
-    // skips rows, that argument dies and the mask path needs the guard.
     let tileInterior = kStart >= minKj && kStart + TILE_K <= maxKj;
 {% for c in range(SLICE_COUNT) %}
     {
@@ -520,35 +630,38 @@ fn main(
       o{{ c }} = acc;
     }
 {% endfor %}
     kStart = kStart + TILE_K;
   }
-  if (qValid) {
 {% if QL == "bhsd" %}
-    let outBase4 = ((b * Q_HEADS + h) * {{ QSEQ }} + qIdx) * HEAD_DIM_V4 + lane8 * SLICE;
 {% else %}
-    let outBase4 = (b * {{ QSEQ }} + qIdx) * Q_HIDDEN_V4 + h * HEAD_DIM_V4 + lane8 * SLICE;
 {% endif %}
 {% if hasHeadSink %}
     // The head sink is a learned logit that competes with the keys but carries
     // no value, so it enters the denominator only and the weighted sum above is
-    // untouched. Renormalizing against max(m, sink) keeps the exponentials in
     // range when the sink dominates a fully-masked row.
     let sink = f32(head_sink[h]);
-    let finalM = max(m, sink);
-    let accScale = exp(m - finalM);
-    let inv = accScale / (exp(sink - finalM) + l * accScale);
 {% else %}
-    // l == 0 means this query had no probability-bearing key: either its
     // causal/window range is empty or its logical bool mask rejects every key.
     // Emit 0 rather than 0/0. Other attention paths enforce the same empty-row
     // contract by selecting on positive global mass.
-    let inv = select(0.0, 1.0 / l, l > 0.0);
 {% endif %}
-{% macro attention_value(c) %}o{{ c }} * inv{% if hasBias %} + load_bias4(2u * {{ Q_HIDDEN }}u + h * HEAD_DIM, lane8 * SLICE + {{ c }}u){% endif %}{% endmacro %}
 {% for c in range(SLICE_COUNT) %}
-    {{ OUTPUT }}[outBase4 + {{ c }}u] = vec4<{{ scalar }}>({{ attention_value(c) }});
 {% endfor %}
   }
 }

+{% 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 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 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
 {% if useSubgroups %}
 enable subgroups;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
 {% set ST = "f16" if usesF16 else "f32" %}
 // K/V tiles are staged as f16 when this specialization uses f16, reducing
 {% 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 %}
 {% set RIGHT = causalRightAlign is defined and causalRightAlign %}
 // Tiled flash prefill attention with configurable-width query clusters for
 // token-major [batch, seq, heads*headDim] attention ops. Each workgroup covers
 // TILE_Q queries with one LPQ-lane cluster per query; K/V are staged in
+// workgroup memory once per TILE_K tile. Within a cluster, each lane holds named
+// q/o slices of HEAD_DIM/LPQ dimensions. Per-lane dot fragments are
 // combined with subgroup shuffles when available, or shared-memory reductions
 // on compatibility variants. Dot products, online softmax, and weighted-value
 // sums accumulate in f32. Optional masks are broadcast-strided; the
 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 %}
 const HALF_LPQ: u32 = LPQ / 2u;
 {% endif %}
 const TILE_K: u32 = {{ TILE_K }}u;
+{% if QPL > 1 %}
+const QPL: u32 = {{ QPL }}u;
+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
 {% 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 %}
 {{ emit_quant_load4(quantCacheFormat, "value", VALUE, "v_scale") }}
 {% endif %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(
   @builtin(workgroup_id) wg: vec3<u32>,
   // portability.
   if (sgSize < LPQ || sgSize % LPQ != 0u) { return; }
 {% endif %}
+{% if QPL > 1 %}
+  let qSub = (tid / LPQ) * QPL;
+{% else %}
   let qSub = tid / LPQ;
+{% endif %}
   let lane8 = tid % LPQ;
+{% for qi in range(QPL) %}
+  let {{ qn("qIdx", qi) }} = wg.x * TILE_Q + qSub{% if qi > 0 %} + {{ qi }}u{% endif %};
+  let {{ qn("qValid", qi) }} = {{ qn("qIdx", qi) }} < {{ QSEQ }} && h < Q_HEADS;
+{% endfor %}
   let hKv = h / (Q_HEADS / KV_HEADS);
   let SCALE = scale_value();
+  // Per-thread q slice and output accumulator, each stored in SLICE named vec4 values.
+{% for qi in range(QPL) %}
+  let {{ qn("qClamped", qi) }} = min({{ qn("qIdx", qi) }}, {{ QSEQ }} - 1u);
 {% if QL == "bhsd" %}
+  let {{ qn("qBase4", qi) }} = ((b * Q_HEADS + h) * {{ QSEQ }} + {{ qn("qClamped", qi) }}) * HEAD_DIM_V4 + lane8 * SLICE;
 {% else %}
+  let {{ qn("qBase4", qi) }} = (b * {{ QSEQ }} + {{ qn("qClamped", qi) }}) * {{ Q_STRIDE }} + h * HEAD_DIM_V4 + lane8 * SLICE;
 {% 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 %}
 {% if FUSED_ROTARY %}
+  // The present-cache keys are already rotary transformed. Transform Q once
+  // before the K loop, avoiding a separate full-size Q scratch pass. NeoX
+  // half-split pairs lanes separated by LPQ/2.
 {% if not useSubgroups %}
   let qHeadBase4 = qBase4 - lane8 * SLICE;
   let pairedLane = (lane8 + HALF_LPQ) % LPQ;
 {% endif %}
   let ropeChunk = (lane8 % HALF_LPQ) * SLICE;
   let pastLenForRope = select(0u, {{ KVSEQ }} - {{ QSEQ }}, {{ KVSEQ }} >= {{ QSEQ }});
+{% for qi in range(QPL) %}
+  let {{ qn("ropePosition", qi) }} = pastLenForRope + {{ qn("qClamped", qi) }};
 {% for c in range(SLICE_COUNT) %}
 {% if useSubgroups %}
+  let {{ qn("pairedQ" ~ c, qi) }} = subgroupShuffleXor({{ qn("qr" ~ c, qi) }}, {{ ROPE_LANE_XOR }}u);
 {% else %}
+  let {{ qn("pairedQ" ~ c, qi) }} = vec4<f32>({{ QUERY }}[qHeadBase4 + pairedLane * SLICE + {{ c }}u]);
+{% endif %}
+  let {{ qn("ropeBase" ~ c, qi) }} = {{ qn("ropePosition", qi) }} * HALF_DIM + (ropeChunk + {{ c }}u) * 4u;
+  let {{ qn("ropeCos" ~ c, qi) }} = vec4<f32>(
+    f32(cos_cache[{{ qn("ropeBase" ~ c, qi) }}]),
+    f32(cos_cache[{{ qn("ropeBase" ~ c, qi) }} + 1u]),
+    f32(cos_cache[{{ qn("ropeBase" ~ c, qi) }} + 2u]),
+    f32(cos_cache[{{ qn("ropeBase" ~ c, qi) }} + 3u])
   );
+  let {{ qn("ropeSin" ~ c, qi) }} = vec4<f32>(
+    f32(sin_cache[{{ qn("ropeBase" ~ c, qi) }}]),
+    f32(sin_cache[{{ qn("ropeBase" ~ c, qi) }} + 1u]),
+    f32(sin_cache[{{ qn("ropeBase" ~ c, qi) }} + 2u]),
+    f32(sin_cache[{{ qn("ropeBase" ~ c, qi) }} + 3u])
   );
+  {{ qn("qr" ~ c, qi) }} = select(
+    {{ qn("qr" ~ c, qi) }} * {{ qn("ropeCos" ~ c, qi) }} + {{ qn("pairedQ" ~ c, qi) }} * {{ qn("ropeSin" ~ c, qi) }},
+    {{ qn("qr" ~ c, qi) }} * {{ qn("ropeCos" ~ c, qi) }} - {{ qn("pairedQ" ~ c, qi) }} * {{ qn("ropeSin" ~ c, qi) }},
     lane8 < HALF_LPQ
   );
 {% endfor %}
+{% endfor %}
 {% endif %}
+{% for qi in range(QPL) %}
+  var {{ qn("m", qi) }}: f32 = NEG_INF;
+  var {{ qn("l", qi) }}: f32 = 0.0;
+{% endfor %}
 {% if useSeqlens %}
   // Resident rows of a windowed cache: the survivors were shifted down to [0, kvActive),
   // so query/key DISTANCE is unchanged and every bound below reads as if the cache were
   // exactly kvActive long. Rotary is excluded from this path (it would need the absolute
   // position, not the cache-relative one), so pastLenForRope keeps the physical length.
+  // Clamping below by the query count matches the cache-update passes and the scalar
+  // path: a right-padded batch (seqlens_k[b]+1 < qSeq) still appends its whole chunk,
+  // so its queries score against all of it.
+  let kvActive = min({{ KVSEQ }}, max({{ QSEQ }}, u32(seqlens_k[b]) + 1u));
 {% endif %}
   // Causal ceiling per query; the key loop runs over the workgroup's union range
   // (uniform trip count), masking out-of-range (query, key) pairs.
 {% if RIGHT or hasWindow %}
   let pastLen = select(0u, {{ KVA }} - {{ QSEQ }}, {{ KVA }} >= {{ QSEQ }});
 {% endif %}
+{% for qi in range(QPL) %}
 {% if hasCausal %}
 {% if RIGHT %}
+  var {{ qn("maxKjV", qi) }} = select({{ KVA }}, min(pastLen + {{ qn("qIdx", qi) }} + 1u, {{ KVA }}), {{ IS_CAUSAL }} != 0u);
 {% else %}
+  var {{ qn("maxKjV", qi) }} = select({{ KVA }}, min({{ qn("qIdx", qi) }} + 1u, {{ KVA }}), {{ IS_CAUSAL }} != 0u);
 {% endif %}
 {% else %}
+  var {{ qn("maxKjV", qi) }} = {{ KVA }};
 {% endif %}
+  var {{ qn("minKjV", qi) }}: u32 = 0u;
+{% endfor %}
   let lastQ = min(wg.x * TILE_Q + TILE_Q - 1u, {{ QSEQ }} - 1u);
 {% if hasCausal %}
 {% if RIGHT %}
 {% endif %}
   var wgStartV: u32 = 0u;
 {% if hasWindow %}
+{% for qi in range(QPL) %}
+  let {{ qn("qAbsP1", qi) }} = pastLen + {{ qn("qIdx", qi) }} + 1u;
+  {{ qn("maxKjV", qi) }} = min({{ qn("maxKjV", qi) }}, {{ qn("qAbsP1", qi) }});
+  if ({{ qn("qAbsP1", qi) }} > params.windowSize) { {{ qn("minKjV", qi) }} = {{ qn("qAbsP1", qi) }} - params.windowSize; }
+{% endfor %}
   let lastQAbsP1 = pastLen + lastQ + 1u;
   wgEndV = min(wgEndV, lastQAbsP1);
   let firstQAbsP1 = pastLen + wg.x * TILE_Q + 1u;
   if (firstQAbsP1 > params.windowSize) { wgStartV = firstQAbsP1 - params.windowSize; }
 {% endif %}
+{% for qi in range(QPL) %}
+  let {{ qn("maxKj", qi) }} = {{ qn("maxKjV", qi) }};
+  let {{ qn("minKj", qi) }} = {{ qn("minKjV", qi) }};
+{% endfor %}
   let wgEnd = wgEndV;
   let wgStart = wgStartV;
 {% else %}
   let minKj: u32 = 0u;
   let wgEnd = {{ KVA }};
   let wgStart: u32 = 0u;
+{% for qi in range(1, QPL) %}
+  let {{ qn("maxKj", qi) }} = maxKj;
+  let {{ qn("minKj", qi) }} = minKj;
+{% endfor %}
 {% endif %}
 {% if KL == "bhsd" %}
 {% if quantCacheFormat == "int4" %}
   loop {
     if (kStart >= wgEnd) { break; }
+    // Cooperative K/V tile load with out-of-bounds keys zero-filled.
     workgroupBarrier();
     for (var i: u32 = tid; i < TILE_K * HEAD_DIM_V4; i = i + WG) {
       let slot = i / HEAD_DIM_V4;
     }
 {% endif %}
     workgroupBarrier();
+    // 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>;
 {% endif %}
       }
     }
+{% elif QPL > 1 %}
+{% for qi in range(QPL) %}
+    var {{ qn("s", qi) }}: array<f32, TILE_K>;
+{% endfor %}
+    for (var kk: u32 = 0u; kk < TILE_K; kk = kk + 1u) {
+      let kj = kStart + kk;
+{% for qi in range(QPL) %}
+      {{ qn("s", qi) }}[kk] = NEG_INF;
+      var {{ qn("part", qi) }}: f32 = 0.0;
+{% endfor %}
+      let kb = kk * HEAD_DIM_V4 + lane8 * SLICE;
+{% for c in range(SLICE_COUNT) %}
+      let kw{{ c }} = vec4<f32>(k_tile[kb + {{ c }}u]);
+{% for qi in range(QPL) %}
+      {{ qn("part", qi) }} = {{ qn("part", qi) }} + dot({{ qn("qr" ~ c, qi) }}, kw{{ c }});
+{% endfor %}
+{% endfor %}
+      // One butterfly carries every query's partial: the shuffles cost the same
+      // whether the payload is one score or QPL of them.
+      var partv = vec{{ QPL }}<f32>({% for qi in range(QPL) %}{{ qn("part", qi) }}{% if not loop.last %}, {% endif %}{% endfor %});
+      partv = partv + subgroupShuffleXor(partv, 1u);
+{% for shift in [2, 4] %}
+{% if LPQ > shift %}
+      partv = partv + subgroupShuffleXor(partv, {{ shift }}u);
+{% endif %}
+{% endfor %}
+{% for qi in range(QPL) %}
+      {
+        let sc = partv[{{ qi }}];
+      if (kj >= {{ qn("minKj", qi) }} && kj < {{ qn("maxKj", qi) }}) {
+  {% if hasMask %}
+  {% if STAGE_MASK %}
+          let maskValue = mask_tile[(qSub + {{ qi }}u) * TILE_K + kk];
+  {% else %}
+          // Broadcast-strided mask address (a 0 stride collapses that axis; rank-2
+          // [q, k] masks set batch/head strides to 0). {{ qn("qClamped", qi) }} keeps the seq index
+          // in-bounds for padding queries in the last tile (their output is dropped).
+          let maskIndex = b * params.maskBatchStride + h * params.maskHeadStride + {{ qn("qClamped", qi) }} * params.maskSeqStride + kj;
+  {% endif %}
+  {% if maskIsKeyKeep %}
+          // A broadcast key mask uses 1 for a retained key and 0 for padding.
+          {{ qn("s", qi) }}[kk] = {{ score_expr("sc") }} + (1.0 - f32({{ MASK_ELEMENT }})) * MASK_NEG;
+  {% elif maskIsBool %}
+          // Logical bool: a rejected key contributes no softmax mass. Leaving
+          // the initialized NEG_INF sentinel in place makes a fully masked row
+          // land on the zero-denominator output guard below.
+          if ({{ MASK_ELEMENT }} != 0u) {
+            {{ qn("s", qi) }}[kk] = {{ score_expr("sc") }};
+          } else {
+            {{ qn("s", qi) }}[kk] = MASK_NEG;
+          }
+  {% else %}
+          {{ qn("s", qi) }}[kk] = {{ score_expr("sc") }} + {{ MASK_ADDITIVE }};
+  {% endif %}
+  {% else %}
+          {{ qn("s", qi) }}[kk] = {{ score_expr("sc") }};
+  {% endif %}
+        }
+      }
+{% endfor %}
+    }
 {% else %}
     var s: array<f32, TILE_K>;
     for (var kk: u32 = 0u; kk < TILE_K; kk = kk + 1u) {
     // 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) %}
+    {
+{% for qi in range(QPL) %}
+      var {{ qn("acc", qi) }} = {{ qn("o" ~ c, qi) }} * {{ qn("corr", qi) }};
+{% endfor %}
+      if (tileInterior) {
+        for (var kk: u32 = 0u; kk < TILE_K; kk = kk + 1u) {
+          let vw = vec4<f32>(v_tile[kk * HEAD_DIM_V4 + lane8 * SLICE + {{ c }}u]);
+{% for qi in range(QPL) %}
+          {{ qn("acc", qi) }} = {{ qn("acc", qi) }} + {{ qn("s", qi) }}[kk] * vw;
+{% endfor %}
+        }
+      } else {
+        for (var kk: u32 = 0u; kk < TILE_K; kk = kk + 1u) {
+          let vw = vec4<f32>(v_tile[kk * HEAD_DIM_V4 + lane8 * SLICE + {{ c }}u]);
+{% for qi in range(QPL) %}
+          {{ qn("acc", qi) }} = {{ qn("acc", qi) }} + {{ qn("s", qi) }}[kk] * select(vec4<f32>(), vw, {{ qn("s", qi) }}[kk] != 0.0);
+{% endfor %}
+        }
+      }
+{% for qi in range(QPL) %}
+      {{ qn("o" ~ c, qi) }} = {{ qn("acc", qi) }};
+{% endfor %}
+    }
+{% endfor %}
+{% else %}
     let tileInterior = kStart >= minKj && kStart + TILE_K <= maxKj;
 {% for c in range(SLICE_COUNT) %}
     {
       o{{ c }} = acc;
     }
 {% endfor %}
+{% endif %}
     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" %}
+    let outBase4 = ((b * Q_HEADS + h) * {{ QSEQ }} + {{ qn("qIdx", qi) }}) * HEAD_DIM_V4 + lane8 * SLICE;
 {% else %}
+    let outBase4 = (b * {{ QSEQ }} + {{ qn("qIdx", qi) }}) * Q_HIDDEN_V4 + h * HEAD_DIM_V4 + lane8 * SLICE;
 {% endif %}
 {% if hasHeadSink %}
     // The head sink is a learned logit that competes with the keys but carries
     // no value, so it enters the denominator only and the weighted sum above is
+    // untouched. Renormalizing against max({{ qn("m", qi) }}, sink) keeps the exponentials in
     // range when the sink dominates a fully-masked row.
     let sink = f32(head_sink[h]);
+    let finalM = max({{ qn("m", qi) }}, sink);
+    let accScale = exp({{ qn("m", qi) }} - finalM);
+    let inv = accScale / (exp(sink - finalM) + {{ qn("l", qi) }} * accScale);
 {% else %}
+    // {{ qn("l", qi) }} == 0 means this query had no probability-bearing key: either its
     // causal/window range is empty or its logical bool mask rejects every key.
     // Emit 0 rather than 0/0. Other attention paths enforce the same empty-row
     // contract by selecting on positive global mass.
+    let inv = select(0.0, 1.0 / {{ qn("l", qi) }}, {{ qn("l", qi) }} > 0.0);
 {% endif %}
 {% for c in range(SLICE_COUNT) %}
+    {{ OUTPUT }}[outBase4 + {{ c }}u] = vec4<{{ scalar }}>({{ attention_value(c, qi) }});
 {% endfor %}
   }
+{% endfor %}
 }

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

@@ -9,12 +9,8 @@
 // cancels in softmax. Arbitrary attention masks use the fallback routes.
 {% set ST = "f16" if usesF16 else "f32" %}
 {% set components = ["x", "y", "z", "w"] %}
-{% set USE_SUBGROUPS = source.useSubgroups if source.useSubgroups is defined else true %}
 {% set Q_STEP = qStep if qStep is defined else 32 %}
-{% set CAUSAL = false if (hasCausal is defined and not hasCausal) else true %}
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {% if USE_SUBGROUPS %}
 enable subgroups;
 {% endif %}
@@ -25,9 +21,6 @@ const Q_HEADS:     u32 = {{ qNumHeads }}u;
 const KV_HEADS:    u32 = {{ kvNumHeads }}u;
 const K_STEP:      u32 = {{ kStep }}u;
 const Q_STEP:      u32 = {{ Q_STEP }}u;
-{% if hasBias %}
-const HEAD_DIM:    u32 = {{ headDim }}u;
-{% endif %}
 const Q_HIDDEN_V4:  u32 = Q_HEADS * HEAD_DIM_V4;
 const KV_HIDDEN_V4: u32 = KV_HEADS * HEAD_DIM_V4;
 const NEG_INF:     f32 = -3.4028234663852886e38;
@@ -50,12 +43,6 @@ fn scale_value() -> f32 {
 }
-{% if hasBias %}
-fn load_bias4(base: u32, d4: u32) -> vec4<f32> {
-  let offset = base + d4 * 4u;
-  return vec4<f32>(bias[offset], bias[offset + 1u], bias[offset + 2u], bias[offset + 3u]);
-}
-{% endif %}
 @compute @workgroup_size({{ Q_STEP }}, 1, 1)
 fn main(
@@ -86,9 +73,6 @@ fn main(
   for (var d4: u32 = 0u; d4 < HEAD_DIM_V4; d4 = d4 + 1u) {
     if (valid_q) {
       q_tile[d4] = query[qBaseV4 + d4];
-{% if hasBias %}
-      q_tile[d4] = vec4<{{ ST }}>(vec4<f32>(q_tile[d4]) + load_bias4(h * HEAD_DIM, d4));
-{% endif %}
     } else {
       q_tile[d4] = vec4<{{ ST }}>(0.0);
     }
@@ -101,14 +85,8 @@ 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.
-{% if CAUSAL %}
-  let lastQ = min(wg.x * Q_STEP + Q_STEP - 1u, params.qSeq - 1u);
-  let kvEnd = select(params.kvSeq, min(lastQ + 1u, params.kvSeq), params.isCausal != 0u);
-  let myMaxKj = select(params.kvSeq, min(qi + 1u, params.kvSeq), params.isCausal != 0u);
-{% else %}
   let kvEnd = params.kvSeq;
   let myMaxKj = params.kvSeq;
-{% endif %}
   let kvBase = b * params.kvSeq * KV_HIDDEN_V4 + hKv * HEAD_DIM_V4;
   let kvKeyStride = KV_HIDDEN_V4;
@@ -236,11 +214,7 @@ fn main(
   if (valid_q) {
     let outBaseV4 = (b * params.qSeq + qi) * Q_HIDDEN_V4 + h * HEAD_DIM_V4;
     for (var d4: u32 = 0u; d4 < HEAD_DIM_V4; d4 = d4 + 1u) {
-{% if hasBias %}
-      output[outBaseV4 + d4] = vec4<{{ scalar }}>(o_tile[d4] + load_bias4(2u * {{ qHidden }}u + h * HEAD_DIM, d4));
-{% else %}
       output[outBaseV4 + d4] = vec4<{{ scalar }}>(o_tile[d4]);
-{% endif %}
     }
   }
 }

 // cancels in softmax. Arbitrary attention masks use the fallback routes.
 {% 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 %}
 const KV_HEADS:    u32 = {{ kvNumHeads }}u;
 const K_STEP:      u32 = {{ kStep }}u;
 const Q_STEP:      u32 = {{ Q_STEP }}u;
 const Q_HIDDEN_V4:  u32 = Q_HEADS * HEAD_DIM_V4;
 const KV_HIDDEN_V4: u32 = KV_HEADS * HEAD_DIM_V4;
 const NEG_INF:     f32 = -3.4028234663852886e38;
 }
 @compute @workgroup_size({{ Q_STEP }}, 1, 1)
 fn main(
   for (var d4: u32 = 0u; d4 < HEAD_DIM_V4; d4 = d4 + 1u) {
     if (valid_q) {
       q_tile[d4] = query[qBaseV4 + d4];
     } else {
       q_tile[d4] = vec4<{{ ST }}>(0.0);
     }
   // 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;
   if (valid_q) {
     let outBaseV4 = (b * params.qSeq + qi) * Q_HIDDEN_V4 + h * HEAD_DIM_V4;
     for (var d4: u32 = 0u; d4 < HEAD_DIM_V4; d4 = d4 + 1u) {
       output[outBaseV4 + d4] = vec4<{{ scalar }}>(o_tile[d4]);
     }
   }
 }

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

@@ -43,10 +43,9 @@ fn exp_shift(value: f32, maxValue: f32) -> f32 {
 @compute @workgroup_size(WG, 1, 1)
 fn main(
-  @builtin(global_invocation_id) gid: vec3<u32>,
-  @builtin(num_workgroups) nwg: vec3<u32>
 ) {
-  let row = gid.x + gid.y * nwg.x * WG;
   if (row >= params.rows) { return; }
   // `row` already runs over (batch, head, query) together, and the partial

 @compute @workgroup_size(WG, 1, 1)
 fn main(
+  @builtin(global_invocation_id) gid: vec3<u32>
 ) {
+  let row = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
   if (row >= params.rows) { return; }
   // `row` already runs over (batch, head, query) together, and the partial

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

@@ -9,22 +9,23 @@ enable subgroup_size_control;
 enable chromium_experimental_subgroup_matrix;
 diagnostic(off, chromium.subgroup_matrix_uniformity);
 {{ env.wgsl.resourceDeclarations }}
-{% set layout = source.layout | default("bsh") %}
 {% set headMajor = layout == "bhsd" %}
-{% set kvHeadMajor = (source.kvLayout | default(layout)) == "bhsd" %}
-{% set CAUSAL_UPPER_LEFT = source.causalUpperLeft is defined and source.causalUpperLeft %}
-{% set CAUSAL = (source.causalRightAlign is defined and source.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 = "HEAD_DIM" if headMajor else "HIDDEN" %}
 {% set KV_ROW_STRIDE = "HEAD_DIM" if kvHeadMajor else "KV_HIDDEN" %}
-{% set scorePhase = source.phase == "score" %}
-{% set SCORE_BIAS = scorePhase and source.scoreBias is defined and source.scoreBias %}
-{% set SCORE_WINDOW = CAUSAL and source.scoreWindow is defined and source.scoreWindow %}
-{% set USE_SEQLENS = source.useSeqlens is defined and source.useSeqlens %}
-{% set FUSED_SOFTMAX = source.fusedSoftmax is defined and source.fusedSoftmax %}
 {% set PRIVATE_ROW_STATS = FUSED_SOFTMAX and (materializedSgmatPrivateRowStats is defined and materializedSgmatPrivateRowStats) %}
 {% macro score_value(index, guard) %}
 {% if FUSED_SOFTMAX %}
@@ -49,7 +50,7 @@ select(0.0, scores[{{ index }}], {{ guard[1] }})
 {% 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 = source.emitRowStats is defined and source.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 %}
@@ -58,17 +59,14 @@ select(0.0, scores[{{ index }}], {{ guard[1] }})
    and (scorePhase or (not hasBias and MT == "f32"))
    and not EMIT_ROW_STATS %}
 {% set ANY_DIRECT_STORE = DIRECT_OUTPUT_STORE or RUNTIME_DIRECT_STORE %}
-{% set SCALE_IN_Q = DIRECT_SCORE_STORE or (RUNTIME_DIRECT_STORE and scorePhase) %}
-{% macro q_tile_value(index) %}{% if hasBias %}(query[{{ index }}] + bias[h * HEAD_DIM + k]){% else %}query[{{ index }}]{% endif %}{% if SCALE_IN_Q %} * score_scale{% endif %}{% endmacro %}
 const HEADS: u32 = {{ qNumHeads }}u;
-{% if kvNumHeads is defined %}
 const KV_HEADS: u32 = {{ kvNumHeads }}u;
-{% else %}
-const KV_HEADS: u32 = HEADS;
-{% endif %}
 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
@@ -158,10 +156,10 @@ fn main(
   let base_A = subtile_idy * SUB_ROWS;
   let base_B = subtile_idx * SUB_COLS;
 {% if USE_SEQLENS %}
-  // Rows the cache-update pass left resident: the survivors live in
-  // [0, kv_active), and every causal/window bound below uses this live length
-  // while params.kvSeq keeps the allocated capacity for strides.
-  let kv_active = min(params.kvSeq, u32(seqlens_k[b]) + 1u);
 {% endif %}
 {% if CAUSAL and scorePhase %}
   // Workgroup-uniform causal skip: a key tile starting past the last row's
@@ -172,7 +170,8 @@ fn main(
   // the same way: later rows only move the floor further right.
 {% endif %}
   let kv_causal_off = i32({% if USE_SEQLENS %}kv_active{% else %}params.kvSeq{% endif %}) - i32(params.qSeq);
-  if (i32(n_base) > kv_causal_off + i32(m_base + TILE_M) - 1{% if SCORE_WINDOW %}
       || i32(n_base + TILE_N) <= kv_causal_off + i32(m_base) + 1 - i32(params.windowSize){% endif %}) {
     for (var idx = li; idx < TILE_M * {{ SUBGROUP_COLS }}u; idx += {{ WORKGROUP_THREADS }}u) {
       let stat_row = m_base + idx / {{ SUBGROUP_COLS }}u;
@@ -195,11 +194,6 @@ fn main(
 {% if scorePhase %}
   let inner = HEAD_DIM;
-{% if SCALE_IN_Q %}
-  // The direct store has no epilogue, so apply the score scale to Q. Direct and
-  // guarded store paths then share one pre-scaled query tile.
-  let score_scale = {{ attentionScaleExpression }};
-{% endif %}
 {% else %}
   let inner = {% if USE_SEQLENS %}kv_active{% else %}params.kvSeq{% endif %};
 {% if CAUSAL %}
@@ -240,7 +234,7 @@ fn main(
 {% 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 source.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 (
@@ -278,10 +272,8 @@ fn main(
 {% endif %}
 {% if TILE_M_VALUE == 64 and TILE_N_VALUE == 64 and TILE_K_VALUE == 32 and WORKGROUP_THREADS == 256 %}
-    // BM64/BN64/WG256 is the high-throughput geometry on wide devices. Four
-    // adjacent threads each own one contiguous eight-element segment of a
-    // row, preserving the original coalesced load schedule while retaining
-    // the generalized K-tail guards.
     let a_row = li / 4u;
     let a_col = (li % 4u) * 8u;
     for (var i = 0u; i < 8u; i++) {
@@ -353,10 +345,10 @@ fn main(
       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.0h" if MT == "f16" else "0.0" }};
       if (row < params.qSeq && k < params.kvSeq) {
         let score_base = (b * HEADS + h) * params.qSeq * params.kvSeq;
-        loaded = {{ "f16(" if MT == "f16" else "" }}scores[score_base + row * params.kvSeq + k]{{ ")" if MT == "f16" else "" }};
       }
 {% endif %}
       tile_A[idx] = loaded;
@@ -412,18 +404,8 @@ fn main(
   }
 {% if ANY_DIRECT_STORE %}
-{% if RUNTIME_DIRECT_STORE %}
-  // Workgroup-uniform (both bases come from workgroup_id alone): an interior tile
-  // publishes through the subgroup-matrix collectives and is done; only an edge
-  // tile falls through to the scratch round trip and its barriers.
-  if (
-    m_base + TILE_M <= params.qSeq &&
-    n_base + TILE_N <= {% if scorePhase %}params.kvSeq{% else %}HEAD_DIM{% endif %}
-  ) {
-{% else %}
   // Every dispatched tile is interior, so the matrices can go straight to
   // storage without the compact scratch/readback epilogue.
-{% endif %}
 {% for row_block in range(ROW_BLOCKS) %}
 {% for col_block in range(COL_BLOCKS) %}
 {% if scorePhase %}
@@ -445,10 +427,6 @@ fn main(
 {% endif %}
 {% endfor %}
 {% endfor %}
-{% if RUNTIME_DIRECT_STORE %}
-    return;
-  }
-{% endif %}
 {% endif %}
 {% if not DIRECT_OUTPUT_STORE %}
   let row_in_block = lane / 4u;
@@ -489,9 +467,7 @@ fn main(
           + 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 %}
@@ -502,11 +478,14 @@ fn main(
       ];
 {% endif %}
       if (i32(col) > kv_causal_off + i32(row)) { scored = -FLT_MAX; }
 {% if SCORE_WINDOW %}
       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[
         (b * HEADS + h) * params.qSeq * params.kvSeq + row * params.kvSeq + col

 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 %}
+{% set USE_SEQLENS = useSeqlens is defined and useSeqlens %}
+{% 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 %}
 {% 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 %}
    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
   let base_A = subtile_idy * SUB_ROWS;
   let base_B = subtile_idx * SUB_COLS;
 {% if USE_SEQLENS %}
+  // Resident rows occupy [0, kv_active). Causal and window bounds use this live
+  // length while params.kvSeq remains the allocated stride. Clamp by qSeq so a
+  // right-padded batch still includes the complete appended query chunk.
+  let kv_active = min(params.kvSeq, max(params.qSeq, u32(seqlens_k[b]) + 1u));
 {% endif %}
 {% if CAUSAL and scorePhase %}
   // Workgroup-uniform causal skip: a key tile starting past the last row's
   // the same way: later rows only move the floor further right.
 {% endif %}
   let kv_causal_off = i32({% if USE_SEQLENS %}kv_active{% else %}params.kvSeq{% endif %}) - i32(params.qSeq);
+  if (i32(n_base) > kv_causal_off + i32(m_base + TILE_M) - 1{% if USE_SEQLENS %}
+      || n_base >= kv_active{% endif %}{% if SCORE_WINDOW %}
       || i32(n_base + TILE_N) <= kv_causal_off + i32(m_base) + 1 - i32(params.windowSize){% endif %}) {
     for (var idx = li; idx < TILE_M * {{ SUBGROUP_COLS }}u; idx += {{ WORKGROUP_THREADS }}u) {
       let stat_row = m_base + idx / {{ SUBGROUP_COLS }}u;
 {% if scorePhase %}
   let inner = HEAD_DIM;
 {% else %}
   let inner = {% if USE_SEQLENS %}kv_active{% else %}params.kvSeq{% endif %};
 {% if CAUSAL %}
 {% 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 (
 {% endif %}
 {% if TILE_M_VALUE == 64 and TILE_N_VALUE == 64 and TILE_K_VALUE == 32 and WORKGROUP_THREADS == 256 %}
+    // In the BM64/BN64/BK32/WG256 geometry, four adjacent threads each own one
+    // contiguous eight-element segment of a row. K-tail guards remain active.
     let a_row = li / 4u;
     let a_col = (li % 4u) * 8u;
     for (var i = 0u; i < 8u; i++) {
       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 %}
       tile_A[idx] = loaded;
   }
 {% if ANY_DIRECT_STORE %}
   // Every dispatched tile is interior, so the matrices can go straight to
   // storage without the compact scratch/readback epilogue.
 {% for row_block in range(ROW_BLOCKS) %}
 {% for col_block in range(COL_BLOCKS) %}
 {% if scorePhase %}
 {% endif %}
 {% endfor %}
 {% endfor %}
 {% endif %}
 {% if not DIRECT_OUTPUT_STORE %}
   let row_in_block = lane / 4u;
           + row_in_block * 8u + col_in_block + pair
       ];
 {% if scorePhase %}
       let scale = {{ attentionScaleExpression }};
 {% if CAUSAL %}
       var scored = result * scale;
 {% if SCORE_BIAS %}
       ];
 {% endif %}
       if (i32(col) > kv_causal_off + i32(row)) { scored = -FLT_MAX; }
+{% if USE_SEQLENS %}
+      if (col >= kv_active) { scored = -FLT_MAX; }
+{% endif %}
 {% if SCORE_WINDOW %}
       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

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

@@ -1,28 +1,14 @@
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
-{% set MASK_BATCH = "batch * params.maskBatchStride + " if source.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. It is the portable
-// path for any workgroup size and device tier.
 // Layout: rank-3 token-major [batch, seq, heads * headDim].
-{% if headsFromParams %}
-// Q/KV head counts are uniforms; the GQA mapping is h / (qHeads / kvHeads).
-{% endif %}
-{% if hasKeyLimit %}
-// params.keyLimit caps the key range when nonzero.
-{% endif %}
-{% if scaleFallbackRsqrt %}
 // An omitted scale uses 1/sqrt(headDim); explicit zero remains zero through
 // the scaleIsExplicitZero specialization.
-{% endif %}
-{% 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).
@@ -30,10 +16,8 @@ enable f16;
 const HEAD_DIM: u32 = {{ headDim }}u;
 const Q_HIDDEN: u32 = {{ qHidden }}u;
 const KV_HIDDEN: u32 = {{ kvHidden }}u;
-{% if not headsFromParams %}
 const Q_HEADS: u32 = {{ qNumHeads }}u;
 const KV_HEADS: u32 = {{ kvNumHeads }}u;
-{% endif %}
 {% set qHeads = "params.qHeads" if headsFromParams else "Q_HEADS" %}
 {% set kvHeads = "params.kvHeads" if headsFromParams else "KV_HEADS" %}
 {% set scale = scale | default("0.0") %}
@@ -80,18 +64,12 @@ fn {{ name }}(value: f32, tid: u32) -> f32 {
 {{ wgsl_tree_reduce_f32("reduce_sum", "add", "partial", "WG") }}
-{% if scaleFallbackRsqrt %}
 {% 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 }}));
 }
-{% else %}
-fn scale_value() -> f32 {
-  return {{ scale }};
-}
-{% endif %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(
@@ -132,12 +110,6 @@ fn main(
   var maxKj = params.kvSeq;
   var minKj: u32 = 0u;
-{% if hasCausal %}
-  maxKj = min(maxKj, select(params.kvSeq, query_token + 1u, params.isCausal != 0u));
-{% endif %}
-{% if hasKeyLimit %}
-  maxKj = min(maxKj, select(params.kvSeq, params.keyLimit, params.keyLimit > 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
@@ -156,11 +128,6 @@ fn main(
     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 = partial_dot + q_value * k_value;
     }
@@ -186,15 +153,14 @@ 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]);
-{% 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();
   }
-  let inv_denom = 1.0 / running_denom;
   for (var d: u32 = tid; d < HEAD_DIM; d = d + WG) {
     output[qBase + d] = {{ outputScalar }}(running_out[d] * inv_denom);
   }

 {{ 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).
 const HEAD_DIM: u32 = {{ headDim }}u;
 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") %}
 {{ 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(
   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
     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;
     }
     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();
   }
+  // A query whose every key is masked to -inf (or whose key range is empty)
+  // leaves running_denom at zero; emit a zero row rather than 0 * (1/0) = NaN.
+  let inv_denom = select(0.0, 1.0 / running_denom, running_denom > 0.0);
   for (var d: u32 = tid; d < HEAD_DIM; d = d + WG) {
     output[qBase + d] = {{ outputScalar }}(running_out[d] * inv_denom);
   }

build/webgpu/bench.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "com.microsoft.GroupQueryAttention",
   "tunableSpace": {
     "PREFILL_QUERY_TILE": [16, 32],
     "PREFILL_LANES_PER_QUERY": [4, 8],
@@ -134,7 +133,7 @@
         ]
       },
       "provenance": {
-        "notes": "Long-context decode beyond the MAX_SPLITS boundary: ceilDiv(kvSeq, QKV_SPLIT_TILE_K)=64 exceeds the cap of 16. Paired with the headDim-32 control, this case checks the tradeoff between additional split parallelism and merge overhead without assuming that one flat cap is optimal for both geometries."
       }
     },
     {
@@ -711,9 +710,7 @@
     {
       "name": "quant-int8-prefill-h32kv8-d128-s512-pathology",
       "preset": "stress",
-      "provenance": {
-        "notes": "Production-sized prefill for the int8 KV-cache contract, retained as a performance witness for quantized-cache attention."
-      },
       "attrs": {
         "num_heads": 32,
         "kv_num_heads": 8,
@@ -768,9 +765,7 @@
     {
       "name": "quant-int4-prefill-h32kv8-d128-s512-pathology",
       "preset": "stress",
-      "provenance": {
-        "notes": "Production-sized prefill for the packed int4 KV-cache contract, retained as a performance witness for quantized-cache attention."
-      },
       "attrs": {
         "num_heads": 32,
         "kv_num_heads": 8,
@@ -1048,7 +1043,7 @@
       "name": "headsink-prefill-h32kv8-d128-s512-generic-pathology",
       "preset": "stress",
       "provenance": {
-        "notes": "Llama-sized GQA prefill with smooth-softmax head sinks, measuring the additional score transformation on a production attention shape."
       },
       "attrs": { "num_heads": 32, "kv_num_heads": 8 },
       "inputs": {
@@ -1091,7 +1086,7 @@
       "name": "bias-headsink-prefill-h32kv8-d128-s512-generic-pathology",
       "preset": "stress",
       "provenance": {
-        "notes": "The valid additive-bias plus smooth-softmax head-sink combination has a distinct manifest variant but no flash-prefill specialization."
       },
       "attrs": { "num_heads": 32, "kv_num_heads": 8 },
       "inputs": {
@@ -1138,9 +1133,7 @@
     {
       "name": "softcap-prefill-h32kv8-d128-s512-generic-pathology",
       "preset": "stress",
-      "provenance": {
-        "notes": "Production prefill with model-used score soft-capping, measuring the capped-score semantic on the past-KV contract."
-      },
       "attrs": { "num_heads": 32, "kv_num_heads": 8, "softcap": 30 },
       "inputs": {
         "queryT": {

 {
   "tunableSpace": {
     "PREFILL_QUERY_TILE": [16, 32],
     "PREFILL_LANES_PER_QUERY": [4, 8],
         ]
       },
       "provenance": {
+        "notes": "Long-context decode yields 64 candidate K partitions, exceeding the 16-partition cap and exercising capped split parallelism plus its merge at head size 128."
       }
     },
     {
     {
       "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,
     {
       "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,
       "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": {
       "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": {
     {
       "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": {

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

@@ -73,9 +73,9 @@ 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(num_workgroups) nwg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
   // One workgroup per query; fold large query grids over x/y.
-  let qi = wid.x + wid.y * nwg.x;
   let total = params.batch * Q_HEADS * params.qSeq;
   if (qi >= total) { return; }
   let lane = lid.x;
@@ -178,7 +178,10 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(num_workgroups) nwg: vec
     smoothTerm = exp(sink - finalM);
   }
   let accScale = exp(gm - finalM);
-  let invDenom = accScale / (smoothTerm + glsum * accScale);
   // Cross-lane sum of the rescaled V accumulators, written cooperatively.
   let oBase = (b * params.qSeq + s) * Q_HIDDEN + h * HEAD_DIM;
@@ -190,13 +193,13 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(num_workgroups) nwg: vec
 }
 {% else %}
 @compute @workgroup_size(WG)
-fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
   // One thread per query; the grid is 2D-folded so the workgroup count stays
-  // within maxComputeWorkgroupsPerDimension for large prefills. Reconstruct the
-  // flat thread index from the folded workgroup grid (wid.x + wid.y*nwg.x) — the
-  // global_invocation_id only spans nwg.x*WG and would drop the y-folded tail.
-  // This reduces to the ordinary flat index when nwg.y == 1.
-  let qi = (wid.x + wid.y * nwg.x) * WG + lid.x;
   let total = params.batch * Q_HEADS * params.qSeq;
   if (qi >= total) { return; }
   let s = qi % params.qSeq;
@@ -277,7 +280,9 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(num_workgroups) nwg: vec
     smoothTerm = exp(sink - finalM);
   }
   let accScale = exp(m - finalM);
-  let invDenom = accScale / (smoothTerm + l * accScale);
   let oBase = (b * params.qSeq + s) * Q_HIDDEN + h * HEAD_DIM;
   for (var d = 0u; d < HEAD_DIM; d = d + 1u) { output[oBase + d] = {{ IO }}(acc[d] * invDenom); }

 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>) {
   // One workgroup per query; fold large query grids over x/y.
+  let qi = wid.x + wid.y * {{ DISPATCH_FOLD_WIDTH }}u;
   let total = params.batch * Q_HEADS * params.qSeq;
   if (qi >= total) { return; }
   let lane = lid.x;
     smoothTerm = exp(sink - finalM);
   }
   let accScale = exp(gm - finalM);
+  // A query whose every key is masked out leaves the denominator at zero (no
+  // head sink to carry mass); emit a zero row instead of 0 * (1/0) = NaN.
+  let denom = smoothTerm + glsum * accScale;
+  let invDenom = select(0.0, accScale / denom, denom > 0.0);
   // Cross-lane sum of the rescaled V accumulators, written cooperatively.
   let oBase = (b * params.qSeq + s) * Q_HIDDEN + h * HEAD_DIM;
 }
 {% else %}
 @compute @workgroup_size(WG)
+fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
   // One thread per query; the grid is 2D-folded so the workgroup count stays
+  // within the per-axis dispatch fold width for large prefills. Reconstruct the flat
+  // thread index from the folded workgroup grid rather than from global_invocation_id,
+  // which only spans the clamped axis and would drop the y-folded tail. With no fold
+  // wid.y is zero, so this reduces to the ordinary flat index.
+  let qi = (wid.x + wid.y * {{ DISPATCH_FOLD_WIDTH }}u) * WG + lid.x;
   let total = params.batch * Q_HEADS * params.qSeq;
   if (qi >= total) { return; }
   let s = qi % params.qSeq;
     smoothTerm = exp(sink - finalM);
   }
   let accScale = exp(m - finalM);
+  // Same all-masked guard as the cooperative epilogue above.
+  let denom = smoothTerm + l * accScale;
+  let invDenom = select(0.0, accScale / denom, denom > 0.0);
   let oBase = (b * params.qSeq + s) * Q_HIDDEN + h * HEAD_DIM;
   for (var d = 0u; d < HEAD_DIM; d = d + 1u) { output[oBase + d] = {{ IO }}(acc[d] * invDenom); }

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

@@ -1,5 +1,5 @@
-{% if source is defined and source.mode is defined %}
-{% set mode = source.mode %}
 {% elif mode is not defined %}
 {% set mode = "transpose" %}
 {% endif %}
@@ -18,7 +18,8 @@ const WG: u32 = {{ copyWorkgroupSize }}u;
 {% if presentVec4 %}
 // Transpose copy [batch, token, head*headDim] -> present [batch, head, token,
 // headDim]. The headDim row is contiguous on both sides and divisible by four,
-// so each invocation copies a vec4 along d with coalesced stores.
 const HEAD_DIM_V4: u32 = {{ headDimV4 }}u;
 const KV_HIDDEN_V4: u32 = {{ kvHiddenV4 }}u;
 {% else %}
@@ -28,12 +29,11 @@ const KV_HIDDEN: u32 = {{ kvHidden }}u;
 @compute @workgroup_size(WG, 1, 1)
 fn main(
-  @builtin(global_invocation_id) gid: vec3<u32>,
-  @builtin(num_workgroups) nwg: 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 * nwg.x * WG;
 {% if presentVec4 %}
   let total = params.batchSize * KV_HEADS * params.kvSeq * HEAD_DIM_V4;
   if (index >= total) {
@@ -111,11 +111,9 @@ fn vscale(d: u32, hk: u32) -> f32 { return v_scale[select(0u, hk * HEAD_DIM + d,
 {% endif %}
 @compute @workgroup_size(WG)
-fn main({% if needsGid %}@builtin(global_invocation_id) gid: vec3<u32>,
-        {% endif %}{% if cooperativeMode %}@builtin(workgroup_id) wid: vec3<u32>,
-        @builtin(local_invocation_id) lid: vec3<u32>,
-        {% endif %}
-        @builtin(num_workgroups) nwg: vec3<u32>) {
 {% macro append_cooperative_walk() %}
   // Compact cooperative append walk: the thread range spans only the appended
   // window (batch x kvHead x keySeq rows), not the full cache capacity; the
@@ -155,7 +153,7 @@ fn main({% if needsGid %}@builtin(global_invocation_id) gid: vec3<u32>,
 {% endmacro %}
 {% macro append_per_row_walk() %}
   // Per-row append walk: one thread streams its row's contiguous bytes.
-  let i = gid.x + gid.y * nwg.x * WG;
   if (i >= params.count) { return; }
   let j = i % params.keySeq;
   let tmp = i / params.keySeq;
@@ -195,7 +193,7 @@ fn main({% if needsGid %}@builtin(global_invocation_id) gid: vec3<u32>,
   // cooperatively: at step s lane l touches element s*WG + l, so adjacent
   // lanes hit adjacent addresses on both source and destination. Same
   // dispatch geometry (ceil(rows/WG) workgroups, folded into x/y).
-  let wgFlat = wid.x + wid.y * nwg.x;
   let elemBase = wgFlat * (WG * HEAD_DIM);
   let totalElems = params.count * HEAD_DIM;
 {% endif %}
@@ -265,7 +263,7 @@ fn main({% if needsGid %}@builtin(global_invocation_id) gid: vec3<u32>,
 {% endif %}
 {% elif mode == "copy" %}
   // Thread i streams one contiguous (batch, kvHead, token) row.
-  let i = gid.x + gid.y * nwg.x * WG;
   if (i >= params.count) { return; }
   let base = i * HEAD_DIM;
   for (var d = 0u; d < HEAD_DIM; d = d + 1u) {
@@ -274,7 +272,7 @@ fn main({% if needsGid %}@builtin(global_invocation_id) gid: vec3<u32>,
   }
 {% elif mode == "merge" %}
   // Per-row merge walk.
-  let i = gid.x + gid.y * nwg.x * WG;
   if (i >= params.count) { return; }
   let t = i % params.seq;
   let tmp = i / params.seq;
@@ -298,7 +296,7 @@ fn main({% if needsGid %}@builtin(global_invocation_id) gid: vec3<u32>,
 {% 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 * nwg.x * WG;
   if (i >= params.count) { return; }
   let t = i % params.seq;
   let tmp = i / params.seq;
@@ -338,7 +336,7 @@ fn main({% if needsGid %}@builtin(global_invocation_id) gid: vec3<u32>,
   // 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 * nwg.x * WG;
   if (i >= params.count) { return; }
   let t = i % params.seq;
   let tmp = i / params.seq;
@@ -377,7 +375,7 @@ fn main({% if needsGid %}@builtin(global_invocation_id) gid: vec3<u32>,
 {% 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 * nwg.x * WG;
   if (i >= params.count) { return; }
   let t = i % params.seq;
   let tmp = i / params.seq;

+{% if modeSpec is defined %}
+{% set mode = modeSpec %}
 {% elif mode is not defined %}
 {% set mode = "transpose" %}
 {% endif %}
 {% if presentVec4 %}
 // Transpose copy [batch, token, head*headDim] -> present [batch, head, token,
 // headDim]. The headDim row is contiguous on both sides and divisible by four,
+// so each invocation copies one vec4 along d and adjacent invocations write
+// adjacent destination vectors.
 const HEAD_DIM_V4: u32 = {{ headDimV4 }}u;
 const KV_HIDDEN_V4: u32 = {{ kvHiddenV4 }}u;
 {% else %}
 @compute @workgroup_size(WG, 1, 1)
 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) {
 {% endif %}
 @compute @workgroup_size(WG)
+fn main({% if needsGid %}@builtin(global_invocation_id) gid: vec3<u32>{% if cooperativeMode %},
+        {% endif %}{% endif %}{% if cooperativeMode %}@builtin(workgroup_id) wid: vec3<u32>,
+        @builtin(local_invocation_id) lid: vec3<u32>{% endif %}) {
 {% macro append_cooperative_walk() %}
   // Compact cooperative append walk: the thread range spans only the appended
   // window (batch x kvHead x keySeq rows), not the full cache capacity; the
 {% 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;
   // cooperatively: at step s lane l touches element s*WG + l, so adjacent
   // lanes hit adjacent addresses on both source and destination. Same
   // dispatch geometry (ceil(rows/WG) workgroups, folded into x/y).
+  let wgFlat = wid.x + wid.y * {{ DISPATCH_FOLD_WIDTH }}u;
   let elemBase = wgFlat * (WG * HEAD_DIM);
   let totalElems = params.count * HEAD_DIM;
 {% endif %}
 {% 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) {
   }
 {% 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;
 {% 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;
   // 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;
 {% 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;

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

@@ -1,18 +1,12 @@
 {% if usesF16 is defined and usesF16 %}enable f16;
 {% endif %}{{ env.wgsl.resourceDeclarations }}
-// Q preprocessing for the flash-prefill path. One thread per (batch, qHead, qToken)
-// applies optional per-head qk-norm (RMS) then NeoX half-split rotary at the query's
-// absolute position, and writes a transformed Q (vec4, same BSH layout) that the
-// flash cluster then consumes unchanged. This mirrors the cooperative/threaded path's Q-prep
-// exactly (the present-cache K is already norm/rotary-applied, so only Q is touched),
-// so the flash result is bit-identical to the cooperative/threaded path.
-//
-// f16 queries are widened before cos/sin and norm/rotary arithmetic, and the
-// rotation itself always runs in f32 — RoPE precision matters at large
-// positions. The store narrows only for a consumer whose operand tiles are
-// f16 anyway (the subgroup-matrix route), where a wider intermediate would be
-// narrowed at staging regardless; every other consumer keeps the f32 output.
 const HEAD_DIM: u32 = {{ headDim }}u;
 const HEAD_DIM_V4: u32 = {{ headDim }}u / 4u;
 const Q_HEADS: u32 = {{ qHeads }}u;
@@ -26,11 +20,10 @@ const QK_EPS: f32 = {{ qkEps }};
 {% endif %}
 @compute @workgroup_size(WG)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>,
-        @builtin(num_workgroups) nwg: vec3<u32>) {
-  // 2D-folded flat index: gid.y carries the high bits past the maxComputeWorkgroupsPerDimension
-  // workgroup-per-dimension dispatch limit. Reduces to gid.x when nwg.y == 1.
-  let qi = gid.x + gid.y * nwg.x * WG;
   let total = params.batch * Q_HEADS * params.qSeq;
   if (qi >= total) { return; }
   let s = qi % params.qSeq;

 {% 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
+// absolute position. Present-cache keys are already transformed, so this pass
+// writes only Q in BSH layout. f16 queries are widened for normalization,
+// trigonometry, and rotation. The result narrows only when the following
+// attention pass consumes f16 operand tiles; other configurations store f32.
 const HEAD_DIM: u32 = {{ headDim }}u;
 const HEAD_DIM_V4: u32 = {{ headDim }}u / 4u;
 const Q_HEADS: u32 = {{ qHeads }}u;
 {% endif %}
 @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;

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,29 +1,99 @@
 {
   "name": "com.microsoft.GroupQueryAttention",
-  "id": "_com_microsoft_groupqueryattention_webgpu_34ecc1b",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "attention-rank4-tiled.wgsl.jinja": "EmYf/jV5lG4hmXr5wP+Qo6hZlY3LDd1aeAQhFXtDJEc=",
-      "attn-flash-decode-splitk-merge.wgsl.jinja": "kQDmryD3bdfkboFSLPK4pF2D40t3MX2xgJrMHx031yc=",
-      "attn-flash-decode-splitk.wgsl.jinja": "M6AWhYqX3pJ/MHz2K6Vy1rx7RjaZktzgQqB9GoY4cj0=",
-      "attn-flash-online.wgsl.jinja": "MStTMF1/l5L9Rzx1+rPz5T+70/46n/vO/+hZQ0lbFws=",
-      "attn-flash-prefill-cluster.wgsl.jinja": "yk1WymJIbxCy0slbiho5DS223uZH9PQ15629yvQ9RFc=",
-      "attn-flash-q32-broadcast.wgsl.jinja": "evSdaiYgvrikYZ8hpb5m4gxf+rIryKkHOTiKgyF7ioI=",
-      "attn-materialized-rowstats-combine-f32.wgsl.jinja": "TtTBeeeP/ufhzeNHoeLWGLxKTfFT9fiEAstnXB79j48=",
-      "attn-materialized-sgmat-f32.wgsl.jinja": "iQeacl5PUCl75lx78qlSBaQ/syBDNxgjZg4UPGki9U0=",
-      "attn-online-scalar.wgsl.jinja": "X4yGNCYu9+uJX6Q3EKkxV6v28CPlIYWme0Nmf8F16L8=",
-      "bench.json": "b4ZLhmCqTm6AgiNfIiRHYVeNp4q1mx12GMcDAextGqI=",
-      "gqa-attention.wgsl.jinja": "0EbwmdstNZMuskS170zYEo8whbQOFtupXLnNP2PaFL0=",
-      "gqa-present.wgsl.jinja": "J9g6oAeY2BtbmWBMlmbIXbxzo4sKHdxYqcE4EAd7BYM=",
-      "gqa-qprep.wgsl.jinja": "RCBDy7LA/yy4ayRAsnyq+7/IaXjsDGmPWidj51c+t8k=",
-      "manifest.json": "oxSz/yE5vrd9vU5lHkrRN2kG8sdBzWQmGP9/jjRoqNU=",
-      "test.json": "zfjYb8H4THOE5qN8K90UPdGq0lL/B2LWarbvokAz8Cg="
     }
   },
-  "provenance": { "kernel": { "sha": "c928d21e6cc1310861cba3bafb75f5f679ecf5f3", "dirty": false } },
-  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.GroupQueryAttention" }
 }

 {
   "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"],
+      "past_kv_bias_materialized_sgmat_f32": ["attn-materialized-rowstats-combine-f32.wgsl.jinja", "attn-materialized-sgmat-f32.wgsl.jinja", "gqa-present.wgsl.jinja"],
+      "past_kv_materialized_sgmat_f16": ["attn-materialized-rowstats-combine-f32.wgsl.jinja", "attn-materialized-sgmat-f32.wgsl.jinja", "gqa-present.wgsl.jinja"],
+      "past_kv_rotary_materialized_sgmat_f32": ["attn-materialized-rowstats-combine-f32.wgsl.jinja", "attn-materialized-sgmat-f32.wgsl.jinja", "gqa-present.wgsl.jinja", "gqa-qprep.wgsl.jinja"],
+      "past_kv_rotary_materialized_sgmat_f16": ["attn-materialized-rowstats-combine-f32.wgsl.jinja", "attn-materialized-sgmat-f32.wgsl.jinja", "gqa-present.wgsl.jinja", "gqa-qprep.wgsl.jinja"],
+      "new_kv_past_materialized_sgmat_f32": ["attn-materialized-rowstats-combine-f32.wgsl.jinja", "attn-materialized-sgmat-f32.wgsl.jinja", "gqa-present.wgsl.jinja"],
+      "window_shift_materialized_sgmat_f32": ["attn-materialized-rowstats-combine-f32.wgsl.jinja", "attn-materialized-sgmat-f32.wgsl.jinja", "gqa-present.wgsl.jinja"],
+      "share_append_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"],
+      "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"],
+      "qkv_present_flash_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"],
+      "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"],
+      "past_kv_softcap_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"],
+      "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"]
+    }
+  }
 }

build/webgpu/test.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "com.microsoft.GroupQueryAttention",
   "fixtureArrays": {
     "ort_past_key_t_pattern": [0.2, 0.4, 0.6, 0.8, 1, 0.2, 0.4, 0.6, 0.8, 1, 0.2, 0.4, 0.6, 0.8, 1, 0.2, 0.4, 0.6, 0.8, 1, 0.2, 0.4, 0.6, 0.8, 1, 0.2, 0.4, 0.6, 0.8, 1, 0.2, 0.4, 0.6, 0.8, 1, 0.2, 0.4, 0.6, 0.8, 1, 0.2, 0.4, 0.6, 0.8, 1, 0.2, 0.4, 0.6, 0.8, 1, 0.2, 0.4, 0.6, 0.8, 1, 0.2, 0.4, 0.6, 0.8, 1, 0.2, 0.4, 0.6, 0.8],
     "ort_past_value_t_pattern": [0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3],
@@ -247,7 +246,7 @@
         "valueT": {
           "dtype": "float32",
           "shape": [1, 512, 64],
-          "data": { "kind": "fillFloat32", "scale": 0.18, "sinStep": 0.23, "cosStep": 0.17 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [512] } }
@@ -467,7 +466,7 @@
         "valueT": {
           "dtype": "float32",
           "shape": [1, 2048, 128],
-          "data": { "kind": "fillFloat32", "scale": 0.18, "sinStep": 0.023, "cosStep": 0.017 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [2047] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [2048] } }
@@ -495,7 +494,7 @@
         "valueT": {
           "dtype": "float32",
           "shape": [1, 1024, 64],
-          "data": { "kind": "fillFloat32", "scale": 0.18, "sinStep": 0.023, "cosStep": 0.017 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1023] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1024] } }
@@ -523,15 +522,15 @@
         "valueT": {
           "dtype": "float16",
           "shape": [1, 128, 64],
-          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.23, "cosStep": 0.17 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [127] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [128] } }
       },
       "outputs": {
-        "outputT": { "dtype": "float16", "shape": [1, 128, 128], "tolerance": 0.03 },
-        "presentKeyT": { "dtype": "float16", "shape": [1, 1, 128, 64], "tolerance": 0.03 },
-        "presentValueT": { "dtype": "float16", "shape": [1, 1, 128, 64], "tolerance": 0.03 }
       }
     },
     {
@@ -551,15 +550,15 @@
         "valueT": {
           "dtype": "float16",
           "shape": [1, 512, 64],
-          "data": { "kind": "fillFloat32", "scale": 0.18, "sinStep": 0.23, "cosStep": 0.17 }
         },
         "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, 1, 256], "tolerance": 0.03 },
-        "presentKeyT": { "dtype": "float16", "shape": [1, 1, 512, 64], "tolerance": 0.03 },
-        "presentValueT": { "dtype": "float16", "shape": [1, 1, 512, 64], "tolerance": 0.03 }
       }
     },
     {
@@ -579,15 +578,15 @@
         "valueT": {
           "dtype": "float16",
           "shape": [1, 2048, 128],
-          "data": { "kind": "fillFloat32", "scale": 0.18, "sinStep": 0.023, "cosStep": 0.017 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [2047] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [2048] } }
       },
       "outputs": {
-        "outputT": { "dtype": "float16", "shape": [1, 1, 512], "tolerance": 0.03 },
-        "presentKeyT": { "dtype": "float16", "shape": [1, 2, 2048, 64], "tolerance": 0.03 },
-        "presentValueT": { "dtype": "float16", "shape": [1, 2, 2048, 64], "tolerance": 0.03 }
       }
     },
     {
@@ -691,7 +690,7 @@
         "valueT": {
           "dtype": "float32",
           "shape": [1, 1024, 128],
-          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.0023, "cosStep": 0.0013 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1023] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1024] } }
@@ -1210,7 +1209,7 @@
         "vScaleT": { "dtype": "float32", "shape": [1], "data": { "kind": "values", "values": [0.0038802] } }
       },
       "outputs": {
-        "outputT": { "dtype": "float32", "shape": [1, 4, 16], "tolerance": 0.05 },
         "presentKeyT": { "dtype": "int8", "shape": [1, 1, 4, 8], "tolerance": 0 },
         "presentValueT": { "dtype": "int8", "shape": [1, 1, 4, 8], "tolerance": 0 }
       }
@@ -1276,7 +1275,7 @@
         }
       },
       "outputs": {
-        "outputT": { "dtype": "float32", "shape": [1, 4, 16], "tolerance": 0.05 },
         "presentKeyT": { "dtype": "int8", "shape": [1, 1, 4, 8], "tolerance": 0 },
         "presentValueT": { "dtype": "int8", "shape": [1, 1, 4, 8], "tolerance": 0 }
       }
@@ -1381,7 +1380,7 @@
         "vScaleT": { "dtype": "float32", "shape": [1], "data": { "kind": "values", "values": [0.0703979] } }
       },
       "outputs": {
-        "outputT": { "dtype": "float32", "shape": [1, 4, 16], "tolerance": 0.15 },
         "presentKeyT": { "dtype": "uint8", "shape": [1, 1, 4, 4], "tolerance": 0 },
         "presentValueT": { "dtype": "uint8", "shape": [1, 1, 4, 4], "tolerance": 0 }
       }
@@ -1447,7 +1446,7 @@
         }
       },
       "outputs": {
-        "outputT": { "dtype": "float32", "shape": [1, 4, 16], "tolerance": 0.15 },
         "presentKeyT": { "dtype": "uint8", "shape": [1, 1, 4, 4], "tolerance": 0 },
         "presentValueT": { "dtype": "uint8", "shape": [1, 1, 4, 4], "tolerance": 0 }
       }
@@ -1508,7 +1507,7 @@
         "vScaleT": { "dtype": "float32", "shape": [1], "data": { "kind": "values", "values": [0.0039328] } }
       },
       "outputs": {
-        "outputT": { "dtype": "float32", "shape": [2, 4, 64], "tolerance": 0.05 },
         "presentKeyT": { "dtype": "int8", "shape": [2, 2, 4, 16], "tolerance": 0 },
         "presentValueT": { "dtype": "int8", "shape": [2, 2, 4, 16], "tolerance": 0 }
       }
@@ -1563,7 +1562,7 @@
         "vScaleT": { "dtype": "float32", "shape": [1], "data": { "kind": "values", "values": [0.0039328] } }
       },
       "outputs": {
-        "outputT": { "dtype": "float32", "shape": [1, 8, 128], "tolerance": 0.05 },
         "presentKeyT": { "dtype": "int8", "shape": [1, 1, 8, 64], "tolerance": 0 },
         "presentValueT": { "dtype": "int8", "shape": [1, 1, 8, 64], "tolerance": 0 }
       }
@@ -1624,7 +1623,7 @@
         "vScaleT": { "dtype": "float32", "shape": [1], "data": { "kind": "values", "values": [0.0703979] } }
       },
       "outputs": {
-        "outputT": { "dtype": "float32", "shape": [1, 4, 64], "tolerance": 0.15 },
         "presentKeyT": { "dtype": "uint8", "shape": [1, 1, 4, 8], "tolerance": 0 },
         "presentValueT": { "dtype": "uint8", "shape": [1, 1, 4, 8], "tolerance": 0 }
       }
@@ -1803,7 +1802,7 @@
         "vScaleT": { "dtype": "float32", "shape": [1], "data": { "kind": "values", "values": [0.0038802] } }
       },
       "outputs": {
-        "outputT": { "dtype": "float32", "shape": [1, 4, 16], "tolerance": 0.05 },
         "presentKeyT": { "dtype": "int8", "shape": [1, 1, 4, 8], "tolerance": 0 },
         "presentValueT": { "dtype": "int8", "shape": [1, 1, 4, 8], "tolerance": 0 }
       }
@@ -1849,6 +1848,102 @@
         "presentValueT": { "dtype": "float32", "shape": [1, 1, 8, 8], "tolerance": 0.0001 }
       }
     },
     {
       "name": "ort_bias_bcast_prompt_b1q8p8_h2kv1d8",
       "attrs": { "num_heads": 2, "kv_num_heads": 1 },
@@ -2002,7 +2097,7 @@
     {
       "name": "flashprefill_pastzero_h8kv2_d128_q16",
       "provenance": {
-        "notes": "Compact companion for bench case sharedkv-prefill-32h8kv-d128-s64, preserving shared-KV prefill with headDim=128."
       },
       "attrs": { "num_heads": 8, "kv_num_heads": 2, "scale": 0.08838834764831845 },
       "inputs": {
@@ -2051,7 +2146,7 @@
         "pastValueT": {
           "dtype": "float32",
           "shape": [1, 2, 48, 64],
-          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [47] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [48] } }
@@ -2136,7 +2231,7 @@
       "name": "qnorm_rotary_scalar_h8kv2_d256_q1",
       "provenance": {
         "source": "scope coverage",
-        "test": "rendered-Jinja scope coverage",
         "notes": "headDim=256 requires 34,048 bytes for the cooperative path, exceeding a 32 KiB workgroup-storage tier and exercising the thread-per-query Q-norm route."
       },
       "attrs": { "num_heads": 8, "kv_num_heads": 2, "scale": 0.0625, "do_rotary": 1, "qk_norm_epsilon": 0.00001 },
@@ -2191,7 +2286,7 @@
       "name": "f16_qnorm_rotary_scalar_f16wts_h8kv2_d256_q1",
       "provenance": {
         "source": "scope coverage",
-        "test": "rendered-Jinja scope coverage",
         "notes": "Float16 counterpart to the thread-per-query Q-norm case. It verifies explicit f32 conversion of q_norm_weight, cos_cache, and sin_cache before scalar arithmetic."
       },
       "attrs": { "num_heads": 8, "kv_num_heads": 2, "scale": 0.0625, "do_rotary": 1, "qk_norm_epsilon": 0.00001 },
@@ -2311,7 +2406,7 @@
         "pastValueT": {
           "dtype": "float32",
           "shape": [1, 2, 1024, 64],
-          "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] } }
@@ -2325,7 +2420,7 @@
     {
       "name": "flashdecode_splitk_h8kv2_d128_kv1024",
       "provenance": {
-        "notes": "Compact companion for bench case sharedkv-decode-32h8kv-d128-kv2048, preserving f32 shared-past decode with headDim=128 and the split-K route."
       },
       "attrs": { "num_heads": 8, "kv_num_heads": 2, "scale": 0.08838834764831845 },
       "inputs": {
@@ -2344,7 +2439,7 @@
         "pastValueT": {
           "dtype": "float32",
           "shape": [1, 2, 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] } }
@@ -2374,13 +2469,13 @@
         "pastValueT": {
           "dtype": "float16",
           "shape": [1, 2, 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, 1, 1024], "tolerance": 0.03 },
         "presentKeyT": { "dtype": "float16", "shape": [1, 2, 512, 128], "tolerance": 0.001 },
         "presentValueT": { "dtype": "float16", "shape": [1, 2, 512, 128], "tolerance": 0.001 }
       }
@@ -2404,13 +2499,13 @@
         "pastValueT": {
           "dtype": "float16",
           "shape": [1, 2, 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": "float16", "shape": [1, 1, 1024], "tolerance": 0.03 },
         "presentKeyT": { "dtype": "float16", "shape": [1, 2, 1024, 128], "tolerance": 0.001 },
         "presentValueT": { "dtype": "float16", "shape": [1, 2, 1024, 128], "tolerance": 0.001 }
       }
@@ -2464,7 +2559,7 @@
         "pastValueT": {
           "dtype": "float16",
           "shape": [1, 2, 512, 64],
-          "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] } },
@@ -2480,7 +2575,7 @@
         }
       },
       "outputs": {
-        "outputT": { "dtype": "float16", "shape": [1, 1, 512], "tolerance": 0.03 },
         "presentKeyT": { "dtype": "float16", "shape": [1, 2, 512, 64], "tolerance": 0.001 },
         "presentValueT": { "dtype": "float16", "shape": [1, 2, 512, 64], "tolerance": 0.001 }
       }
@@ -2504,7 +2599,7 @@
         "pastValueT": {
           "dtype": "float16",
           "shape": [1, 2, 32, 64],
-          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [31] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [32] } },
@@ -2544,7 +2639,7 @@
         "pastValueT": {
           "dtype": "float16",
           "shape": [1, 2, 32, 64],
-          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [31] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [32] } },
@@ -2630,7 +2725,7 @@
         "valueT": {
           "dtype": "float32",
           "shape": [1, 32, 128],
-          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.019, "cosStep": 0.023 }
         },
         "pastKeyT": {
           "dtype": "float32",
@@ -2640,7 +2735,7 @@
         "pastValueT": {
           "dtype": "float32",
           "shape": [1, 2, 32, 64],
-          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [63] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [64] } }
@@ -2668,7 +2763,7 @@
         "valueT": {
           "dtype": "float32",
           "shape": [1, 1, 128],
-          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.019, "cosStep": 0.023 }
         },
         "pastKeyT": {
           "dtype": "float32",
@@ -2678,7 +2773,7 @@
         "pastValueT": {
           "dtype": "float32",
           "shape": [1, 2, 1024, 64],
-          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1024] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1025] } }
@@ -2706,7 +2801,7 @@
         "valueT": {
           "dtype": "float16",
           "shape": [1, 32, 128],
-          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.019, "cosStep": 0.023 }
         },
         "pastKeyT": {
           "dtype": "float16",
@@ -2716,13 +2811,13 @@
         "pastValueT": {
           "dtype": "float16",
           "shape": [1, 2, 32, 64],
-          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
         },
         "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, 32, 512], "tolerance": 0.03 },
         "presentKeyT": { "dtype": "float16", "shape": [1, 2, 64, 64], "tolerance": 0.001 },
         "presentValueT": { "dtype": "float16", "shape": [1, 2, 64, 64], "tolerance": 0.001 }
       }
@@ -2872,7 +2967,7 @@
         "valueT": {
           "dtype": "float32",
           "shape": [1, 511, 128],
-          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.007, "cosStep": 0.041 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [510] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } }
@@ -2901,7 +2996,7 @@
         "valueT": {
           "dtype": "float32",
           "shape": [1, 64, 256],
-          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.007, "cosStep": 0.041 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [63] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [64] } }
@@ -2932,7 +3027,7 @@
         "valueT": {
           "dtype": "float32",
           "shape": [1, 512, 256],
-          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.029, "cosStep": 0.017 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [512] } }
@@ -3063,7 +3158,7 @@
     {
       "name": "qkv_prefill_f32_h8kv2_d256_register_boundary",
       "provenance": {
-        "notes": "Upper-head routing boundary for QKV prefill. On subgroup tiers the generic online kernel beats both register-heavy Q32 broadcast and the clustered d256 topology; no-subgroup tiers retain the shared-memory cluster as their fastest fallback."
       },
       "attrs": { "num_heads": 8, "kv_num_heads": 2, "scale": 0.0625, "causal": 0 },
       "inputs": {
@@ -3261,7 +3356,7 @@
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/bert/flash_attention_decode_qkv.wgsl.template",
         "test": "head_size=8 split-reduce OOB race class (PR #29593)",
-        "notes": "head_size below the flash minimum routes to the cooperative fallback; this pins the tiny-headDim decode regime where tile-width-guarded reductions have raced in other engines."
       },
       "attrs": { "num_heads": 4, "kv_num_heads": 2 },
       "inputs": {
@@ -3280,7 +3375,7 @@
         "pastValueT": {
           "dtype": "float32",
           "shape": [1, 2, 512, 8],
-          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.31, "scale": 0.5 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [512] } }
@@ -3296,7 +3391,7 @@
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/bert/flash_attention_decode_qkv.wgsl.template",
         "test": "head_size=16 split-reduce OOB race class (PR #29593)",
-        "notes": "Sibling of the head_size=8 case at the next tiny head width."
       },
       "attrs": { "num_heads": 4, "kv_num_heads": 2 },
       "inputs": {
@@ -3315,7 +3410,7 @@
         "pastValueT": {
           "dtype": "float32",
           "shape": [1, 2, 512, 16],
-          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.31, "scale": 0.5 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [512] } }
@@ -3541,7 +3636,7 @@
       "provenance": {
         "source": "onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc",
         "test": "GQA head sink on the buffer-sharing append path",
-        "notes": "CASE: the decode step of the same geometry -- one new K/V row appended at 11 with 20 rows of capacity slack the attention bound must exclude, so seqlens_k rather than the bound capacity sets both the append offset and the window origin."
       }
     },
     {
@@ -3627,7 +3722,7 @@
     {
       "name": "prefill_tiled_q31_h8kv2_d36_non_cluster_stride",
       "provenance": {
-        "notes": "Head size 36 is a multiple of 4 but not of 4*PREFILL_LANES_PER_QUERY=16, so the cooperative cluster prefill cannot map its per-query lanes; qSeq 31 clears the prefill minimum and qSeq*heads=248 clears the flash occupancy floor. Locks the blocked tiled prefill kernel on tiers without subgroups."
       },
       "attrs": { "num_heads": 8, "kv_num_heads": 2, "scale": 0.16666666666666666, "causal": 0 },
       "inputs": {
@@ -3685,7 +3780,7 @@
         "presentValueT": { "dtype": "float32", "shape": [1, 2, 32, 64], "tolerance": 0.0001 }
       },
       "provenance": {
-        "notes": "Logit soft-capping on the past_kv prefill path. Every flash prefill variant used to gate `standardSoftmax`, so a softcap sent the whole prefill to the scalar kernel at ~3% of peak; capping is a per-score transform, so the tiled kernel takes it unchanged."
       }
     },
     {
@@ -3772,7 +3867,7 @@
     {
       "name": "flashprefill_quant_int8_h4kv2_d64_q32",
       "provenance": {
-        "notes": "INT8 KV cache on the tiled prefill path. The scalar kernel dequantized a cache element per (query, key); staging the tile dequantizes once per key and every query in the tile reads the result, so the unpack cost divides by the tile height."
       },
       "attrs": {
         "num_heads": 4,
@@ -3883,7 +3978,7 @@
       "provenance": {
         "source": "onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc",
         "test": "GQA sliding_window_cache (ORT #29904), CPU/CUDA reference semantics",
-        "notes": "Windowed KV cache: the bound past/present buffer is a fixed CAPACITY holding only the most recent min(T, C) tokens contiguously at rows [0, L). Appending S tokens onto past P evicts E = max(0, min(P,C) + S - C) rows and slides survivors down by E. Attention reads keys at absolute positions translated by the cache origin T - L; RoPE keeps absolute positions. ORT's own WebGPU EP rejects this attribute (CUDA/CPU only). CASE: T=6 <= C=8, E=0, origin=0 — degenerates to append; also pins the unread-row clear."
       }
     },
     {
@@ -3918,7 +4013,7 @@
       "provenance": {
         "source": "onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc",
         "test": "GQA sliding_window_cache (ORT #29904), CPU/CUDA reference semantics",
-        "notes": "Windowed KV cache: the bound past/present buffer is a fixed CAPACITY holding only the most recent min(T, C) tokens contiguously at rows [0, L). Appending S tokens onto past P evicts E = max(0, min(P,C) + S - C) rows and slides survivors down by E. Attention reads keys at absolute positions translated by the cache origin T - L; RoPE keeps absolute positions. ORT's own WebGPU EP rejects this attribute (CUDA/CPU only). CASE: T=12 > C=8, E=1, appendStart=7, origin=4 — the compaction shift."
       }
     },
     {
@@ -3953,7 +4048,7 @@
       "provenance": {
         "source": "onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc",
         "test": "GQA sliding_window_cache (ORT #29904), CPU/CUDA reference semantics",
-        "notes": "Windowed KV cache: the bound past/present buffer is a fixed CAPACITY holding only the most recent min(T, C) tokens contiguously at rows [0, L). Appending S tokens onto past P evicts E = max(0, min(P,C) + S - C) rows and slides survivors down by E. Attention reads keys at absolute positions translated by the cache origin T - L; RoPE keeps absolute positions. ORT's own WebGPU EP rejects this attribute (CUDA/CPU only). CASE: S=4, T=10 > C=8, E=2, appendStart=4, origin=2 — multi-row eviction."
       }
     },
     {
@@ -3988,7 +4083,7 @@
       "provenance": {
         "source": "onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc",
         "test": "GQA sliding_window_cache (ORT #29904), CPU/CUDA reference semantics",
-        "notes": "Windowed KV cache: the bound past/present buffer is a fixed CAPACITY holding only the most recent min(T, C) tokens contiguously at rows [0, L). Appending S tokens onto past P evicts E = max(0, min(P,C) + S - C) rows and slides survivors down by E. Attention reads keys at absolute positions translated by the cache origin T - L; RoPE keeps absolute positions. ORT's own WebGPU EP rejects this attribute (CUDA/CPU only). CASE: b0 T=4 (E=0, origin=0), b1 T=12 (E=1, origin=4) — per-batch origin."
       }
     },
     {
@@ -4023,7 +4118,7 @@
       "provenance": {
         "source": "onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc",
         "test": "GQA sliding_window_cache (ORT #29904), CPU/CUDA reference semantics",
-        "notes": "Windowed KV cache: the bound past/present buffer is a fixed CAPACITY holding only the most recent min(T, C) tokens contiguously at rows [0, L). Appending S tokens onto past P evicts E = max(0, min(P,C) + S - C) rows and slides survivors down by E. Attention reads keys at absolute positions translated by the cache origin T - L; RoPE keeps absolute positions. ORT's own WebGPU EP rejects this attribute (CUDA/CPU only). CASE: w=4 <= C=8 with E=1 — mask floor and cache origin interact."
       }
     },
     {
@@ -4058,7 +4153,7 @@
       "provenance": {
         "source": "onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc",
         "test": "GQA sliding_window_cache (ORT #29904), CPU/CUDA reference semantics",
-        "notes": "Windowed KV cache: the bound past/present buffer is a fixed CAPACITY holding only the most recent min(T, C) tokens contiguously at rows [0, L). Appending S tokens onto past P evicts E = max(0, min(P,C) + S - C) rows and slides survivors down by E. Attention reads keys at absolute positions translated by the cache origin T - L; RoPE keeps absolute positions. ORT's own WebGPU EP rejects this attribute (CUDA/CPU only). CASE: C == w == 4, T=10 — the shape a sliding-window layer actually allocates."
       }
     },
     {
@@ -4093,7 +4188,7 @@
       "provenance": {
         "source": "onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc",
         "test": "GQA sliding_window_cache (ORT #29904), CPU/CUDA reference semantics",
-        "notes": "Windowed KV cache: the bound past/present buffer is a fixed CAPACITY holding only the most recent min(T, C) tokens contiguously at rows [0, L). Appending S tokens onto past P evicts E = max(0, min(P,C) + S - C) rows and slides survivors down by E. Attention reads keys at absolute positions translated by the cache origin T - L; RoPE keeps absolute positions. ORT's own WebGPU EP rejects this attribute (CUDA/CPU only). CASE: T=64 with C=16 — origin 48, far past the buffer."
       }
     },
     {
@@ -4115,7 +4210,7 @@
         "pastValueT": {
           "dtype": "float32",
           "shape": [1, 1, 1024, 64],
-          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.13, "scale": 0.5 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [4095] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [4096] } }
@@ -4128,7 +4223,7 @@
       "provenance": {
         "source": "onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc",
         "test": "GQA sliding_window_cache (ORT #29904), CPU/CUDA reference semantics",
-        "notes": "Windowed cache at a capacity that clears CACHED_DECODE_MIN_KV_TOKENS, so the split-K decode variant is selected rather than the scalar fallback. T=4096 > C=1024, so the step genuinely evicts and the split-K kernel reads a shifted cache."
       }
     },
     {
@@ -4502,7 +4597,7 @@
       "provenance": {
         "source": "onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc",
         "test": "GQA buffer-sharing append chunk prefill on the flash cluster kernel",
-        "notes": "Share-append binds one fixed-capacity cache for past and present, so the resident length is seqlens_k + 1 rather than the bound capacity — the same split between attention bound and batch stride the windowed cache needs. headDim >= 64 with qSeq >= PREFILL_QUERY_TILE routes these to the flash prefill kernel. CASE: T=48 < C=64 — 16 rows of capacity slack the flash bound must exclude."
       }
     },
     {
@@ -4574,7 +4669,206 @@
       "provenance": {
         "source": "onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc",
         "test": "GQA buffer-sharing append chunk prefill on the flash cluster kernel",
-        "notes": "Share-append binds one fixed-capacity cache for past and present, so the resident length is seqlens_k + 1 rather than the bound capacity — the same split between attention bound and batch stride the windowed cache needs. headDim >= 64 with qSeq >= PREFILL_QUERY_TILE routes these to the flash prefill kernel. CASE: b0 T=32 (append fills exactly), b1 T=48 — bound is per batch."
       }
     },
     {
@@ -4611,7 +4905,7 @@
     {
       "name": "splitk_decode_h8kv2_d64_kv2048_f16_ramp_value_scale_lock",
       "provenance": {
-        "notes": "The f16 compile of the present-KV decode split-K kernel had exactly one fixture, splitk_decode_h8kv2_d64_kv2048_f16, and its 0.03 absolute tolerance against a 2.1e-4 expected output is blind to a 146x scale error - the worst ratio in either attention op. Same shape and route with a ramped V so the f16 accumulator, its rescale and the final divide are all scale-locked."
       },
       "attrs": { "num_heads": 8, "kv_num_heads": 2, "scale": 0.125, "causal": 0 },
       "inputs": {
@@ -4642,7 +4936,7 @@
     {
       "name": "flash_prefill_h2kv1_d64_s128_f16_ramp_value_scale_lock",
       "provenance": {
-        "notes": "qkv_present_flash_q32_broadcast and qkv_present_flash_q32_shared are reached by one fixture only, flash_prefill_h2kv1_d64_s128_f16, whose zero-mean V averages to 2.3e-3 under a 0.03 absolute tolerance - blind to a 13x scale error. Ramped V makes every causal row an O(1) function of its own weighted mean key index, so the register-blocked q32 epilogue divide is under test and each of the 128 query rows carries a different expected value."
       },
       "attrs": { "num_heads": 2, "kv_num_heads": 1, "scale": 0.125, "causal": 0 },
       "inputs": {
@@ -4673,7 +4967,7 @@
     {
       "name": "newkv_past_splitk_h8kv2_d64_q1p1024_ramp_value_scale_lock",
       "provenance": {
-        "notes": "new_kv_past_decode_splitk and its _nosg twin - the append-new-KV-then-decode route - had one fixture, newkv_past_splitk_h8kv2_d64_q1p1024, whose zero-mean past and new V average to 4.7e-4 under a 0.005 tolerance: a 10x scale error passes. Ramping both value tensors makes the output an O(1) function of the weighted mean key index across the final 256 cached/new keys, so the local-window floor, split-K combine, softmax denominator, and seam between cached and new V are all scale-locked."
       },
       "attrs": { "num_heads": 8, "kv_num_heads": 2, "scale": 0.125, "local_window_size": 256 },
       "inputs": {
@@ -4710,7 +5004,7 @@
     {
       "name": "newkv_past_f16_flash_h8kv2_d64_q32p32_ramp_value_scale_lock",
       "provenance": {
-        "notes": "new_kv_past_flash_prefill and its _nosg twin are reached only by newkv_past_flash_h8kv2_d64_q32p32 and its f16 sibling; the f16 one averages zero-mean V to 0.034 under a 0.03 tolerance, so an 89% scale error passes. Ramping both value tensors makes each of the 32 chunked-prefill rows an O(1) function of its 16-key causal local window; early rows cross the past/new seam, locking the window floor, flash normalization, and cache transition on the f16 compile."
       },
       "attrs": { "num_heads": 8, "kv_num_heads": 2, "scale": 0.125, "local_window_size": 16 },
       "inputs": {
@@ -4751,7 +5045,7 @@
     {
       "name": "window_cache_decode_splitk_cap1024_b1q1_h2kv1d64_ramp_value_scale_lock",
       "provenance": {
-        "notes": "window_shift_decode_splitk and its _nosg twin are reached by one fixture, window_cache_decode_splitk_cap1024_b1q1_h2kv1d64, which already ramps the one appended V but leaves the 1024-slot shifted cache zero-mean, so the average is 1.7e-3 under a 0.005 tolerance and a 3x scale error passes. Ramping the past cache as well makes the output an O(1) function of the weighted mean slot in the shifted window, so the split-K combine and the eviction offset both move it: T=4096 > C=1024, so the cache genuinely wraps."
       },
       "attrs": { "num_heads": 2, "kv_num_heads": 1, "sliding_window_cache": 1, "local_window_size": 1024 },
       "inputs": {
@@ -4784,7 +5078,7 @@
     {
       "name": "quant_int8_scalar_prompt_ramp_value_scale_lock",
       "provenance": {
-        "notes": "The scalar INT8 KV-cache route is reached only by ORT-derived fixtures carrying a 0.05 absolute tolerance against an O(0.5) output, so a 10% uniform scale error passes every one of them; unlike the other attention holes this is loose tolerance rather than a cancelled average. A synthetic sibling on the same route reproduces the reference to under 1e-7 - the kernel and the oracle quantize identically, so the ORT tolerance was never needed here - which locks the vScale dequant, the softmax denominator and the apply divide at 2e-5. V ramps 0.5 to 2.0 so each causal row lands on a different O(1) value."
       },
       "attrs": {
         "num_heads": 2,
@@ -4821,7 +5115,7 @@
     {
       "name": "quant_int4_scalar_prompt_ramp_value_scale_lock",
       "provenance": {
-        "notes": "Same hole as the INT8 sibling but wider: the scalar INT4 KV-cache route only has ORT-derived fixtures at a 0.15 absolute tolerance over an O(0.5) output, blind to a 30% scale error. vScale 0.3 spans the ramp inside the signed 4-bit range, and the kernel matches the oracle to under 1e-7, so a 2e-5 tolerance locks the +8-biased nibble unpack, the dequant scale and the normalization."
       },
       "attrs": {
         "num_heads": 2,
@@ -5235,6 +5529,40 @@
         "presentValueT": { "dtype": "float16", "shape": [1, 1, 512, 64], "tolerance": 0.001 }
       }
     },
     {
       "name": "past_kv_materialized_sgmat_f16_floor_q256_p256_h2kv1_d64",
       "attrs": { "num_heads": 2, "kv_num_heads": 1 },
@@ -5374,6 +5702,39 @@
         "presentKeyT": { "dtype": "float32", "shape": [2, 1, 256, 64], "tolerance": 0.000001 },
         "presentValueT": { "dtype": "float32", "shape": [2, 1, 256, 64], "tolerance": 0.000001 }
       }
     }
   ]
 }

 {
   "fixtureArrays": {
     "ort_past_key_t_pattern": [0.2, 0.4, 0.6, 0.8, 1, 0.2, 0.4, 0.6, 0.8, 1, 0.2, 0.4, 0.6, 0.8, 1, 0.2, 0.4, 0.6, 0.8, 1, 0.2, 0.4, 0.6, 0.8, 1, 0.2, 0.4, 0.6, 0.8, 1, 0.2, 0.4, 0.6, 0.8, 1, 0.2, 0.4, 0.6, 0.8, 1, 0.2, 0.4, 0.6, 0.8, 1, 0.2, 0.4, 0.6, 0.8, 1, 0.2, 0.4, 0.6, 0.8, 1, 0.2, 0.4, 0.6, 0.8, 1, 0.2, 0.4, 0.6, 0.8],
     "ort_past_value_t_pattern": [0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3, 0.6, 0.9, 0.3],
         "valueT": {
           "dtype": "float32",
           "shape": [1, 512, 64],
+          "data": { "kind": "fillFloat32", "scale": 0.18, "sinStep": 0.09817477042, "cosStep": 0.000383495197 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [512] } }
         "valueT": {
           "dtype": "float32",
           "shape": [1, 2048, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.18, "sinStep": 0.04908738521, "cosStep": 0.00004793689962 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [2047] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [2048] } }
         "valueT": {
           "dtype": "float32",
           "shape": [1, 1024, 64],
+          "data": { "kind": "fillFloat32", "scale": 0.18, "sinStep": 0.09817477042, "cosStep": 0.0001917475985 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1023] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1024] } }
         "valueT": {
           "dtype": "float16",
           "shape": [1, 128, 64],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.09817477042, "cosStep": 0.001533980788 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [127] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [128] } }
       },
       "outputs": {
+        "outputT": { "dtype": "float16", "shape": [1, 128, 128], "tolerance": 0.004 },
+        "presentKeyT": { "dtype": "float16", "shape": [1, 1, 128, 64], "tolerance": 0.004 },
+        "presentValueT": { "dtype": "float16", "shape": [1, 1, 128, 64], "tolerance": 0.004 }
       }
     },
     {
         "valueT": {
           "dtype": "float16",
           "shape": [1, 512, 64],
+          "data": { "kind": "fillFloat32", "scale": 0.18, "sinStep": 0.09817477042, "cosStep": 0.000383495197 }
         },
         "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, 1, 256], "tolerance": 0.004 },
+        "presentKeyT": { "dtype": "float16", "shape": [1, 1, 512, 64], "tolerance": 0.004 },
+        "presentValueT": { "dtype": "float16", "shape": [1, 1, 512, 64], "tolerance": 0.004 }
       }
     },
     {
         "valueT": {
           "dtype": "float16",
           "shape": [1, 2048, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.18, "sinStep": 0.04908738521, "cosStep": 0.00004793689962 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [2047] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [2048] } }
       },
       "outputs": {
+        "outputT": { "dtype": "float16", "shape": [1, 1, 512], "tolerance": 0.004 },
+        "presentKeyT": { "dtype": "float16", "shape": [1, 2, 2048, 64], "tolerance": 0.004 },
+        "presentValueT": { "dtype": "float16", "shape": [1, 2, 2048, 64], "tolerance": 0.004 }
       }
     },
     {
         "valueT": {
           "dtype": "float32",
           "shape": [1, 1024, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.04908738521, "cosStep": 0.00009587379924 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1023] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1024] } }
         "vScaleT": { "dtype": "float32", "shape": [1], "data": { "kind": "values", "values": [0.0038802] } }
       },
       "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 4, 16], "tolerance": 0.01 },
         "presentKeyT": { "dtype": "int8", "shape": [1, 1, 4, 8], "tolerance": 0 },
         "presentValueT": { "dtype": "int8", "shape": [1, 1, 4, 8], "tolerance": 0 }
       }
         }
       },
       "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 4, 16], "tolerance": 0.01 },
         "presentKeyT": { "dtype": "int8", "shape": [1, 1, 4, 8], "tolerance": 0 },
         "presentValueT": { "dtype": "int8", "shape": [1, 1, 4, 8], "tolerance": 0 }
       }
         "vScaleT": { "dtype": "float32", "shape": [1], "data": { "kind": "values", "values": [0.0703979] } }
       },
       "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 4, 16], "tolerance": 0.01 },
         "presentKeyT": { "dtype": "uint8", "shape": [1, 1, 4, 4], "tolerance": 0 },
         "presentValueT": { "dtype": "uint8", "shape": [1, 1, 4, 4], "tolerance": 0 }
       }
         }
       },
       "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 4, 16], "tolerance": 0.01 },
         "presentKeyT": { "dtype": "uint8", "shape": [1, 1, 4, 4], "tolerance": 0 },
         "presentValueT": { "dtype": "uint8", "shape": [1, 1, 4, 4], "tolerance": 0 }
       }
         "vScaleT": { "dtype": "float32", "shape": [1], "data": { "kind": "values", "values": [0.0039328] } }
       },
       "outputs": {
+        "outputT": { "dtype": "float32", "shape": [2, 4, 64], "tolerance": 0.01 },
         "presentKeyT": { "dtype": "int8", "shape": [2, 2, 4, 16], "tolerance": 0 },
         "presentValueT": { "dtype": "int8", "shape": [2, 2, 4, 16], "tolerance": 0 }
       }
         "vScaleT": { "dtype": "float32", "shape": [1], "data": { "kind": "values", "values": [0.0039328] } }
       },
       "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 8, 128], "tolerance": 0.01 },
         "presentKeyT": { "dtype": "int8", "shape": [1, 1, 8, 64], "tolerance": 0 },
         "presentValueT": { "dtype": "int8", "shape": [1, 1, 8, 64], "tolerance": 0 }
       }
         "vScaleT": { "dtype": "float32", "shape": [1], "data": { "kind": "values", "values": [0.0703979] } }
       },
       "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 4, 64], "tolerance": 0.01 },
         "presentKeyT": { "dtype": "uint8", "shape": [1, 1, 4, 8], "tolerance": 0 },
         "presentValueT": { "dtype": "uint8", "shape": [1, 1, 4, 8], "tolerance": 0 }
       }
         "vScaleT": { "dtype": "float32", "shape": [1], "data": { "kind": "values", "values": [0.0038802] } }
       },
       "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 4, 16], "tolerance": 0.01 },
         "presentKeyT": { "dtype": "int8", "shape": [1, 1, 4, 8], "tolerance": 0 },
         "presentValueT": { "dtype": "int8", "shape": [1, 1, 4, 8], "tolerance": 0 }
       }
         "presentValueT": { "dtype": "float32", "shape": [1, 1, 8, 8], "tolerance": 0.0001 }
       }
     },
+    {
+      "name": "ort_bias_decode_all_neg_inf_row_zero",
+      "provenance": {
+        "notes": "Cached decode uses two query heads, one KV head, head size 8, eight active cache positions, and no new K/V values. Bias for query head 1 is -Infinity at every key while head 0 remains finite. The masked head's output is zero; the finite head is its softmax-weighted V. `seqlens_k = 7` means eight active keys, not an empty row. Present K/V outputs copy the input cache unchanged."
+      },
+      "attrs": { "num_heads": 2, "kv_num_heads": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 16],
+          "data": {
+            "kind": "values",
+            "values": { "$ref": "#/fixtureArrays/ort_sharedkv_decode_b1q1p8_h2kv1d8_input_queryT" }
+          }
+        },
+        "keyT": { "dtype": "float32", "shape": [1, 0, 8], "data": { "kind": "values", "values": [] } },
+        "valueT": { "dtype": "float32", "shape": [1, 0, 8], "data": { "kind": "values", "values": [] } },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 8, 8],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/ort_past_key_t_pattern" } }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 8, 8],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/ort_past_value_t_pattern" } }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [7] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [8] } },
+        "attentionBiasT": {
+          "dtype": "float32",
+          "shape": [1, 2, 1, 8],
+          "data": {
+            "kind": "values",
+            "values": [-0.05, -0.0020574, 0.0341471, 0.0497495, 0.0409297, 0.0098472, -0.035888, -0.0850783, "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity"]
+          }
+        }
+      },
+      "outputs": {
+        "outputT": {
+          "dtype": "float32",
+          "shape": [1, 1, 16],
+          "tolerance": 0.0001,
+          "data": {
+            "kind": "values",
+            "values": [0.6027476, 0.5584169, 0.6388355, 0.6027476, 0.5584169, 0.6388355, 0.6027476, 0.5584169, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
+          }
+        },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 1, 8, 8], "tolerance": 0.0001 },
+        "presentValueT": { "dtype": "float32", "shape": [1, 1, 8, 8], "tolerance": 0.0001 }
+      }
+    },
+    {
+      "name": "ort_bias_decode_single_finite_key_control",
+      "provenance": {
+        "notes": "Attention bias has one finite key per head and -Infinity elsewhere. Each softmax therefore selects exactly one cached value: head 0 returns row 5 and head 1 returns row 3. This distinguishes a valid one-key denominator from the fully masked behavior in `ort_bias_decode_all_neg_inf_row_zero`."
+      },
+      "attrs": { "num_heads": 2, "kv_num_heads": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 16],
+          "data": {
+            "kind": "values",
+            "values": { "$ref": "#/fixtureArrays/ort_sharedkv_decode_b1q1p8_h2kv1d8_input_queryT" }
+          }
+        },
+        "keyT": { "dtype": "float32", "shape": [1, 0, 8], "data": { "kind": "values", "values": [] } },
+        "valueT": { "dtype": "float32", "shape": [1, 0, 8], "data": { "kind": "values", "values": [] } },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 8, 8],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/ort_past_key_t_pattern" } }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 8, 8],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/ort_past_value_t_pattern" } }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [7] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [8] } },
+        "attentionBiasT": {
+          "dtype": "float32",
+          "shape": [1, 2, 1, 8],
+          "data": {
+            "kind": "values",
+            "values": ["-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", 0.0, "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", 0.0, "-Infinity", "-Infinity", "-Infinity", "-Infinity"]
+          }
+        }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 1, 16], "tolerance": 0.0001 },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 1, 8, 8], "tolerance": 0.0001 },
+        "presentValueT": { "dtype": "float32", "shape": [1, 1, 8, 8], "tolerance": 0.0001 }
+      }
+    },
     {
       "name": "ort_bias_bcast_prompt_b1q8p8_h2kv1d8",
       "attrs": { "num_heads": 2, "kv_num_heads": 1 },
     {
       "name": "flashprefill_pastzero_h8kv2_d128_q16",
       "provenance": {
+        "notes": "A compact shared-KV prefill with 32 query heads, eight KV heads, head size 128, and sequence length 64 exercises grouped attention."
       },
       "attrs": { "num_heads": 8, "kv_num_heads": 2, "scale": 0.08838834764831845 },
       "inputs": {
         "pastValueT": {
           "dtype": "float32",
           "shape": [1, 2, 48, 64],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.09817477042, "cosStep": 0.002045307717 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [47] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [48] } }
       "name": "qnorm_rotary_scalar_h8kv2_d256_q1",
       "provenance": {
         "source": "scope coverage",
+        "test": "headDim=256 scalar Q-norm and rotary specialization",
         "notes": "headDim=256 requires 34,048 bytes for the cooperative path, exceeding a 32 KiB workgroup-storage tier and exercising the thread-per-query Q-norm route."
       },
       "attrs": { "num_heads": 8, "kv_num_heads": 2, "scale": 0.0625, "do_rotary": 1, "qk_norm_epsilon": 0.00001 },
       "name": "f16_qnorm_rotary_scalar_f16wts_h8kv2_d256_q1",
       "provenance": {
         "source": "scope coverage",
+        "test": "headDim=256 float16-weight scalar Q-norm and rotary specialization",
         "notes": "Float16 counterpart to the thread-per-query Q-norm case. It verifies explicit f32 conversion of q_norm_weight, cos_cache, and sin_cache before scalar arithmetic."
       },
       "attrs": { "num_heads": 8, "kv_num_heads": 2, "scale": 0.0625, "do_rotary": 1, "qk_norm_epsilon": 0.00001 },
         "pastValueT": {
           "dtype": "float32",
           "shape": [1, 2, 1024, 64],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.09817477042, "cosStep": 0.00009587379924 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1023] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1024] } }
     {
       "name": "flashdecode_splitk_h8kv2_d128_kv1024",
       "provenance": {
+        "notes": "A compact float32 shared-past decode with 32 query heads, eight KV heads, head size 128, and KV length 2,048 exercises split-K grouped attention."
       },
       "attrs": { "num_heads": 8, "kv_num_heads": 2, "scale": 0.08838834764831845 },
       "inputs": {
         "pastValueT": {
           "dtype": "float32",
           "shape": [1, 2, 1024, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.04908738521, "cosStep": 0.00004793689962 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1023] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1024] } }
         "pastValueT": {
           "dtype": "float16",
           "shape": [1, 2, 512, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.04908738521, "cosStep": 0.00009587379924 }
         },
         "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, 1, 1024], "tolerance": 0.004 },
         "presentKeyT": { "dtype": "float16", "shape": [1, 2, 512, 128], "tolerance": 0.001 },
         "presentValueT": { "dtype": "float16", "shape": [1, 2, 512, 128], "tolerance": 0.001 }
       }
         "pastValueT": {
           "dtype": "float16",
           "shape": [1, 2, 1024, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.04908738521, "cosStep": 0.00004793689962 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1023] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1024] } }
       },
       "outputs": {
+        "outputT": { "dtype": "float16", "shape": [1, 1, 1024], "tolerance": 0.004 },
         "presentKeyT": { "dtype": "float16", "shape": [1, 2, 1024, 128], "tolerance": 0.001 },
         "presentValueT": { "dtype": "float16", "shape": [1, 2, 1024, 128], "tolerance": 0.001 }
       }
         "pastValueT": {
           "dtype": "float16",
           "shape": [1, 2, 512, 64],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.09817477042, "cosStep": 0.0001917475985 }
         },
         "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, 1, 512], "tolerance": 0.004 },
         "presentKeyT": { "dtype": "float16", "shape": [1, 2, 512, 64], "tolerance": 0.001 },
         "presentValueT": { "dtype": "float16", "shape": [1, 2, 512, 64], "tolerance": 0.001 }
       }
         "pastValueT": {
           "dtype": "float16",
           "shape": [1, 2, 32, 64],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.09817477042, "cosStep": 0.003067961576 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [31] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [32] } },
         "pastValueT": {
           "dtype": "float16",
           "shape": [1, 2, 32, 64],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.09817477042, "cosStep": 0.003067961576 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [31] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [32] } },
         "valueT": {
           "dtype": "float32",
           "shape": [1, 32, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.04908738521, "cosStep": 0.003067961576 }
         },
         "pastKeyT": {
           "dtype": "float32",
         "pastValueT": {
           "dtype": "float32",
           "shape": [1, 2, 32, 64],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.09817477042, "cosStep": 0.003067961576 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [63] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [64] } }
         "valueT": {
           "dtype": "float32",
           "shape": [1, 1, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.04908738521, "cosStep": 0.09817477042 }
         },
         "pastKeyT": {
           "dtype": "float32",
         "pastValueT": {
           "dtype": "float32",
           "shape": [1, 2, 1024, 64],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.09817477042, "cosStep": 0.00009587379924 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1024] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1025] } }
         "valueT": {
           "dtype": "float16",
           "shape": [1, 32, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.04908738521, "cosStep": 0.003067961576 }
         },
         "pastKeyT": {
           "dtype": "float16",
         "pastValueT": {
           "dtype": "float16",
           "shape": [1, 2, 32, 64],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.09817477042, "cosStep": 0.003067961576 }
         },
         "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, 32, 512], "tolerance": 0.004 },
         "presentKeyT": { "dtype": "float16", "shape": [1, 2, 64, 64], "tolerance": 0.001 },
         "presentValueT": { "dtype": "float16", "shape": [1, 2, 64, 64], "tolerance": 0.001 }
       }
         "valueT": {
           "dtype": "float32",
           "shape": [1, 511, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.04908738521, "cosStep": 0.0001921228384 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [510] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } }
         "valueT": {
           "dtype": "float32",
           "shape": [1, 64, 256],
+          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.02454369261, "cosStep": 0.0007669903939 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [63] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [64] } }
         "valueT": {
           "dtype": "float32",
           "shape": [1, 512, 256],
+          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.02454369261, "cosStep": 0.00009587379924 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [512] } }
     {
       "name": "qkv_prefill_f32_h8kv2_d256_register_boundary",
       "provenance": {
+        "notes": "A QKV prefill with eight query heads, two KV heads, and head size 256 exercises generic online attention on subgroup tiers and shared-memory clustered attention on tiers without subgroups."
       },
       "attrs": { "num_heads": 8, "kv_num_heads": 2, "scale": 0.0625, "causal": 0 },
       "inputs": {
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/bert/flash_attention_decode_qkv.wgsl.template",
         "test": "head_size=8 split-reduce OOB race class (PR #29593)",
+        "notes": "Diverges from the upstream test's inputs (inputs.pastValueT fillFloat32 -> fillFloat32); the expected output is recomputed by the CPU reference for the new inputs. A head size below the flash minimum exercises the cooperative fallback and its tile-width-guarded reduction."
       },
       "attrs": { "num_heads": 4, "kv_num_heads": 2 },
       "inputs": {
         "pastValueT": {
           "dtype": "float32",
           "shape": [1, 2, 512, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.7853981634, "cosStep": 0.001533980788, "scale": 0.5 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [512] } }
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/bert/flash_attention_decode_qkv.wgsl.template",
         "test": "head_size=16 split-reduce OOB race class (PR #29593)",
+        "notes": "Diverges from the upstream test's inputs (inputs.pastValueT fillFloat32 -> fillFloat32); the expected output is recomputed by the CPU reference for the new inputs. Head size 16 exercises the cooperative fallback at the next four-wide tiny-head boundary."
       },
       "attrs": { "num_heads": 4, "kv_num_heads": 2 },
       "inputs": {
         "pastValueT": {
           "dtype": "float32",
           "shape": [1, 2, 512, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.3926990817, "cosStep": 0.0007669903939, "scale": 0.5 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [512] } }
       "provenance": {
         "source": "onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc",
         "test": "GQA head sink on the buffer-sharing append path",
+        "notes": "One K/V row is appended at index 11 in a capacity-32 cache. The 20 unused rows must not enter attention; `seqlens_k` determines both the append offset and window origin."
       }
     },
     {
     {
       "name": "prefill_tiled_q31_h8kv2_d36_non_cluster_stride",
       "provenance": {
+        "notes": "Head size 36 is divisible by four but not by the 16 values mapped across each query cluster. With 31 queries and eight heads, tiers without subgroups exercise blocked tiled prefill."
       },
       "attrs": { "num_heads": 8, "kv_num_heads": 2, "scale": 0.16666666666666666, "causal": 0 },
       "inputs": {
         "presentValueT": { "dtype": "float32", "shape": [1, 2, 32, 64], "tolerance": 0.0001 }
       },
       "provenance": {
+        "notes": "Logit soft-capping on the past-KV prefill path exercises the tiled attention route with a per-score cap before softmax."
       }
     },
     {
     {
       "name": "flashprefill_quant_int8_h4kv2_d64_q32",
       "provenance": {
+        "notes": "INT8 KV cache on the tiled prefill path. Each key is dequantized once into the staged tile and reused by every query in that tile."
       },
       "attrs": {
         "num_heads": 4,
       "provenance": {
         "source": "onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc",
         "test": "GQA sliding_window_cache (ORT #29904), CPU/CUDA reference semantics",
+        "notes": "A capacity-8 windowed cache reaches `T = 6` without eviction, so its origin remains zero and unused rows must be cleared."
       }
     },
     {
       "provenance": {
         "source": "onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc",
         "test": "GQA sliding_window_cache (ORT #29904), CPU/CUDA reference semantics",
+        "notes": "At `T = 12` with capacity 8, appending one token evicts one row, writes at row 7, and advances the absolute cache origin to 4."
       }
     },
     {
       "provenance": {
         "source": "onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc",
         "test": "GQA sliding_window_cache (ORT #29904), CPU/CUDA reference semantics",
+        "notes": "Appending four tokens to a capacity-8 cache at `T = 10` evicts two rows, begins the append at row 4, and advances the absolute cache origin to 2."
       }
     },
     {
       "provenance": {
         "source": "onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc",
         "test": "GQA sliding_window_cache (ORT #29904), CPU/CUDA reference semantics",
+        "notes": "Two batches require independent cache origins: batch 0 reaches `T = 4` without eviction, while batch 1 reaches `T = 12`, evicts one row, and uses origin 4."
       }
     },
     {
       "provenance": {
         "source": "onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc",
         "test": "GQA sliding_window_cache (ORT #29904), CPU/CUDA reference semantics",
+        "notes": "A local window of four within a capacity-8 cache combines one-row eviction with a nonzero attention floor derived from the absolute cache origin."
       }
     },
     {
       "provenance": {
         "source": "onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc",
         "test": "GQA sliding_window_cache (ORT #29904), CPU/CUDA reference semantics",
+        "notes": "Cache capacity and local-window size are both four at `T = 10`, exercising a fully occupied sliding-window allocation with a nonzero origin."
       }
     },
     {
       "provenance": {
         "source": "onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc",
         "test": "GQA sliding_window_cache (ORT #29904), CPU/CUDA reference semantics",
+        "notes": "At `T = 64` with capacity 16, the cache origin is 48. The four-to-one grouped-query layout and head size 16 exercise absolute-position translation far beyond the physical buffer."
       }
     },
     {
         "pastValueT": {
           "dtype": "float32",
           "shape": [1, 1, 1024, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.09817477042, "cosStep": 0.0001917475985, "scale": 0.5 }
         },
         "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [4095] } },
         "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [4096] } }
       "provenance": {
         "source": "onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc",
         "test": "GQA sliding_window_cache (ORT #29904), CPU/CUDA reference semantics",
+        "notes": "Diverges from the upstream test's inputs (inputs.pastValueT fillFloat32 -> fillFloat32); the expected output is recomputed by the CPU reference for the new inputs. Windowed cache at a capacity that clears CACHED_DECODE_MIN_KV_TOKENS, so the split-K decode variant is selected rather than the scalar fallback. T=4096 > C=1024, so the step genuinely evicts and the split-K kernel reads a shifted cache."
       }
     },
     {
       "provenance": {
         "source": "onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc",
         "test": "GQA buffer-sharing append chunk prefill on the flash cluster kernel",
+        "notes": "A fixed-capacity cache has active length 48 and capacity 64 before a 32-row prefill. The flash route must use the live length for attention and the capacity for batch strides."
       }
     },
     {
       "provenance": {
         "source": "onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc",
         "test": "GQA buffer-sharing append chunk prefill on the flash cluster kernel",
+        "notes": "Two batches share capacity-64 caches but have live lengths 32 and 48. The flash route must track each attention bound independently while retaining the fixed batch stride."
+      }
+    },
+    {
+      "name": "share_append_sgmat_chunk_slack_b1q128cap384_h2kv1d64",
+      "attrs": { "num_heads": 2, "kv_num_heads": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.27, "scale": 0.5 }
+        },
+        "keyT": { "dtype": "float32", "shape": [1, 128, 64], "data": { "kind": "linspace", "start": -1.0, "end": 1.0 } },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 128, 64],
+          "data": { "kind": "linspace", "start": 0.0, "end": 2.0 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 384, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 384, 64],
+          "data": { "kind": "linspace", "start": 0.0, "end": 4.0 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [255] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 128, 128], "tolerance": 0.005 },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 1, 384, 64], "tolerance": 0.0001 },
+        "presentValueT": { "dtype": "float32", "shape": [1, 1, 384, 64], "tolerance": 0.0001 }
+      },
+      "provenance": {
+        "source": "onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc",
+        "test": "GQA buffer-sharing chunk append on the materialized subgroup-matrix route",
+        "notes": "128 past rows plus a 128-row appended chunk in a 384-capacity shared buffer: the live length (256) is shorter than the capacity, so the causal bound and the tile skips must come from seqlens_k while every stride stays capacity-sized."
+      }
+    },
+    {
+      "name": "share_append_sgmat_chunk_full_b1q128cap256_h2kv1d128",
+      "attrs": { "num_heads": 2, "kv_num_heads": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 128, 256],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.27, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 128, 128],
+          "data": { "kind": "linspace", "start": -1.0, "end": 1.0 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 128, 128],
+          "data": { "kind": "linspace", "start": 0.0, "end": 2.0 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 256, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 256, 128],
+          "data": { "kind": "linspace", "start": 0.0, "end": 4.0 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [255] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 128, 256], "tolerance": 0.005 },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 1, 256, 128], "tolerance": 0.0001 },
+        "presentValueT": { "dtype": "float32", "shape": [1, 1, 256, 128], "tolerance": 0.0001 }
+      },
+      "provenance": {
+        "source": "onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc",
+        "test": "GQA buffer-sharing chunk append on the materialized subgroup-matrix route",
+        "notes": "The append exactly fills the shared capacity at d=128 (two apply column tiles); live length equals capacity, so the seqlens bound and the compile-time bound coincide."
+      }
+    },
+    {
+      "name": "share_append_sgmat_chunk_mixed_batch_b2q128cap384_h2kv1d64",
+      "attrs": { "num_heads": 2, "kv_num_heads": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.27, "scale": 0.5 }
+        },
+        "keyT": { "dtype": "float32", "shape": [2, 128, 64], "data": { "kind": "linspace", "start": -1.0, "end": 1.0 } },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [2, 128, 64],
+          "data": { "kind": "linspace", "start": 0.0, "end": 2.0 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 1, 384, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 1, 384, 64],
+          "data": { "kind": "linspace", "start": 0.0, "end": 4.0 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [255, 127] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [2, 128, 128], "tolerance": 0.005 },
+        "presentKeyT": { "dtype": "float32", "shape": [2, 1, 384, 64], "tolerance": 0.0001 },
+        "presentValueT": { "dtype": "float32", "shape": [2, 1, 384, 64], "tolerance": 0.0001 }
+      },
+      "provenance": {
+        "source": "onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc",
+        "test": "GQA buffer-sharing chunk append on the materialized subgroup-matrix route",
+        "notes": "Per-batch live lengths from seqlens_k: batch 0 has 128 past rows, batch 1 appends into an empty cache, and the same dispatch serves both."
+      }
+    },
+    {
+      "name": "share_append_sgmat_chunk_right_padded_b2q128cap384_h2kv1d64",
+      "attrs": { "num_heads": 2, "kv_num_heads": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.27, "scale": 0.5 }
+        },
+        "keyT": { "dtype": "float32", "shape": [2, 128, 64], "data": { "kind": "linspace", "start": -1.0, "end": 1.0 } },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [2, 128, 64],
+          "data": { "kind": "linspace", "start": 0.0, "end": 2.0 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 1, 384, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 1, 384, 64],
+          "data": { "kind": "linspace", "start": 0.0, "end": 4.0 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [255, 63] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [2, 128, 128], "tolerance": 0.005 },
+        "presentKeyT": { "dtype": "float32", "shape": [2, 1, 384, 64], "tolerance": 0.0001 },
+        "presentValueT": { "dtype": "float32", "shape": [2, 1, 384, 64], "tolerance": 0.0001 }
+      },
+      "provenance": {
+        "source": "onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc",
+        "test": "GQA buffer-sharing chunk append on the materialized subgroup-matrix route",
+        "notes": "Batch 1's live length (64) is shorter than the 128-row query chunk, so its queries left-align against the live region exactly as the flash routes' saturated past length does; batch 0 stays right-aligned."
+      }
+    },
+    {
+      "name": "share_append_sgmat_chunk_qtail_b1q160cap384_h2kv1d64",
+      "attrs": { "num_heads": 2, "kv_num_heads": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 160, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.27, "scale": 0.5 }
+        },
+        "keyT": { "dtype": "float32", "shape": [1, 160, 64], "data": { "kind": "linspace", "start": -1.0, "end": 1.0 } },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 160, 64],
+          "data": { "kind": "linspace", "start": 0.0, "end": 2.0 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 384, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 384, 64],
+          "data": { "kind": "linspace", "start": 0.0, "end": 4.0 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [255] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 160, 128], "tolerance": 0.005 },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 1, 384, 64], "tolerance": 0.0001 },
+        "presentValueT": { "dtype": "float32", "shape": [1, 1, 384, 64], "tolerance": 0.0001 }
+      },
+      "provenance": {
+        "source": "onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc",
+        "test": "GQA buffer-sharing chunk append on the materialized subgroup-matrix route",
+        "notes": "q=160 is not a query-tile multiple, so the guarded query tails run together with the live-length ceiling."
       }
     },
     {
     {
       "name": "splitk_decode_h8kv2_d64_kv2048_f16_ramp_value_scale_lock",
       "provenance": {
+        "notes": "A monotone value ramp keeps float16 decode outputs at order-one magnitude, making split-K accumulation, cross-partition rescaling, and the final normalization observable."
       },
       "attrs": { "num_heads": 8, "kv_num_heads": 2, "scale": 0.125, "causal": 0 },
       "inputs": {
     {
       "name": "flash_prefill_h2kv1_d64_s128_f16_ramp_value_scale_lock",
       "provenance": {
+        "notes": "A monotone value ramp gives each of 128 causal float16 prefill rows a distinct order-one output, exposing the register-blocked q32 epilogue division and row selection."
       },
       "attrs": { "num_heads": 2, "kv_num_heads": 1, "scale": 0.125, "causal": 0 },
       "inputs": {
     {
       "name": "newkv_past_splitk_h8kv2_d64_q1p1024_ramp_value_scale_lock",
       "provenance": {
+        "notes": "Monotone cached and new value ramps make decode output an order-one function of the final 256 key positions. This exposes the local-window floor, split-K combine, softmax denominator, and cached/new seam."
       },
       "attrs": { "num_heads": 8, "kv_num_heads": 2, "scale": 0.125, "local_window_size": 256 },
       "inputs": {
     {
       "name": "newkv_past_f16_flash_h8kv2_d64_q32p32_ramp_value_scale_lock",
       "provenance": {
+        "notes": "Monotone cached and new value ramps make each of 32 float16 prefill rows an order-one function of its 16-key causal window. Early rows cross the cached/new seam, exposing the window floor and flash normalization."
       },
       "attrs": { "num_heads": 8, "kv_num_heads": 2, "scale": 0.125, "local_window_size": 16 },
       "inputs": {
     {
       "name": "window_cache_decode_splitk_cap1024_b1q1_h2kv1d64_ramp_value_scale_lock",
       "provenance": {
+        "notes": "A monotone past-cache ramp makes the decode output depend at order-one scale on the weighted slot within the shifted window. With `T = 4096` and capacity 1024, both the eviction offset and split-K combine affect the result."
       },
       "attrs": { "num_heads": 2, "kv_num_heads": 1, "sliding_window_cache": 1, "local_window_size": 1024 },
       "inputs": {
     {
       "name": "quant_int8_scalar_prompt_ramp_value_scale_lock",
       "provenance": {
+        "notes": "An INT8 value ramp from 0.5 to 2.0 gives each causal row a distinct order-one output. Tight expected values exercise value-scale dequantization, the softmax denominator, and the scalar apply division."
       },
       "attrs": {
         "num_heads": 2,
     {
       "name": "quant_int4_scalar_prompt_ramp_value_scale_lock",
       "provenance": {
+        "notes": "An INT4 value ramp with scale 0.3 spans the signed nibble range and gives order-one outputs. Tight expected values exercise +8-biased nibble unpacking, dequantization, and normalization."
       },
       "attrs": {
         "num_heads": 2,
         "presentValueT": { "dtype": "float16", "shape": [1, 1, 512, 64], "tolerance": 0.001 }
       }
     },
+    {
+      "name": "past_kv_materialized_sgmat_f16_cooperative_q512_p512_h2kv1_d64",
+      "provenance": {
+        "notes": "A 128-row query tile uses 512 workgroup threads and the generic cooperative loader. Both GEMM phases stage tiles with a workgroup-strided linear walk, and the fused-softmax apply narrows each normalized score to float16 when staging it."
+      },
+      "tunables": { "MATERIALIZED_SGMAT_QUERY_TILE": 128 },
+      "attrs": { "num_heads": 2, "kv_num_heads": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [1, 512, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": { "dtype": "float16", "shape": [1, 0, 64], "data": { "kind": "values", "values": [] } },
+        "valueT": { "dtype": "float16", "shape": [1, 0, 64], "data": { "kind": "values", "values": [] } },
+        "pastKeyT": {
+          "dtype": "float16",
+          "shape": [1, 1, 512, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "pastValueT": {
+          "dtype": "float16",
+          "shape": [1, 1, 512, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41 }
+        },
+        "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, 128], "tolerance": 0.03 },
+        "presentKeyT": { "dtype": "float16", "shape": [1, 1, 512, 64], "tolerance": 0.001 },
+        "presentValueT": { "dtype": "float16", "shape": [1, 1, 512, 64], "tolerance": 0.001 }
+      }
+    },
     {
       "name": "past_kv_materialized_sgmat_f16_floor_q256_p256_h2kv1_d64",
       "attrs": { "num_heads": 2, "kv_num_heads": 1 },
         "presentKeyT": { "dtype": "float32", "shape": [2, 1, 256, 64], "tolerance": 0.000001 },
         "presentValueT": { "dtype": "float32", "shape": [2, 1, 256, 64], "tolerance": 0.000001 }
       }
+    },
+    {
+      "name": "flashdecode_splitk_h8kv2_d16_kv1024",
+      "attrs": { "num_heads": 8, "kv_num_heads": 2, "scale": 0.25 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 128],
+          "data": { "kind": "fillFloat32", "scale": 2.0, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": { "dtype": "float32", "shape": [1, 0, 32], "data": { "kind": "values", "values": [] } },
+        "valueT": { "dtype": "float32", "shape": [1, 0, 32], "data": { "kind": "values", "values": [] } },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 2, 1024, 16],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.017, "cosStep": 0.019 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 2, 1024, 16],
+          "data": { "kind": "fillFloat32", "scale": 3.0, "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, 1, 128], "tolerance": 0.0001 },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 2, 1024, 16], "tolerance": 0.0001 },
+        "presentValueT": { "dtype": "float32", "shape": [1, 2, 1024, 16], "tolerance": 0.0001 }
+      },
+      "provenance": {
+        "notes": "Cached decode at head size 16 exercises the split-K route exactly at its minimum head-size boundary."
+      }
     }
   ]
 }